[Model] Add PW CUDA graph support for Inkling [2/N] (#48822)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Jul 16, 2026 at 10:28 UTC
251f7e478e8eb0c90a01eb8fff40056da2aa3ff7
10 files changed
+98
-32
tests/models/inkling/test_contract_validation.py
+5
-5
@@ -34,17 +34,17 @@ def test_inkling_raw_2d_audio_is_rejected_as_ambiguous():
34
parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32))
35
36
37
-def test_inkling_supports_full_decode_only_cudagraphs():
37
+def test_inkling_supports_piecewise_cudagraphs():
38
support = InklingSconvMetadataBuilder.get_cudagraph_support
39
- assert support(None, None) == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
39
+ assert support(None, None) == AttentionCGSupport.UNIFORM_BATCH
40
41
compilation_config = CompilationConfig(
42
- cudagraph_mode=CUDAGraphMode.FULL,
42
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
43
splitting_ops=[],
44
)
45
resolved_mode = compilation_config.resolve_cudagraph_mode_and_sizes(
46
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE,
46
+ AttentionCGSupport.UNIFORM_BATCH,
47
"InklingSconvBackend",
48
)
49
50
- assert resolved_mode == CUDAGraphMode.FULL_DECODE_ONLY
50
+ assert resolved_mode == CUDAGraphMode.PIECEWISE
tests/v1/cudagraph/test_breakable_cudagraph.py
+50
@@ -8,6 +8,8 @@ from __future__ import annotations
8
9
import os
10
import threading
11
+from contextlib import nullcontext
12
+from unittest.mock import patch
13
14
import pytest
15
import torch
@@ -15,6 +17,54 @@ import torch
17
os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
18
19
20
+def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
21
+ from vllm.config import CUDAGraphMode
22
+ from vllm.v1.worker.gpu.cudagraph_utils import (
23
+ BatchExecutionDescriptor,
24
+ CudaGraphManager,
25
+ )
26
+
27
+ manager = CudaGraphManager.__new__(CudaGraphManager)
28
+ desc = BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 8, None)
29
+ manager.device = torch.device("cpu")
30
+ manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]}
31
+ manager._graphs_captured = False
32
+ manager.use_breakable_cg = True
33
+
34
+ create_calls = []
35
+ forward_calls = []
36
+
37
+ def create_forward_fn(desc_arg, warmup):
38
+ assert desc_arg == desc
39
+ metadata = {"layer": object()}
40
+ create_calls.append((warmup, metadata))
41
+
42
+ def forward_fn(cg_mode):
43
+ assert metadata
44
+ forward_calls.append((warmup, cg_mode, metadata))
45
+
46
+ return forward_fn
47
+
48
+ with (
49
+ patch(
50
+ "vllm.v1.worker.gpu.cudagraph_utils.graph_capture",
51
+ return_value=nullcontext(),
52
+ ),
53
+ patch(
54
+ "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank",
55
+ return_value=False,
56
+ ),
57
+ ):
58
+ manager.capture(create_forward_fn)
59
+
60
+ assert [warmup for warmup, _ in create_calls] == [True, False]
61
+ assert [mode for _, mode, _ in forward_calls] == [
62
+ CUDAGraphMode.NONE,
63
+ CUDAGraphMode.PIECEWISE,
64
+ ]
65
+ assert create_calls[0][1] is not create_calls[1][1]
66
+
67
+
68
@pytest.fixture(autouse=True)
69
def _reset_breakable_tls():
70
"""Defensively clear thread-local capture state between tests so a
vllm/config/vllm.py
+2
@@ -1168,6 +1168,8 @@ class VllmConfig:
1168
in (
1169
"DeepseekV4ForCausalLM",
1170
"DeepSeekV4MTPModel",
1171
+ "InklingForCausalLM",
1172
+ "InklingForConditionalGeneration",
1173
"MiniMaxM3SparseForCausalLM",
1174
"MiniMaxM3SparseForConditionalGeneration",
1175
)
vllm/models/inkling/nvidia/attention.py
+2
@@ -7,6 +7,7 @@ from typing import cast
7
import torch
8
from torch import nn
9
10
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
from vllm.config import VllmConfig, get_current_vllm_config
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.forward_context import get_forward_context
@@ -286,6 +287,7 @@ class InklingAttention(nn.Module, AttentionLayerBase):
287
output, _ = self.wo_ud(flat)
288
return output
289
290
+ @eager_break_during_capture
291
def _attention(
292
self,
293
q: torch.Tensor,
vllm/models/inkling/nvidia/ops/sconv.py
+4
-4
@@ -21,8 +21,8 @@ spec alike.
21
All kernels address the cache purely by ``(slot, absolute_position)`` and
22
allocate nothing inside the captured region; their grids depend only on the
23
token count (``fused_sconv`` on a fixed token/channel tiling), so the same
24
-decode path replays correctly under a full CUDA graph without any
25
-data-dependent shape or branch.
24
+path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs
25
+without any data-dependent shape or branch.
26
"""
27
28
from __future__ import annotations
@@ -167,8 +167,8 @@ def fused_sconv(
167
"""Single-launch insert + depthwise causal conv1d over the paged cache.
168
169
Reads same-forward taps from ``x`` and pre-forward taps from the cache, so
170
- it is race-free in one launch for prefill / decode / spec and supports full
171
- CUDA-graph capture for decode.
170
+ it is race-free in one launch for prefill / decode / spec and cudagraph-safe
171
+ under eager / piecewise / full capture.
172
"""
173
T = x.shape[0]
174
out = torch.empty_like(x)
vllm/models/inkling/nvidia/sconv_swa_attn.py
+1
-3
@@ -49,9 +49,7 @@ class InklingSconvMetadata(AttentionMetadata):
49
50
51
class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]):
52
- _cudagraph_support: ClassVar[AttentionCGSupport] = (
53
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
54
- )
52
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
53
54
def __init__(
55
self,
vllm/models/inkling/nvidia/short_conv.py
+2
-2
@@ -17,8 +17,8 @@ Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` /
17
the owner's prefix in the forward context; the absolute ``positions`` are
18
threaded in from the model. The insert + conv run in a single ``fused_sconv``
19
launch (same path for prefill / decode / mixed / spec). All inputs are
20
-fixed-address persistent buffers and the grid is fixed, so decode can replay
21
-under a full CUDA graph.
20
+fixed-address persistent buffers and the grid is fixed, so the conv replays
21
+correctly under eager, PIECEWISE, and FULL cudagraphs.
22
"""
23
24
from __future__ import annotations
vllm/v1/worker/gpu/cudagraph_utils.py
+23
-16
@@ -138,8 +138,6 @@ class CudaGraphManager:
138
self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
139
self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
140
141
- self._init_candidates()
142
-
141
# Breakable CUDA graph (PW CUDA graph without torch.compile)
142
self.use_breakable_cg = (
143
is_breakable_cudagraph_enabled()
@@ -147,6 +145,8 @@ class CudaGraphManager:
145
)
146
self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
147
148
+ self._init_candidates()
149
+
150
def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:
151
"""Precompute actual num_active_loras -> effective captured case.
152
@@ -255,11 +255,11 @@ class CudaGraphManager:
255
# for PIECEWISE graphs there is no limit on requests when replaying
256
# i.e. no request padding is needed
257
# so we leave it as None
258
- num_reqs = (
259
- min(num_tokens, self.max_num_reqs)
260
- if mixed_mode == CUDAGraphMode.FULL
261
- else None
262
- )
258
+ num_reqs = None
259
+ if mixed_mode == CUDAGraphMode.FULL or (
260
+ mixed_mode == CUDAGraphMode.PIECEWISE and self.use_breakable_cg
261
+ ):
262
+ num_reqs = min(num_tokens, self.max_num_reqs)
263
desc = BatchExecutionDescriptor(
264
cg_mode=mixed_mode,
265
num_tokens=num_tokens,
@@ -301,12 +301,10 @@ class CudaGraphManager:
301
302
Args:
303
create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and
304
- returns a forward_fn. For FULL cudagraph mode, it is invoked
305
- once with warmup=True for the warmup pass, and again with
306
- warmup=False for the captured pass. For attention backends
307
- that perform lazy metadata initialization (e.g. FlashMLA),
308
- FULL cudagraph capture requires distinct metadatas for warmup
309
- and capture.
304
+ returns a forward_fn. For FULL and breakable PIECEWISE modes,
305
+ it is invoked once with warmup=True and again with warmup=False
306
+ because attention backends may mutate or lazily initialize
307
+ metadata during warmup.
308
"""
309
with graph_capture(device=self.device):
310
# Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger
@@ -330,11 +328,17 @@ class CudaGraphManager:
328
logger.debug(
329
"CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc
330
)
333
- if desc.cg_mode == CUDAGraphMode.PIECEWISE:
331
+ if (
332
+ desc.cg_mode == CUDAGraphMode.PIECEWISE
333
+ and not self.use_breakable_cg
334
+ ):
335
forward_fn(CUDAGraphMode.PIECEWISE)
336
else:
337
# Capture with fresh attention state.
338
forward_fn = create_forward_fn(desc, warmup=False)
339
+ if desc.cg_mode == CUDAGraphMode.PIECEWISE:
340
+ forward_fn(CUDAGraphMode.PIECEWISE)
341
+ continue
342
assert desc not in self.graphs, (
343
f"Graph already captured for {desc}"
344
)
@@ -489,7 +493,10 @@ class ModelCudaGraphManager(CudaGraphManager):
493
block_tables,
494
attn_groups,
495
kv_cache_config,
492
- skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
496
+ skip_attn=(
497
+ desc.cg_mode == CUDAGraphMode.PIECEWISE
498
+ and not self.use_breakable_cg
499
+ ),
500
)
501
502
# Capture with dummy rows marked as padding.
@@ -498,7 +505,7 @@ class ModelCudaGraphManager(CudaGraphManager):
505
def forward_fn(cg_mode: CUDAGraphMode) -> None:
506
batch_descriptor = None
507
if cg_mode == CUDAGraphMode.PIECEWISE:
501
- assert attn_metadata is None
508
+ assert (attn_metadata is not None) == self.use_breakable_cg
509
batch_descriptor = BatchDescriptor(
510
num_tokens=num_tokens,
511
has_lora=has_lora,
vllm/v1/worker/gpu/model_states/default.py
+5
-1
@@ -5,6 +5,7 @@ from typing import Any
5
import torch
6
import torch.nn as nn
7
8
+from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled
9
from vllm.config import VllmConfig
10
from vllm.config.compilation import CUDAGraphMode
11
from vllm.v1.core.sched.output import NewRequestData
@@ -139,7 +140,10 @@ class DefaultModelState(ModelState):
140
kv_cache_config: KVCacheConfig,
141
for_capture: bool = False,
142
) -> dict[str, Any]:
142
- if cudagraph_mode == CUDAGraphMode.FULL:
143
+ if cudagraph_mode == CUDAGraphMode.FULL or (
144
+ cudagraph_mode == CUDAGraphMode.PIECEWISE
145
+ and is_breakable_cudagraph_enabled()
146
+ ):
147
# Use padded sizes - padding is handled by model_runner.prepare_attn.
148
num_reqs = input_batch.num_reqs_after_padding
149
num_tokens = input_batch.num_tokens_after_padding
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
- skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
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(