[Model] Add PW CUDA graph support for Inkling [2/N] (#48822)

Signed-off-by: Woosuk Kwon <[email protected]>
Co-authored-by: Bugen Zhao <[email protected]>
Co-authored-by: Giancarlo Delfin <[email protected]>
Co-authored-by: Isotr0py <[email protected]>
Co-authored-by: Isotr0py <[email protected]>
Co-authored-by: Jee Jee Li <[email protected]>
Co-authored-by: Roger Wang <[email protected]>
Co-authored-by: Yifan Qiao <[email protected]>
Co-authored-by: Claude Fable 5 <[email protected]>
Co-authored-by: OpenAI Codex <[email protected]>

tests/models/inkling/test_contract_validation.py

@@ -34,17 +34,17 @@ def test_inkling_raw_2d_audio_is_rejected_as_ambiguous():
         parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32))
 
 
-def test_inkling_supports_full_decode_only_cudagraphs():
+def test_inkling_supports_piecewise_cudagraphs():
     support = InklingSconvMetadataBuilder.get_cudagraph_support
-    assert support(None, None) == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
+    assert support(None, None) == AttentionCGSupport.UNIFORM_BATCH
 
     compilation_config = CompilationConfig(
-        cudagraph_mode=CUDAGraphMode.FULL,
+        cudagraph_mode=CUDAGraphMode.PIECEWISE,
         splitting_ops=[],
     )
     resolved_mode = compilation_config.resolve_cudagraph_mode_and_sizes(
-        AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE,
+        AttentionCGSupport.UNIFORM_BATCH,
         "InklingSconvBackend",
     )
 
-    assert resolved_mode == CUDAGraphMode.FULL_DECODE_ONLY
+    assert resolved_mode == CUDAGraphMode.PIECEWISE

tests/v1/cudagraph/test_breakable_cudagraph.py

@@ -8,6 +8,8 @@ from __future__ import annotations
 
 import os
 import threading
+from contextlib import nullcontext
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -15,6 +17,54 @@ import torch
 os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
 
 
+def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
+    from vllm.config import CUDAGraphMode
+    from vllm.v1.worker.gpu.cudagraph_utils import (
+        BatchExecutionDescriptor,
+        CudaGraphManager,
+    )
+
+    manager = CudaGraphManager.__new__(CudaGraphManager)
+    desc = BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 8, None)
+    manager.device = torch.device("cpu")
+    manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]}
+    manager._graphs_captured = False
+    manager.use_breakable_cg = True
+
+    create_calls = []
+    forward_calls = []
+
+    def create_forward_fn(desc_arg, warmup):
+        assert desc_arg == desc
+        metadata = {"layer": object()}
+        create_calls.append((warmup, metadata))
+
+        def forward_fn(cg_mode):
+            assert metadata
+            forward_calls.append((warmup, cg_mode, metadata))
+
+        return forward_fn
+
+    with (
+        patch(
+            "vllm.v1.worker.gpu.cudagraph_utils.graph_capture",
+            return_value=nullcontext(),
+        ),
+        patch(
+            "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank",
+            return_value=False,
+        ),
+    ):
+        manager.capture(create_forward_fn)
+
+    assert [warmup for warmup, _ in create_calls] == [True, False]
+    assert [mode for _, mode, _ in forward_calls] == [
+        CUDAGraphMode.NONE,
+        CUDAGraphMode.PIECEWISE,
+    ]
+    assert create_calls[0][1] is not create_calls[1][1]
+
+
 @pytest.fixture(autouse=True)
 def _reset_breakable_tls():
     """Defensively clear thread-local capture state between tests so a

vllm/config/vllm.py

@@ -1168,6 +1168,8 @@ class VllmConfig:
                 in (
                     "DeepseekV4ForCausalLM",
                     "DeepSeekV4MTPModel",
+                    "InklingForCausalLM",
+                    "InklingForConditionalGeneration",
                     "MiniMaxM3SparseForCausalLM",
                     "MiniMaxM3SparseForConditionalGeneration",
                 )

vllm/models/inkling/nvidia/attention.py

@@ -7,6 +7,7 @@ from typing import cast
 import torch
 from torch import nn
 
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
@@ -286,6 +287,7 @@ class InklingAttention(nn.Module, AttentionLayerBase):
         output, _ = self.wo_ud(flat)
         return output
 
+    @eager_break_during_capture
     def _attention(
         self,
         q: torch.Tensor,

vllm/models/inkling/nvidia/ops/sconv.py

@@ -21,8 +21,8 @@ spec alike.
 All kernels address the cache purely by ``(slot, absolute_position)`` and
 allocate nothing inside the captured region; their grids depend only on the
 token count (``fused_sconv`` on a fixed token/channel tiling), so the same
-decode path replays correctly under a full CUDA graph without any
-data-dependent shape or branch.
+path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs
+without any data-dependent shape or branch.
 """
 
 from __future__ import annotations
@@ -167,8 +167,8 @@ def fused_sconv(
     """Single-launch insert + depthwise causal conv1d over the paged cache.
 
     Reads same-forward taps from ``x`` and pre-forward taps from the cache, so
-    it is race-free in one launch for prefill / decode / spec and supports full
-    CUDA-graph capture for decode.
+    it is race-free in one launch for prefill / decode / spec and cudagraph-safe
+    under eager / piecewise / full capture.
     """
     T = x.shape[0]
     out = torch.empty_like(x)

vllm/models/inkling/nvidia/sconv_swa_attn.py

@@ -49,9 +49,7 @@ class InklingSconvMetadata(AttentionMetadata):
 
 
 class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]):
