[MRV2] Always build attn metadata at capture time (#49364) (#49995)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Nick Hill committed
Jul 27, 2026 at 10:25 UTC
8112b6c9972067012e08ce4f0dbf1d1d9906dbe3
2 files changed
+35
-21
vllm/v1/worker/gpu/cudagraph_utils.py
+34
-17
@@ -498,10 +498,7 @@ class ModelCudaGraphManager(CudaGraphManager):
498
block_tables,
499
attn_groups,
500
kv_cache_config,
501
- skip_attn=(
502
- desc.cg_mode == CUDAGraphMode.PIECEWISE
503
- and not self.use_breakable_cg
504
- ),
501
+ full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
502
)
503
504
# Capture with dummy rows marked as padding.
@@ -510,7 +507,6 @@ class ModelCudaGraphManager(CudaGraphManager):
507
def forward_fn(cg_mode: CUDAGraphMode) -> None:
508
batch_descriptor = None
509
if cg_mode == CUDAGraphMode.PIECEWISE:
513
- assert (attn_metadata is not None) == self.use_breakable_cg
510
batch_descriptor = BatchDescriptor(
511
num_tokens=num_tokens,
512
has_lora=has_lora,
@@ -593,7 +589,7 @@ def prepare_inputs_to_capture(
589
block_tables: BlockTables,
590
attn_groups: list[list[AttentionGroup]],
591
kv_cache_config: KVCacheConfig,
596
- skip_attn: bool = False,
592
+ full_cudagraph: bool,
593
) -> AttentionState:
594
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
595
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
@@ -614,15 +610,36 @@ def prepare_inputs_to_capture(
610
)
611
input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
612
617
- attn_metadata = None
618
- if not skip_attn:
619
- attn_metadata = model_state.prepare_attn(
620
- input_batch,
621
- CUDAGraphMode.NONE,
622
- input_block_tables,
623
- slot_mappings,
624
- attn_groups,
625
- kv_cache_config,
626
- for_capture=True,
627
- )
613
+ # NOTE(woosuk): Attention metadata is required not just by standard attention
614
+ # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
615
+ # DSV4 compressor), which maintain their own states and require special metadata
616
+ # such as block tables.
617
+ # During CUDA graph capture:
618
+ # - For FULL CUDA graphs: We set for_capture=True so that both attention and
619
+ # attention-like ops produce capturable metadata compatible with CUDA graphs.
620
+ # - For PIECEWISE CUDA graphs: We still build attention metadata, but set
621
+ # for_capture=False. This is because:
622
+ # * Attention-like ops (such as sconv or DSV4 compressor) may not be used as
623
+ # breakpoints in PIECEWISE CUDA graphs, so we must generate their attention
624
+ # metadata so they can execute and be captured during graph capture.
625
+ # * Standard attention ops that are treated as breakpoints will be executed
626
+ # eagerly at capture time (not included in the graph itself), and for these,
627
+ # setting for_capture=False is essential. Some attention backends
628
+ # (like linear attention) cannot generate capturable metadata for prefill,
629
+ # so for_capture=False ensures they execute without issue.
630
+ # * We assume that attention-like operations intended for capture will still
631
+ # produce capturable metadata, even when for_capture=False. While this
632
+ # assumption is brittle, it currently works in practice.
633
+ # In summary: We always generate attention metadata for both FULL and PIECEWISE
634
+ # CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False
635
+ # for PIECEWISE graphs, to ensure correct execution and capture.
636
+ attn_metadata = model_state.prepare_attn(
637
+ input_batch,
638
+ CUDAGraphMode.NONE,
639
+ input_block_tables,
640
+ slot_mappings,
641
+ attn_groups,
642
+ kv_cache_config,
643
+ for_capture=full_cudagraph,
644
+ )
645
return AttentionState(attn_metadata, slot_mappings_by_layer)
vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
+1
-4
@@ -56,10 +56,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
56
block_tables,
57
attn_groups,
58
kv_cache_config,
59
- skip_attn=(
60
- desc.cg_mode == CUDAGraphMode.PIECEWISE
61
- and not self.use_breakable_cg
62
- ),
59
+ full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
60
)
61
62
return lambda cg_mode: forward_fn(