[Bugfix] Skip generic attention during the PIECEWISE capture warmup too

skip_attention was keyed on the runtime cg_mode argument passed to
forward_fn. CudaGraphManager.capture() runs a warmup pass
forward_fn(CUDAGraphMode.NONE) before the PIECEWISE capture, so during
warmup skip_attention evaluated to False and generic FlashAttention ran
eagerly against the capture-time dummy inputs -- crashing on FA3/Hopper
(CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS) at the first
"Capturing CUDA graphs (PIECEWISE)" step, which is still seen on this PR.

Key skip_attention on the descriptor's target mode (desc.cg_mode) instead,
matching how the capture metadata is already built
(full_cudagraph=desc.cg_mode == FULL), so both the warmup and capture
passes of a PIECEWISE desc skip generic attention.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -16,6 +16,8 @@ from vllm.config import (
     VllmConfig,
 )
 from vllm.distributed.device_communicators import pynccl_allocator
+from vllm.forward_context import ForwardContext, override_forward_context
+from vllm.model_executor.layers.attention.attention import get_attention_context
 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 
@@ -46,6 +48,29 @@ def _create_vllm_config() -> MagicMock:
     return vllm_config
 
 
+def test_skip_attention_preserves_raw_metadata():
+    layer_name = "model.layers.0.self_attn"
+    attn_metadata = object()
+    attn_layer = SimpleNamespace(kv_cache=object())
+    slot_mapping = object()
+    forward_context = ForwardContext(
+        no_compile_layers={layer_name: attn_layer},
+        attn_metadata={layer_name: attn_metadata},
+        slot_mapping={layer_name: slot_mapping},
+        skip_attention=True,
+    )
+
+    with override_forward_context(forward_context):
+        resolved_metadata, resolved_layer, _, resolved_slot_mapping = (
+            get_attention_context(layer_name)
+        )
+
+    assert forward_context.attn_metadata[layer_name] is attn_metadata
+    assert resolved_metadata is None
+    assert resolved_layer is attn_layer
+    assert resolved_slot_mapping is slot_mapping
+
+
 def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
     """FULL capture must set graph_pool_id before entering torch.cuda.graph().
 

vllm/forward_context.py

@@ -158,6 +158,10 @@ class ForwardContext:
     # If True, bypass the compiled model call, e.g. by using .forward() directly
     skip_compiled: bool = False
 
+    # If True, generic attention ops use their profiling path while the full
+    # metadata dictionary remains available to model-specific attention-like ops.
+    skip_attention: bool = False
+
     # For torch.compile cold start times, we need to avoid hard-coding
     # any strings into the graph. Right now, the vllm.moe_forward
     # and vllm.moe_forward_shared custom operators hard-code strings into
@@ -220,6 +224,7 @@ def create_forward_context(
     additional_kwargs: dict[str, Any] | None = None,
     skip_compiled: bool = False,
     is_padding: torch.Tensor | None = None,
+    skip_attention: bool = False,
 ):
     if vllm_config.compilation_config.fast_moe_cold_start:
         all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
@@ -236,6 +241,7 @@ def create_forward_context(
         batch_descriptor=batch_descriptor,
         ubatch_slices=ubatch_slices,
         skip_compiled=skip_compiled,
+        skip_attention=skip_attention,
         additional_kwargs=additional_kwargs or {},
         is_padding=is_padding,
     )
@@ -268,6 +274,7 @@ def set_forward_context(
     slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
     skip_compiled: bool = False,
     is_padding: torch.Tensor | None = None,
+    skip_attention: bool = False,
 ):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
@@ -337,6 +344,7 @@ def set_forward_context(
         additional_kwargs,
         skip_compiled,
         is_padding=is_padding,
+        skip_attention=skip_attention,
     )
 
     try:

vllm/model_executor/layers/attention/attention.py

@@ -752,8 +752,10 @@ def get_attention_context(
     """
     forward_context: ForwardContext = get_forward_context()
     attn_metadata_raw = forward_context.attn_metadata
