fix(SPEC-DFLASH2): key the draft context by request id, so a condensed batch row cannot take it from its owner (#2008)

[localai-bot]

DFlash2 served exactly one sequence. At concurrency 2 the engine refused on its own position invariant one request in, EngineCore stopped, and every later request on that server came back [request submitted to a stopped AsyncLLM]. This is the blocking defect for the #1574 three-engine comparison: vLLM and SGLang were both measured at c=1..32 with the same DFlash2 draft at k=8, and we could produce c=1 and nothing above it.

Closes #2008.

Why the accumulation desynchronises

The runner held the draft context in four arrays indexed by batch row. A row index is not stable for a request's lifetime here: InputBatch::condense slides a live request down into the hole a finished neighbour left (input_batch.cpp:686-706), and swap_states exchanges two live rows (:762-847). Both permute every per-slot array they own — req_ids, num_computed_tokens_cpu, num_accepted_tokens, last_sampled_tokens, prefill_len, the block-table rows, the index-keyed sampling maps — and neither knows these four exist, because they live on GPUModelRunner rather than in InputBatch.

So when the first of two concurrent requests finished, the survivor was moved onto the departed request's bookkeeping. The reused-slot test at runner.cpp:2895-2906 read a changed occupant, allocated a fresh empty store and set ctx_len = 0 — under a request still using it — and the invariant at :2939-2945 then refused rather than drafting from a foreign context. ok=1 at c=2 is that mechanism's signature, not an incidental count.

At c=1 the only row is row 0: a request finishes, the next is admitted into the same row, and the reset is correct, because the new occupant is a fresh prefill whose first position is 0. The condense move that breaks the state exists only when a second live request has to be slid over a departed one.

The operator's isolation had already settled the layer: with --speculative-config omitted and everything else identical, the same two concurrent requests both complete.

This tree already names the bug class, in a comment above the log written to solve it for a different array (input_batch.h:240-252): "Upstream needs no equivalent because it never condenses ... the slot index is stable for the request's lifetime. This log is the price of our condensed dense batch." The DFlash2 arrays never got that treatment.

Upstream

Read beyond our parity pin at b389ac2946 (the clone is shallow and checked out at the pin, so anchors were read with git show). Upstream has no analogue of the host-side counter at all. DFlash and DFlash2 address the draft KV by absolute position (dflash/speculator.py:562-590), re-read their one anchor from the target's own positions every step (:553), and index every cross-step tensor by the persistent request slot through req_state_idx = idx_mapping[req_idx] (:536, dflash2/speculator.py:95) — in a V2 runner that has no condense at all (gpu/states.py:29,100,132). Its legacy V1 runner does condense, and there the draft's block-table row moves with the request (gpu_input_batch.py:786 → block_table.py:367-373).

Both upstream shapes key the draft context to the request. Ours keyed it to neither.

The change

One unordered_map<string, DflashReqCtx> holding the store, the context length and the disabled flag. Every row permutation the batch can perform — today's two and any future one — is now a no-op for the draft context.

The reused-slot test goes away with it. "Has this row's occupant changed" was only ever a proxy for "is this state this request's", and the proxy is precisely what broke: after a condense move the answer was yes for a request whose context was perfectly valid. Keyed by id the question cannot be asked wrongly; a first sight of an id constructs the entry, and that construction is the reset. dflash_ctx_disabled_'s own comment — "the flag is a property of the REQUEST, not of the row" — becomes literally true. The decode-first reorder is fixed for free, and it is a second live trigger rather than a hypothetical one. reorder_batch_to_split_decodes_and_prefills runs unconditionally every step (runner.cpp:1324) and swaps live rows through swap_states. It emits no swap only while arrival order already puts decodes before prefills — the common case, since condense keeps older decoding requests at low rows and a new arrival appends at the end as a prefill — which is why the #2008 measurement met the condense move first. A batch whose regions are out of order does swap, and on the pre-change code that exchanged two live requests' draft contexts. (The commit body and an earlier draft of the spec called this path "inert for a Qwen3 target today"; that was wrong and the spec now says so — the reorder has no model-family gate at its call site.)

Entries are pruned each propose against InputBatch's membership — the authority on residency, unlike exec_state_.req_ids, which lists only the rows scheduled this step. Each entry owns a device allocation, so pruning is part of the change rather than a follow-up.

One detail a reviewer will want checked: the prune is guarded by size() > num_reqs(), and a departure paired with an arrival in the same step can leave the sizes equal with a stale entry present. That entry survives exactly one further step and cannot be read in the meantime, because lookup is by current request id. It is then reclaimed: the resolve loop inserts the newcomer, so the next prune sees size() > num_reqs() and fires. The lag is bounded at one step and is self-correcting; it does not accumulate.

Adopting upstream's paged shape — the draft context as a real KV cache group carried by MultiGroupBlockTable::move_row — is the right end state and is recorded as owed, not attempted. It replaces DflashDeviceKVStore, changes how #1919 sizes the draft budget and how #2007 splits the pools, and would land the concurrency repair behind a rewrite.

The gate

tests/vllm/v1/spec_decode/test_dflash2_concurrency.cpp drives two concurrent requests through the synchronous production front and reproduces the discontinuity on CPU, which is what converts this from something only a GPU run can see.

Its load-bearing leg is not "it did not throw". Marking the moved row disabled — the #1919 fallback, which continues before the invariant — also makes the throw stop and emits identical tokens, because the verify is lossless and a request that stops speculating just runs slower. So the gate asserts that the survivor keeps proposing, and at every step it is alive for.

Mutations

Each rival repair was built on the pre-change code and measured, not argued about.

Mutation	Result
A′ — the fallback repair. On a position mismatch, disable the row instead of asserting.	CAUGHT. Case 1 goes green and the propose count reads none_lines == 9 — the survivor stopped proposing for all nine remaining steps.
B′ — "silence the check." Delete the position invariant, let the reset context stand.	NOT CAUGHT. Filed as #2009.
Reachability. Delete the propose_drafts_block call site in propose_drafts_dflash.	Red, by name.

One leg was written, run, and removed rather than shipped. Comparing the drafted blocks against a solo control is a tautology on this fixture: with the invariant deleted and the context reset at every row move, the synthetic draft still emits 12 12 12 at every step of both runs — seeded-noise weights over a 24-token vocabulary, and the selector walk collapses to one id. Nine passing string comparisons measured nothing. The test file and the spec both record it, because the next person to reach for that leg will find the same fixture. That B′ is consequently ungated is #2009, not a thing this row quietly left out.

What this does not do, and what must be measured before it is trusted

A concurrency change does not merge here on a green compile. The end-to-end ladder is the operator's to run — rungs c = 1, 2, 4, 8, 16 at the #2008 flags. Predictions and falsifiers are in the spec under ## Gates. In short: c=1 must be inert (24.70 out tok/s, TPOT 37.90 ms), no rung may VOID or report failed > 0, and c=2 must not come in below c=1. A rung that lands low but completes is #2007 and the P == 1 capture gate being measured for the first time — this change is what lets a batch reach P > 1 at all — and those are named, owned and out of scope.

Gates run

Gate	Result
cmake --build build -j 4 (CI's CPU configuration)	rc 0, no error: lines
ctest --test-dir build --output-on-failure	rc 0 — 628 tests, 0 failed, 5 pre-existing skips
test_dflash2_concurrency in that suite	Passed, 37.73 s
test_dflash2_runner_reach (sibling DFlash2 gate)	Passed, 45.78 s
scripts/agent-preflight.sh --staged	rc 0, 0 gates failed
check-agent-record.py	rc 0
check-commit-trailers.py / check-commit-style.py	rc 0

Exit codes were captured explicitly rather than read off the tail, which is how the first preflight run's role-undeclared failure was caught behind a clean-looking output. The staged preflight's trailer check initially SKIPPED because origin/main had moved and was no longer an ancestor of HEAD — a diff-scoped gate that skips still exits zero. origin/main is merged (bringing #1999) and both checks were then rerun against the new base.

Found and not fixed

• DFlash2's draft-context position invariant is ungated: deleting it leaves the whole suite green #2009 — the position invariant is ungated; deleting it leaves this row's own gate green. Closing it needs a fixture whose draft is sensitive to its context, which this one is not.
• UNVERIFIED hypothesis, labelled rather than asserted: a prefix-cache hit or a resumed request may trip the same invariant at c=1. The reasoning is that such a request is admitted with num_computed_tokens > 0 and has no draft context for the cache-supplied tokens, so its first propose would read positions[rows[0]] > 0 against L == 0. It was not reproduced. A throwaway probe forced enable_prefix_caching = true and issued the same 20-token prompt twice; both completed, nothing threw, and the engine's own prefix_cache_metrics() reported queries=40 hits=0 — the cache never engaged, so the probe measured nothing about the hypothesis. The fixture's target is a GDN hybrid, the family upstream defaults prefix caching OFF for. Reported as a hypothesis, not a defect; confirming it needs a decoder-only DFlash2-capable target.
• Attention and recurrent state are allocated from two separate pools, so c=32 at k=8 is unservable where vLLM and SGLang serve it #2007 — the two-pool allocation. Noted, not widened into.

Re-merge onto main (b758127ec)

origin/main moved to b758127ec (#1994) and GitHub reported this PR CONFLICTING. Re-merged; head is 2ea92d2fc. The sole conflict under GitHub's merge is .agents/issue-index.md — the merge=union driver resolves it locally and GitHub does not honour that driver. tests/CMakeLists.txt merges clean either way, and no source file conflicted, so no take-both resolution exists on this branch.

.agents/issue-index.md was union-appended and then checked by row-ID set difference rather than by reading the diff: base 730 + ours 2 (#2008, #2009) + main 15 = 747 expected, 747 actual, 0 lost / 0 invented / 0 duplicated, with every merged row byte-identical to a row in one of the three sources.

tests/CMakeLists.txt was rebuilt from origin/main (blob 6a616d781) with this row's edit re-applied after the test_dflash2_ctx_capacity source line, an anchor asserted to occur exactly once. The rebuild is byte-identical to the three-way result; git diff origin/main is 7 added, 0 deleted.

Nothing this row owns moved. Comparing the reviewed pre-merge delta against the post-merge delta, normalised for hunk offsets, the added and removed line sets are identical (792 added, 74 removed) — the only difference is two context lines, where main's #1982 and #1992 index rows now sit adjacent to this row's. Main's delta touches neither runner.h nor runner.cpp, and both, along with test_dflash2_concurrency.cpp, are blob-identical to the reviewed head 3d895a202, so the request-id keying of the draft context is untouched.

Re-merged again onto 2c27f0d57, which carries #2000 and #1997. Head is 53e144f37. #2000 is the first commit on main to touch runner.h and runner.cpp; #1997 touches neither. Both runner files three-way merged with no conflict because the edits are disjoint against the common base b758127ec — #2000 at runner.h 357-474 and runner.cpp 334-823, this row at runner.h 836-870 and runner.cpp 2752-3148 — and the merge is verified in both directions, each side's per-file delta byte-identical to its own source diff at shifted offsets. Brace balance checked arithmetically anyway: runner.h expected 55, merged 55; runner.cpp expected 337, merged 337; both for { and }. Index re-checked against the current base: 745 + 2 ours (#2008, #2009) + 5 theirs (#1963, #1966, #1984, #1985, #2002) = 752 expected, 752 actual, 0 lost / 0 invented / 0 duplicated. tests/CMakeLists.txt rebuilt from blob d4d6cff71 at the same uniqueness-asserted anchor, 7 added / 0 deleted, with #2000's and #1997's own lines intact. The reviewed change is unchanged across all three re-merges: added/removed line sets identical to git diff 2a42cb369 3d895a202 (786/68 excluding file headers, 792/74 including them).

The full ctest suite has NO local verdict on this head, and CI carries it. The box was at 98% disk; a full build drove the root filesystem to 100% with 4.1 GiB free, which is the ENOSPC state that makes unrelated checkers emit false policy refusals, so it was aborted and the space reclaimed. The focused suites were re-run after the abort from a clean rebuild — ctest -R dflash 27/27, test_dflash2_concurrency 2/2 (21 assertions), test_dflash2_runner_reach 8/8, test_dflash2_ctx_capacity 6/6, test_runner 20/20, every one rc=0 — and scripts/agent-preflight.sh --staged is green.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: CLAUDE-CODE:claude-opus-5 [Claude Code]