Conversation
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.plandereferencedctx.slot_lease.sloton the unleased step and took the worker main loop down ('NoneType' object has no attribute 'slot'), killing every in-flight request.KVManager.admit/planwould have handed the staged request's reserved pages and cached plan to the other step's segments.FlashInferManager.plan(and likewise the position and cross-attention managers) promote whatever is staged whenplanis 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_planrecords what it staged (walk, padded rows, lease slot, capture key, every resource's segments), andadmit/planclear the stage on every resource before a different step proceeds.Engine.reset_pre_plan_for_batchgoes 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 realKVManager, 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 aStepRunnerover 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, andclear_preplanreaches 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.ruffandtest/modularpass.