[merceod]

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.

[stephen-dwq]

LGTM

[NSagan271]

The guard looks good and is correct to have, but unless I'm missing something, the worker loop is not intended to spuriously schedule other batches while the preplan is active. The speculative step always runs unless there is an admit failure, other error, or prepare_inputs returns all None.

Where I think the gap is that we skip rids if prepare_inputs returns None (which cosmos does if it's the last step of a batch, and I believe the engine's exec does not reset preplan in the case that all requests have been dropped. nevermind, this is handled by preplan_is_stale.

Some potential places it could be (?):

• _handle_main_loop_error never clears preplan, and also doesn't drain the plan_future.
• _execute_on_gpu_thread's finally block doesn't clear preplan either

[NSagan271]

I'd change this to a logger.warning; it something that shouldn't be the typical path

[merceod]

Done (1545386). It also names the staged step and the one that displaced it, rows included.

[merceod]

Looked at both places. _handle_main_loop_error was a real gap: a speculation that was armed (plan future running) but not yet submitted when the pending step raised kept its stage, and the next batch off the scheduler would have inherited it. The handler now drains that future, drops the stage through the engine and hands the fresh rids back to their queues (c076d17). _execute_on_gpu_thread drops the stage when prepare_inputs raises for a pre-planned batch, and Engine.exec does the same when declare or admit raise before the plan is promoted (c076d17, e0da66c). The guard log is a warning now and names both steps (1545386); tests in 06b8cda and a134ffd.

On where the 09-17 stale pre-plan came from: the log we have shows a one-row decode pre-plan displaced by a five-row decode step of the same node, but I could not pin down which path armed it. With the warning naming both steps, the next occurrence will say.

[NSagan271]

LGTM