-    attn_metadata: AttentionMetadata
-    if isinstance(attn_metadata_raw, dict):
+    attn_metadata: AttentionMetadata | None
+    if forward_context.skip_attention:
+        attn_metadata = None
+    elif isinstance(attn_metadata_raw, dict):
         attn_metadata = attn_metadata_raw[layer_name]
     elif isinstance(attn_metadata_raw, list):
         # list[dict[str, AttentionMetadata]]: used in speculative decoding

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -521,6 +521,10 @@ class ModelCudaGraphManager(CudaGraphManager):
                     slot_mapping=slot_mappings,
                     batch_descriptor=batch_descriptor,
                     is_padding=input_buffers.is_padding[:num_tokens],
+                    skip_attention=(
+                        desc.cg_mode == CUDAGraphMode.PIECEWISE
+                        and not self.use_breakable_cg
+                    ),
                 ):
                     if cg_mode == CUDAGraphMode.PIECEWISE:
                         # PIECEWISE graph (compiled PW or breakable, chosen inside

vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py

@@ -66,6 +66,10 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
                 slot_mappings,
                 num_tokens_across_dp,
                 cg_mode,
+                skip_attention=(
+                    desc.cg_mode == CUDAGraphMode.PIECEWISE
+                    and not self.use_breakable_cg
+                ),
             )
 
         super().capture(create_forward_fn, progress_bar_desc)

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -282,6 +282,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         num_tokens_across_dp: torch.Tensor | None,
         cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        skip_attention: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
         with set_forward_context(
@@ -292,6 +293,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             num_tokens_across_dp=num_tokens_across_dp,
             slot_mapping=slot_mappings,
             batch_descriptor=batch_descriptor,
+            skip_attention=skip_attention,
         ):
             inputs_embeds = None
             if self.supports_mm_inputs:
@@ -341,6 +343,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         num_tokens_across_dp: torch.Tensor | None,
         cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
+        skip_attention: bool = False,
     ) -> None:
         last_token_indices = self.last_token_indices[:num_reqs]
         positions = self.input_buffers.positions[last_token_indices]
@@ -353,6 +356,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             num_tokens_across_dp=num_tokens_across_dp,
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
+            skip_attention=skip_attention,
         )
         sample_hidden_states = last_hidden_states[last_token_indices]
 
@@ -431,6 +435,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         slot_mappings: dict[str, torch.Tensor] | None,
         num_tokens_across_dp: torch.Tensor | None,
         cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        skip_attention: bool = False,
     ) -> None:
         self._prepare_eplb_forward(num_reqs)
 
@@ -443,6 +448,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             slot_mappings,
             num_tokens_across_dp,
             cudagraph_runtime_mode,
+            skip_attention=skip_attention,
         )
         last_hidden_states = last_hidden_states[:num_reqs]
 

vllm/v1/worker/gpu_ubatch_wrapper.py

@@ -353,6 +353,7 @@ class UBatchWrapper:
         dp_metadata,
         batch_descriptor,
         cudagraph_runtime_mode,
+        skip_attention,
     ) -> list[UbatchMetadata]:
         # Create one forward context per ubatch
         forward_contexts = []
@@ -368,6 +369,7 @@ class UBatchWrapper:
                     batch_descriptor=batch_descriptor,
                     cudagraph_runtime_mode=cudagraph_runtime_mode,
                     slot_mapping=slot_mapping[i] if has_slot_mapping else None,
+                    skip_attention=skip_attention,
                 )
             )
 
@@ -506,6 +508,7 @@ class UBatchWrapper:
                 dp_metadata=ubatch_dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE,
+                skip_attention=forward_context.skip_attention,
             )
             with self.sm_control:
                 return self._capture_ubatches(ubatch_metadata, self.runnable)
@@ -532,6 +535,7 @@ class UBatchWrapper:
                 dp_metadata=ubatch_dp_metadata,
                 batch_descriptor=batch_descriptor,
                 cudagraph_runtime_mode=CUDAGraphMode.NONE,
+                skip_attention=forward_context.skip_attention,
             )
             with self.sm_control:
                 return self._run_ubatches(ubatch_metadata, self.runnable)