Revert "[MRV2] Always build attn metadata at capture time" (#49364) (#49451)
Co-authored-by: vllm-agent CI bot <ci-bot@vllm-agent.local>
vllm-agent committed
Jul 23, 2026 at 09:51 UTC
494845e79f8ae40f356902ba8ac181bf545603b5
2 files changed
+21
-35
vllm/v1/worker/gpu/cudagraph_utils.py
+17
-34
@@ -498,7 +498,10 @@ class ModelCudaGraphManager(CudaGraphManager):
498
block_tables,
499
attn_groups,
500
kv_cache_config,
501
- full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
501
+ skip_attn=(
502
+ desc.cg_mode == CUDAGraphMode.PIECEWISE
503
+ and not self.use_breakable_cg
504
+ ),
505
)
506
507
# Capture with dummy rows marked as padding.
@@ -507,6 +510,7 @@ class ModelCudaGraphManager(CudaGraphManager):
510
def forward_fn(cg_mode: CUDAGraphMode) -> None:
511
batch_descriptor = None
512
if cg_mode == CUDAGraphMode.PIECEWISE:
513
+ assert (attn_metadata is not None) == self.use_breakable_cg
514
batch_descriptor = BatchDescriptor(
515
num_tokens=num_tokens,
516
has_lora=has_lora,
@@ -589,7 +593,7 @@ def prepare_inputs_to_capture(
593
block_tables: BlockTables,
594
attn_groups: list[list[AttentionGroup]],
595
kv_cache_config: KVCacheConfig,
592
- full_cudagraph: bool,
596
+ skip_attn: bool = False,
597
) -> AttentionState:
598
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
599
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
@@ -610,36 +614,15 @@ def prepare_inputs_to_capture(
614
)
615
input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
616
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
- )
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
+ )
628
return AttentionState(attn_metadata, slot_mappings_by_layer)
vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py
+4
-1
@@ -56,7 +56,10 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
56
block_tables,
57
attn_groups,
58
kv_cache_config,
59
- full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
59
+ skip_attn=(
60
+ desc.cg_mode == CUDAGraphMode.PIECEWISE
61
+ and not self.use_breakable_cg
62
+ ),
63
)
64
65
return lambda cg_mode: forward_fn(