[DFlash] Capture the context K/V precompute in the draft CUDA graph (#57632)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Juntian Liu committed
Sep 24, 2026 at 14:38 UTC
04730e82700d0c9b957cda188dd1bbc8073e1fcd
2 files changed
+62
-30
vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+16
-8
@@ -84,8 +84,12 @@ class DFlashCudaGraphManager(CudaGraphManager):
84
kv_cache_config: KVCacheConfig,
85
max_model_len: int,
86
causal: bool | Mapping[int, bool],
87
+ precompute_context_kv: Callable[[int], None],
88
progress_bar_desc: str = "Capturing CUDA graphs",
89
) -> None:
90
+ """``precompute_context_kv(num_reqs)`` is captured ahead of the query
91
+ forward."""
92
+
93
def create_forward_fn(
94
desc: BatchExecutionDescriptor,
95
warmup: bool,
@@ -110,13 +114,17 @@ class DFlashCudaGraphManager(CudaGraphManager):
114
)
115
attn_metadata, slot_mappings = attn_state
116
113
- return lambda cg_mode: forward_fn(
114
- num_reqs,
115
- num_tokens,
116
- attn_metadata,
117
- slot_mappings,
118
- num_tokens_across_dp,
119
- cg_mode,
120
- )
117
+ def forward(cg_mode: CUDAGraphMode) -> None:
118
+ precompute_context_kv(num_reqs)
119
+ forward_fn(
120
+ num_reqs,
121
+ num_tokens,
122
+ attn_metadata,
123
+ slot_mappings,
124
+ num_tokens_across_dp,
125
+ cg_mode,
126
+ )
127
+
128
+ return forward
129
130
super().capture(create_forward_fn, progress_bar_desc)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+46
-22
@@ -79,9 +79,7 @@ class DFlashSpeculator(DraftModelSpeculator):
79
)
80
self.sample_from_anchor = False
81
82
- # Context positions for the K/V precompute. Populated by
83
- # prepare_dflash_inputs, and processed by the model's
84
- # precompute_and_store_context_kv method. NOT captured by CUDA graphs.
82
+ # Context positions for the K/V precompute, populated by prepare_dflash_inputs.
83
self.context_positions = torch.zeros(
84
self.max_num_tokens, dtype=torch.int64, device=device
85
)
@@ -151,6 +149,8 @@ class DFlashSpeculator(DraftModelSpeculator):
149
self.sample_indices.zero_()
150
self.sample_pos.zero_()
151
self.sample_idx_mapping.fill_(-1)
152
+ # Capture must not write context K/V.
153
+ self._context_slot_mappings.fill_(PAD_SLOT_ID)
154
assert self.query_cudagraph_manager is not None
155
self.query_cudagraph_manager.capture(
156
self._generate_draft,
@@ -160,6 +160,9 @@ class DFlashSpeculator(DraftModelSpeculator):
160
self.kv_cache_config,
161
self.max_model_len,
162
causal=self._group_causal,
163
+ precompute_context_kv=lambda num_reqs: self._precompute_context_kv(
164
+ 0, self._num_graph_context_tokens(num_reqs)
165
+ ),
166
progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs",
167
)
168
@@ -297,6 +300,28 @@ class DFlashSpeculator(DraftModelSpeculator):
300
num_reqs, self.num_speculative_steps
301
)
302
303
+ def _num_graph_context_tokens(self, num_reqs: int) -> int:
304
+ # Context rows a captured draft step stores: one full verify per request.
305
+ return min(num_reqs * (self.num_speculative_steps + 1), self.max_num_tokens)
306
+
307
+ def _precompute_context_kv(
308
+ self, start: int, end: int, dummy_run: bool = False
309
+ ) -> None:
310
+ if dummy_run:
311
+ context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
312
+ elif self._layer_group_idx is not None:
313
+ context_slots = [
314
+ self._context_slot_mappings[gidx][start:end]
315
+ for gidx in self._layer_group_idx
316
+ ]
317
+ else:
318
+ context_slots = self._context_slot_mappings[0][start:end]
319
+ self.model.precompute_and_store_context_kv(
320
+ self.hidden_states[start:end],
321
+ self.context_positions[start:end],
322
+ context_slots,
323
+ )
324
+
325
@torch.inference_mode()
326
def propose(
327
self,
@@ -407,25 +432,6 @@ class DFlashSpeculator(DraftModelSpeculator):
432
self.sample_from_anchor,
433
)
434
410
- # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph
411
- # because the context shape varies per step. During dummy runs the block tables
412
- # are placeholders, so we skip the cache write to avoid clobbering real entries.
413
- # Each layer uses the context slots of its own kv-cache group.
414
- if dummy_run:
415
- context_slots: torch.Tensor | list[torch.Tensor | None] | None = None
416
- elif self._layer_group_idx is not None:
417
- context_slots = [
418
- self._context_slot_mappings[gidx][:num_target_tokens]
419
- for gidx in self._layer_group_idx
420
- ]
421
- else:
422
- context_slots = self._context_slot_mappings[0][:num_target_tokens]
423
- self.model.precompute_and_store_context_kv(
424
- self.hidden_states[:num_target_tokens],
425
- self.context_positions[:num_target_tokens],
426
- context_slots,
427
- )
428
-
435
batch_sync, num_batch_tokens = (
436
self._build_uniform_batch_dp_sync(dp_sync, num_reqs, self.num_query_per_req)
437
if dp_sync is not None
@@ -447,6 +453,24 @@ class DFlashSpeculator(DraftModelSpeculator):
453
batch_sync.num_tokens_across_dp if batch_sync is not None else None
454
)
455
456
+ if batch_desc.cg_mode == CUDAGraphMode.FULL:
457
+ # The graph stores the first num_context context rows.
458
+ assert batch_desc.num_reqs is not None
459
+ num_context = self._num_graph_context_tokens(batch_desc.num_reqs)
460
+ if dummy_run:
461
+ # Dummy block tables are placeholders: write no context K/V.
462
+ self._context_slot_mappings[:, :num_context].fill_(PAD_SLOT_ID)
463
+ elif num_target_tokens <= num_context:
464
+ # Rows past the batch keep stale positions but write no K/V.
465
+ self._context_slot_mappings[:, num_target_tokens:num_context].fill_(
466
+ PAD_SLOT_ID
467
+ )
468
+ else:
469
+ # Prefill context beyond the graph's rows is stored before replay.
470
+ self._precompute_context_kv(num_context, num_target_tokens)
471
+ else:
472
+ self._precompute_context_kv(0, num_target_tokens, dummy_run)
473
+
474
# Rebuild the draft attention metadata even when replaying the FULL
475
# graph so that any attention metadata builder state is updated.
476
draft_attn_metadata = self._build_uniform_attn_metadata(