diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 26f331fe84..3a174cba80 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -498,10 +498,7 @@ class ModelCudaGraphManager(CudaGraphManager): block_tables, attn_groups, kv_cache_config, - skip_attn=( - desc.cg_mode == CUDAGraphMode.PIECEWISE - and not self.use_breakable_cg - ), + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, ) # Capture with dummy rows marked as padding. @@ -510,7 +507,6 @@ class ModelCudaGraphManager(CudaGraphManager): def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None if cg_mode == CUDAGraphMode.PIECEWISE: - assert (attn_metadata is not None) == self.use_breakable_cg batch_descriptor = BatchDescriptor( num_tokens=num_tokens, has_lora=has_lora, @@ -593,7 +589,7 @@ def prepare_inputs_to_capture( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, - skip_attn: bool = False, + full_cudagraph: bool, ) -> AttentionState: input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) @@ -614,15 +610,36 @@ def prepare_inputs_to_capture( ) input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs] - attn_metadata = None - if not skip_attn: - attn_metadata = model_state.prepare_attn( - input_batch, - CUDAGraphMode.NONE, - input_block_tables, - slot_mappings, - attn_groups, - kv_cache_config, - for_capture=True, - ) + # NOTE(woosuk): Attention metadata is required not just by standard attention + # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv, + # DSV4 compressor), which maintain their own states and require special metadata + # such as block tables. + # During CUDA graph capture: + # - For FULL CUDA graphs: We set for_capture=True so that both attention and + # attention-like ops produce capturable metadata compatible with CUDA graphs. + # - For PIECEWISE CUDA graphs: We still build attention metadata, but set + # for_capture=False. This is because: + # * Attention-like ops (such as sconv or DSV4 compressor) may not be used as + # breakpoints in PIECEWISE CUDA graphs, so we must generate their attention + # metadata so they can execute and be captured during graph capture. + # * Standard attention ops that are treated as breakpoints will be executed + # eagerly at capture time (not included in the graph itself), and for these, + # setting for_capture=False is essential. Some attention backends + # (like linear attention) cannot generate capturable metadata for prefill, + # so for_capture=False ensures they execute without issue. + # * We assume that attention-like operations intended for capture will still + # produce capturable metadata, even when for_capture=False. While this + # assumption is brittle, it currently works in practice. + # In summary: We always generate attention metadata for both FULL and PIECEWISE + # CUDA graphs, setting for_capture=True for FULL graphs, and for_capture=False + # for PIECEWISE graphs, to ensure correct execution and capture. + attn_metadata = model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + input_block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=full_cudagraph, + ) return AttentionState(attn_metadata, slot_mappings_by_layer) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py index 19919043c8..ef3b6e2ed5 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py @@ -56,10 +56,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager): block_tables, attn_groups, kv_cache_config, - skip_attn=( - desc.cg_mode == CUDAGraphMode.PIECEWISE - and not self.use_breakable_cg - ), + full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL, ) return lambda cg_mode: forward_fn(