Skip to content

Cosmos3-edge PR 4: a staged pre-plan only promotes into the step it was planned for - #295

Open
merceod wants to merge 7 commits into
mainfrom
engine/preplan-guard
Open

merceod wants to merge 7 commits into
mainfrom
engine/preplan-guard

Conversation

@merceod

@merceod merceod commented Sep 22, 2026

Copy link
Copy Markdown
Collaborator

Fixes two crashes and a silent-corruption path in the speculative pre-plan, found while running concurrent Cosmos3-Edge reasoner requests.

The plan thread pre-plans a node's speculated next step while the current one runs. If the scheduler then dispatches a different batch on that node first (a new request's eager prefill while a decode step sits pre-planned), the pre-planning resources promoted the staged plan into the wrong step.

  • SamplerResource.plan dereferenced ctx.slot_lease.slot on the unleased step and took the worker main loop down ('NoneType' object has no attribute 'slot'), killing every in-flight request.
  • KVManager.admit / plan would have handed the staged request's reserved pages and cached plan to the other step's segments.
  • With only those two guarded, the stage was no longer all-or-nothing. The KV manager dropped its stale plan and planned the foreign step afresh, but FlashInferManager.plan (and likewise the position and cross-attention managers) promote whatever is staged when plan is next called, so attention ran with wrappers laid out for the other step's rows (q implies q_len_per_req=5 but plan() used 1). With equal row counts it would have attended over the other request's pages without any error.

The guard now lives in StepRunner, the one place that sees the whole step. pre_plan records what it staged (walk, padded rows, lease slot, capture key, every resource's segments), and admit / plan clear the stage on every resource before a different step proceeds. Engine.reset_pre_plan_for_batch goes through the same call. The KV and sampler checks stay as a second line of defense, and the pre-planned step itself just plans inline when its turn comes.

Tests in test/modular/test_preplan_guard.py (CPU). With the real KVManager, a foreign step admits and plans on its own pages while a pre-plan is staged and the staged plan is dropped, and a pre-plan still promotes into its own step. With a StepRunner over two blindly-promoting fake resources, a foreign step (other rows, the same rows without their lease, another slot, another span) clears the stage on both, the staged step promotes on both, and clear_preplan reaches every resource. Also verified on H100 against the Cosmos3-Edge server with deterministic prompts at batch sizes 3 and 5 (padded to buckets 4 and 8), where every concurrent answer matched its sequential one byte for byte. ruff and test/modular pass.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant