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

include/vllm/v1/worker/gpu/runner.h

@@ -854,41 +854,76 @@ class GPUModelRunner final : public ModelRunnerBase {
   bool dspark_sample_from_anchor_ = true;
   bool use_dspark() const { return dspark_weights_ != nullptr; }
   // Per-request PERSISTENT context KV store (D9 persistent paged draft-KV — the
-  // perf form of vLLM's incrementally-written draft KV cache). dflash_kv_store_[i]
-  // holds request i's per-layer bf16 context K/V (K normed+RoPE'd, V raw) for its
-  // committed positions 0..L_i-1 (L_i = dflash_ctx_len_[i]). Each verify step
-  // projects ONLY the newly-accepted rows (AppendContextKVHost) and APPENDS them,
-  // instead of re-projecting the whole growing context (the D5/D7 O(context^2)
-  // recompute). Bit-identical to the recompute by per-row projection independence.
-  // dflash_ctx_reqid_[i] tracks the occupant so a reused batch slot resets its
-  // store; rejected drafts' rows are never appended (rollback = don't-append).
-  // Indexed by the runner's condensed-dense batch row. Sized on set_dflash_draft.
-  // D11 A-wire: the store is now the DEVICE-RESIDENT append-only draft-KV store
-  // (DflashDeviceKVStore, opaque, one shared_ptr per condensed-dense batch row).
-  // AppendContextKVDevice keeps the projected bf16 K/V on-device (no D<->H round
-  // trip) and ForwardBlockLogitsWithDeviceKV runs the block forward straight off
-  // the device store — bit-identical to the D9 host path, and the capture-ready
-  // substrate for Parts B/C. shared_ptr-to-incomplete is safe: MakeDeviceKVStore
-  // constructs the control block (with its deleter) in qwen3_dflash.cpp.
-  std::vector<std::shared_ptr<vllm::DflashDeviceKVStore>> dflash_kv_store_;
-  std::vector<int32_t> dflash_ctx_len_;
-  std::vector<std::string> dflash_ctx_reqid_;
-  // #1919: the store's resolved capacity, taken ONCE at set_dflash_draft from
-  // this engine's own max_model_len, and the per-row "this request no longer
-  // fits" flag.
+  // perf form of vLLM's incrementally-written draft KV cache). One entry holds
+  // that request's per-layer bf16 context K/V (K normed+RoPE'd, V raw) for its
+  // committed positions 0..ctx_len-1. Each verify step projects ONLY the
+  // newly-accepted rows (AppendContextKVDeviceRows) and APPENDS them, instead of
+  // re-projecting the whole growing context (the D5/D7 O(context^2) recompute).
+  // Bit-identical to the recompute by per-row projection independence; rejected
+  // drafts' rows are never appended (rollback = don't-append).
+  //
+  // D11 A-wire: the store is the DEVICE-RESIDENT append-only draft-KV store
+  // (DflashDeviceKVStore, opaque). AppendContextKVDevice keeps the projected bf16
+  // K/V on-device (no D<->H round trip) and ForwardBlockLogitsWithDeviceKV runs
+  // the block forward straight off it — bit-identical to the D9 host path, and
+  // the capture-ready substrate for Parts B/C. shared_ptr-to-incomplete is safe:
+  // MakeDeviceKVStore constructs the control block (with its deleter) in
+  // qwen3_dflash.cpp.
+  //
+  // KEYED BY REQUEST ID, NOT BY BATCH ROW (#2008). These three fields were three
+  // arrays indexed by the runner's condensed-dense batch row, with a fourth
+  // recording each row's occupant so a reused slot could reset. 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, and `swap_states`
+  // exchanges two live rows. Both permute every per-slot array they own —
+  // including the block-table rows — and neither knows these exist, because they
+  // live on the runner rather than in `InputBatch`. So the survivor of a
+  // completed pair met the departed request's bookkeeping, the occupant test
+  // read a changed id, the store was reset to EMPTY under a request still using
+  // it, and `propose_drafts_block`'s position invariant then refused. #2008
+  // measured what that costs: DFlash2 served c=1 at 24.70 out tok/s and VOIDed
+  // at c=2 with ok=1, after which every later request on that server came back
+  // `[request submitted to a stopped AsyncLLM]`.
   //
-  // The flag is STICKY for the lifetime of the request occupying the row, and
-  // that is forced rather than chosen. `propose_drafts_block` keeps
-  // `dflash_ctx_len_` in lockstep with the store's `num_ctx` and asserts both
-  // against the target's committed positions; a step that declines to append
-  // breaks that lockstep, so every later step for the same request must decline
-  // too. It is cleared where the store is rebuilt — when a reused dense slot
-  // changes occupant — so a later, shorter request on the same row speculates
-  // normally. Upstream's own skip is monotone in the same way: its
-  // `num_tokens >= max_model_len` condition only ever becomes true
-  // (`vllm/v1/spec_decode/ngram_proposer.py:156-159`).
+  // Upstream keys the same state to the request on both of its paths. Its V2
+  // runner, where DFlash2 lives, has no `condense` at all — a finished request's
+  // slot returns to a free list and stays that request's for its lifetime
+  // (`vllm/v1/worker/gpu/states.py:29,100,132` @ `b389ac2946`) — and every
+  // cross-step speculator tensor is indexed through
+  // `req_state_idx = idx_mapping[req_idx]`
+  // (`vllm/v1/worker/gpu/spec_decode/dflash/speculator.py:536`). Its legacy V1
+  // runner does condense, and there the draft's block-table row moves with the
+  // request (`vllm/v1/worker/gpu_input_batch.py:786` ->
+  // `vllm/v1/worker/block_table.py:367-373`). A request id is the key that
+  // survives any reordering of OUR batch, so every permutation the batch can
+  // perform is a no-op here — which is the property that makes upstream's
+  // speculator indifferent to row order in the first place.
+  //
+  // Entries are pruned each propose against `InputBatch`'s own membership, so a
+  // finished or preempted request releases its device store on the step after it
+  // leaves the batch.
+  struct DflashReqCtx {
+    std::shared_ptr<vllm::DflashDeviceKVStore> store;
+    // Committed context length L. Kept in lockstep with the store's own num_ctx,
+    // and asserted against it every propose (SPEC-DFLASH2 W8, #1838).
+    int32_t ctx_len = 0;
+    // #1919: this request no longer fits the store and runs on the target alone.
+    // STICKY for the request's lifetime, and that is forced rather than chosen:
+    // `propose_drafts_block` keeps `ctx_len` in lockstep with the store's
+    // `num_ctx` and asserts both against the target's committed positions, so a
+    // step that declines to append breaks that lockstep and every later step for
+    // the same request must decline too. Upstream's own skip is monotone in the
+    // same way: its `num_tokens >= max_model_len` condition only ever becomes
+    // true (`vllm/v1/spec_decode/ngram_proposer.py:156-159`). Being a property
+    // of the REQUEST — which the row-indexed form could only approximate, and
+    // its comment already claimed — it now simply ends with the request.
+    bool disabled = false;
+  };
+  std::unordered_map<std::string, DflashReqCtx> dflash_ctx_;
+  // #1919: the store's resolved capacity, taken ONCE at set_dflash_draft from
+  // this engine's own max_model_len. The "no longer fits" flag it pairs with is
+  // `DflashReqCtx::disabled` above.
   vllm::Qwen3DFlashModel::DflashCtxStoreSizing dflash_ctx_sizing_;
-  std::vector<bool> dflash_ctx_disabled_;
   // Draft KV cache (`fa_draft` group) backing storage, owned by the runner and
   // allocated in initialize_kv_cache when spec is on. draft_attn_kv_ (declared
   // above) views into these buffers. Empty on the default path.

src/vllm/v1/worker/gpu/runner.cpp

@@ -2855,10 +2855,7 @@ void GPUModelRunner::set_dflash_draft(const vllm::Qwen3DFlashWeights* weights,
       dflash_tap_layer_ids_.push_back(id.get<int32_t>());
     }
   }
-  dflash_kv_store_.clear();
-  dflash_ctx_len_.clear();
-  dflash_ctx_reqid_.clear();
-  dflash_ctx_disabled_.clear();
+  dflash_ctx_.clear();
 
   // #1919: resolve the draft context store's capacity from THIS engine's
   // advertised context, and say so. Before this, the capacity was a
@@ -2957,11 +2954,38 @@ void GPUModelRunner::propose_drafts_block(
   const int num_mask_rows = num_query_per_req - 1;
   const StepInputs& step = exec_state_.step;
 
-  if (static_cast<int>(dflash_ctx_len_.size()) < num_reqs) {
-    dflash_kv_store_.resize(static_cast<size_t>(num_reqs));
-    dflash_ctx_len_.resize(static_cast<size_t>(num_reqs), 0);
-    dflash_ctx_reqid_.resize(static_cast<size_t>(num_reqs));
-    dflash_ctx_disabled_.resize(static_cast<size_t>(num_reqs), false);
+  // #2008 — THE DRAFT CONTEXT BELONGS TO THE REQUEST, NOT TO THE BATCH ROW.
+  //
+  // Release the entries of requests that have left the batch, then resolve one
+  // pointer per row for the rest of this function. `InputBatch` is the authority
+  // on residency, not `exec_state_.req_ids`, which lists only the rows SCHEDULED
+  // this step; a resident request that happens not to be scheduled must keep its
+  // context. The scan is bounded by `max_num_seqs` and runs once per step, off
+  // the per-token path.
+  //
+  // The pruning is not housekeeping: each entry owns a device allocation sized
+  // from `dflash_ctx_sizing_.slots`, so a leaked entry is leaked device memory.
+  if (dflash_ctx_.size() > static_cast<size_t>(input_batch_.num_reqs())) {
+    for (auto it = dflash_ctx_.begin(); it != dflash_ctx_.end();) {
+      it = (input_batch_.req_id_to_index.count(it->first) == 0)
+               ? dflash_ctx_.erase(it)
+               : std::next(it);
+    }
+  }
+  // Resolving here replaces the reused-slot test this loop used to open with.
+  // "Has this row's occupant changed" was only ever a proxy for "is this state
+  // this request's", and the proxy is what #2008 broke: after a condense move
+  // the answer was yes for a request whose context was perfectly valid, and the
+  // reset threw it away. Keyed by id the question cannot be asked wrongly — a
+  // first sight of a request id constructs its entry, and THAT is the reset.
+  std::vector<DflashReqCtx*> row_ctx(static_cast<size_t>(num_reqs), nullptr);
+  for (int i = 0; i < num_reqs; ++i) {
+    DflashReqCtx& c = dflash_ctx_[exec_state_.req_ids[static_cast<size_t>(i)]];
+    if (c.store == nullptr) {
+      c.store = Qwen3DFlashModel::MakeDeviceKVStore(config, queue_,
+                                                    dflash_ctx_sizing_.slots);
+    }
+    row_ctx[static_cast<size_t>(i)] = &c;
   }
 
   // SPEC-DFLASH2 W8 (#1838): the propose pre-phase timer. Before W8 everything
@@ -2998,29 +3022,21 @@ void GPUModelRunner::propose_drafts_block(
   std::vector<bool> fell_back(static_cast<size_t>(num_reqs), false);
 
   for (int i = 0; i < num_reqs; ++i) {
-    // Reset a reused dense slot (a new request now occupies this row).
-    if (dflash_ctx_reqid_[static_cast<size_t>(i)] !=
-        exec_state_.req_ids[static_cast<size_t>(i)]) {
-      dflash_kv_store_[static_cast<size_t>(i)] = Qwen3DFlashModel::MakeDeviceKVStore(
-          config, queue_, dflash_ctx_sizing_.slots);
-      dflash_ctx_len_[static_cast<size_t>(i)] = 0;
-      dflash_ctx_reqid_[static_cast<size_t>(i)] =
-          exec_state_.req_ids[static_cast<size_t>(i)];
-      // A new occupant starts speculating again (#1919). The flag is a property
-      // of the REQUEST, not of the row.
-      dflash_ctx_disabled_[static_cast<size_t>(i)] = false;
-    }
+    // #2008: no reused-slot reset here any more. This request's state was
+    // resolved by id above, so it is this request's whether or not the batch
+    // moved it, and a request seen for the first time started empty.
+    DflashReqCtx& ctx = *row_ctx[static_cast<size_t>(i)];
     // #1919: a request whose context has outgrown the store stops speculating
     // for the rest of its life, and this test comes BEFORE the two invariants
     // below because a disabled row stops maintaining both: it neither appends
-    // nor advances `dflash_ctx_len_`, so its counter and the target's committed
+    // nor advances its `ctx_len`, so that counter and the target's committed
     // positions legitimately diverge from here on.
     //
     // What such a row PROPOSES is decided in section 4, which is also where the
     // two scheduling modes part company; `fell_back` is how this loop tells it
     // which rows are in that state and are decoding rather than still
     // prefilling.
-    if (dflash_ctx_disabled_[static_cast<size_t>(i)]) {
+    if (ctx.disabled) {
       fell_back[static_cast<size_t>(i)] =
           !(i < static_cast<int>(exec_state_.discard.size()) &&
             exec_state_.discard[static_cast<size_t>(i)]);
@@ -3045,18 +3061,18 @@ void GPUModelRunner::propose_drafts_block(
     // Invariant: this step's first committed token sits at absolute position L
     // (== current context length). A violation means the accumulation lost sync
     // (the I5e async-input-combine bug class) — assert rather than corrupt.
-    const int64_t L = dflash_ctx_len_[static_cast<size_t>(i)];
+    const int64_t L = ctx.ctx_len;
     VT_CHECK(step.positions[static_cast<size_t>(rows[0])] == L,
              "propose_drafts_block: context position discontinuity (accumulation "
              "out of sync with the target's committed positions)");
     // SPEC-DFLASH2 W8 (#1838): the runner's counter and the DEVICE store must
     // agree, or the append is dead. The W8 mutation run proved the check above
-    // cannot see that state: with the append call deleted, `dflash_ctx_len_`
+    // cannot see that state: with the append call deleted, `ctx_len`
     // kept advancing, the store stayed empty, every propose ran CONTEXT-FREE,
     // and every gate stayed green — well-formed drafts, lossless verify, only
     // ACCEPTANCE falls, the exact invisible-defect class this row exists to
     // remove. This host integer comparison is what makes that state loud.
-    VT_CHECK(L == Qwen3DFlashModel::DeviceKVNumCtx(*dflash_kv_store_[static_cast<size_t>(i)]),
+    VT_CHECK(L == Qwen3DFlashModel::DeviceKVNumCtx(*ctx.store),
              "propose_drafts_block: the runner's context length and the device "
              "store's num_ctx disagree — the context-KV append is dead or "
              "double-run (SPEC-DFLASH2 W8, #1838)");
@@ -3087,10 +3103,9 @@ void GPUModelRunner::propose_drafts_block(
     // worth of context and keeps the W11 paged fast route, which
     // `ClassifyDflashBlockAttn` would otherwise drop below its capacity
     // conjunct.
-    const int64_t capacity = Qwen3DFlashModel::DeviceKVCapacity(
-        *dflash_kv_store_[static_cast<size_t>(i)]);
+    const int64_t capacity = Qwen3DFlashModel::DeviceKVCapacity(*ctx.store);
     if (L + append + num_query_per_req > capacity) {
-      dflash_ctx_disabled_[static_cast<size_t>(i)] = true;
+      ctx.disabled = true;
       std::cerr << "vllm.cpp: request " << exec_state_.req_ids[static_cast<size_t>(i)]
                 << " has outgrown the draft speculative context (" << (L + append)
                 << " context tokens plus a " << num_query_per_req
@@ -3111,10 +3126,9 @@ void GPUModelRunner::propose_drafts_block(
       new_rows[static_cast<size_t>(j)] = static_cast<int32_t>(rows[static_cast<size_t>(j)]);
       new_pos[static_cast<size_t>(j)] = static_cast<int32_t>(L + j);
     }
-    Qwen3DFlashModel::AppendContextKVDeviceRows(*dflash_kv_store_[static_cast<size_t>(i)],
-                                                combined.tensor, new_rows, new_pos, backbone,
-                                                config, queue_);
-    dflash_ctx_len_[static_cast<size_t>(i)] = static_cast<int32_t>(L + append);
+    Qwen3DFlashModel::AppendContextKVDeviceRows(*ctx.store, combined.tensor, new_rows,
+                                                new_pos, backbone, config, queue_);
+    ctx.ctx_len = static_cast<int32_t>(L + append);
 
     // A discarded (still-prefilling chunk) row commits its chunk's features but
     // proposes no draft — it has no valid last_sampled anchor yet.
@@ -3124,7 +3138,7 @@ void GPUModelRunner::propose_drafts_block(
 
     // Block: anchor = last_sampled (the bonus/last committed token, re-embedded),
     // then k mask tokens; positions L' .. L'+k (L' = the new context length).
-    const int64_t Lp = dflash_ctx_len_[static_cast<size_t>(i)];
+    const int64_t Lp = ctx.ctx_len;
     const int32_t anchor = input_batch_.last_sampled_tokens[static_cast<size_t>(i)];
     blk_ids.push_back(anchor);
     blk_pos.push_back(static_cast<int32_t>(Lp));
@@ -3251,7 +3265,7 @@ void GPUModelRunner::propose_drafts_block(
     int64_t total_ctx = 0;
     for (int r = 0; r < P; ++r) {
       vllm::DflashDeviceKVStore* st =
-          dflash_kv_store_[static_cast<size_t>(propose_rows[static_cast<size_t>(r)])].get();
+          row_ctx[static_cast<size_t>(propose_rows[static_cast<size_t>(r)])]->store.get();
       stores.push_back(st);
       total_ctx += Qwen3DFlashModel::DeviceKVNumCtx(*st);
       ctx_cu.push_back(static_cast<int32_t>(total_ctx));

tests/CMakeLists.txt

@@ -1678,6 +1678,13 @@ target_include_directories(test_dflash2_embed_dedup_reach PRIVATE ${CMAKE_SOURCE
 # above 4096, which no other binary drives.
 vllm_cpp_add_test(test_dflash2_ctx_capacity
   vllm/v1/spec_decode/test_dflash2_ctx_capacity.cpp)
+# SPEC-DFLASH2 (#2008): the draft context must follow the REQUEST, not the batch
+# ROW. Its OWN binary for the reason every fixture binary here has one --
+# `VT_SPEC_TRACE` is latched once per process by a function-local static on the
+# first propose, and this one needs level 1 with `max_num_seqs` above 1, which no
+# other binary drives.
+vllm_cpp_add_test(test_dflash2_concurrency
+  vllm/v1/spec_decode/test_dflash2_concurrency.cpp)
 target_include_directories(test_dflash_propose PRIVATE ${CMAKE_SOURCE_DIR}/src)
 # SPEC-NGRAM (ROAD-V1-D3) — the draft-free n-gram matcher unit gate (ports
 # vllm/tests/v1/spec_decode/test_ngram.py, host-side, runs everywhere).