-    _cudagraph_support: ClassVar[AttentionCGSupport] = (
-        AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
-    )
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
     def __init__(
         self,

vllm/models/inkling/nvidia/short_conv.py

@@ -17,8 +17,8 @@ Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` /
 the owner's prefix in the forward context; the absolute ``positions`` are
 threaded in from the model. The insert + conv run in a single ``fused_sconv``
 launch (same path for prefill / decode / mixed / spec). All inputs are
-fixed-address persistent buffers and the grid is fixed, so decode can replay
-under a full CUDA graph.
+fixed-address persistent buffers and the grid is fixed, so the conv replays
+correctly under eager, PIECEWISE, and FULL cudagraphs.
 """
 
 from __future__ import annotations

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -138,8 +138,6 @@ class CudaGraphManager:
         self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
         self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
 
-        self._init_candidates()
-
         # Breakable CUDA graph (PW CUDA graph without torch.compile)
         self.use_breakable_cg = (
             is_breakable_cudagraph_enabled()
@@ -147,6 +145,8 @@ class CudaGraphManager:
         )
         self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
 
+        self._init_candidates()
+
     def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:
         """Precompute actual num_active_loras -> effective captured case.
 
@@ -255,11 +255,11 @@ class CudaGraphManager:
                 # for PIECEWISE graphs there is no limit on requests when replaying
                 # i.e. no request padding is needed
                 # so we leave it as None
-                num_reqs = (
-                    min(num_tokens, self.max_num_reqs)
-                    if mixed_mode == CUDAGraphMode.FULL
-                    else None
-                )
+                num_reqs = None
+                if mixed_mode == CUDAGraphMode.FULL or (
+                    mixed_mode == CUDAGraphMode.PIECEWISE and self.use_breakable_cg
+                ):
+                    num_reqs = min(num_tokens, self.max_num_reqs)
                 desc = BatchExecutionDescriptor(
                     cg_mode=mixed_mode,
                     num_tokens=num_tokens,
@@ -301,12 +301,10 @@ class CudaGraphManager:
 
         Args:
             create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
-                returns a forward_fn. For FULL cudagraph mode, it is invoked
-                once with warmup=True for the warmup pass, and again with
-                warmup=False for the captured pass. For attention backends
-                that perform lazy metadata initialization (e.g. FlashMLA),
-                FULL cudagraph capture requires distinct metadatas for warmup
-                and capture.
+                returns a forward_fn. For FULL and breakable PIECEWISE modes,
+                it is invoked once with warmup=True and again with warmup=False
+                because attention backends may mutate or lazily initialize
+                metadata during warmup.
         """
         with graph_capture(device=self.device):
             # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
@@ -330,11 +328,17 @@ class CudaGraphManager:
                     logger.debug(
                         "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
                     )
-                    if desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                    if (
+                        desc.cg_mode == CUDAGraphMode.PIECEWISE
+                        and not self.use_breakable_cg
+                    ):
                         forward_fn(CUDAGraphMode.PIECEWISE)
                     else:
                         # Capture with fresh attention state.
                         forward_fn = create_forward_fn(desc, warmup=False)
+                        if desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                            forward_fn(CUDAGraphMode.PIECEWISE)
+                            continue
                         assert desc not in self.graphs, (
                             f"Graph already captured for {desc}"
                         )
@@ -489,7 +493,10 @@ class ModelCudaGraphManager(CudaGraphManager):
                 block_tables,
                 attn_groups,
                 kv_cache_config,
-                skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
+                skip_attn=(
+                    desc.cg_mode == CUDAGraphMode.PIECEWISE
+                    and not self.use_breakable_cg
+                ),
             )
 
             # Capture with dummy rows marked as padding.
@@ -498,7 +505,7 @@ class ModelCudaGraphManager(CudaGraphManager):
             def forward_fn(cg_mode: CUDAGraphMode) -> None:
                 batch_descriptor = None
                 if cg_mode == CUDAGraphMode.PIECEWISE:
-                    assert attn_metadata is None
+                    assert (attn_metadata is not None) == self.use_breakable_cg
                     batch_descriptor = BatchDescriptor(
                         num_tokens=num_tokens,
                         has_lora=has_lora,

vllm/v1/worker/gpu/model_states/default.py

@@ -5,6 +5,7 @@ from typing import Any
 import torch
 import torch.nn as nn
 
+from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import NewRequestData
@@ -139,7 +140,10 @@ class DefaultModelState(ModelState):
         kv_cache_config: KVCacheConfig,
         for_capture: bool = False,
     ) -> dict[str, Any]:
-        if cudagraph_mode == CUDAGraphMode.FULL:
+        if cudagraph_mode == CUDAGraphMode.FULL or (
+            cudagraph_mode == CUDAGraphMode.PIECEWISE
+            and is_breakable_cudagraph_enabled()
+        ):
             # Use padded sizes - padding is handled by model_runner.prepare_attn.
             num_reqs = input_batch.num_reqs_after_padding
             num_tokens = input_batch.num_tokens_after_padding

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

@@ -56,7 +56,10 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
                 block_tables,
                 attn_groups,
                 kv_cache_config,
-                skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
+                skip_attn=(
+                    desc.cg_mode == CUDAGraphMode.PIECEWISE
+                    and not self.use_breakable_cg
+                ),
             )
 
             return lambda cg_mode: forward_fn(