From 251f7e478e8eb0c90a01eb8fff40056da2aa3ff7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Thu, 16 Jul 2026 10:28:56 -0700 Subject: [PATCH] [Model] Add PW CUDA graph support for Inkling [2/N] (#48822) Signed-off-by: Woosuk Kwon Co-authored-by: Bugen Zhao Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Jee Jee Li Co-authored-by: Roger Wang Co-authored-by: Yifan Qiao Co-authored-by: Claude Fable 5 Co-authored-by: OpenAI Codex --- .../inkling/test_contract_validation.py | 10 ++-- .../v1/cudagraph/test_breakable_cudagraph.py | 50 +++++++++++++++++++ vllm/config/vllm.py | 2 + vllm/models/inkling/nvidia/attention.py | 2 + vllm/models/inkling/nvidia/ops/sconv.py | 8 +-- vllm/models/inkling/nvidia/sconv_swa_attn.py | 4 +- vllm/models/inkling/nvidia/short_conv.py | 4 +- vllm/v1/worker/gpu/cudagraph_utils.py | 39 +++++++++------ vllm/v1/worker/gpu/model_states/default.py | 6 ++- .../autoregressive/cudagraph_utils.py | 5 +- 10 files changed, 98 insertions(+), 32 deletions(-) diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py index 535359996dd..6d0ae51c9d1 100644 --- a/tests/models/inkling/test_contract_validation.py +++ b/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 diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index f856d91b639..742aafd3890 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 4afaf30f375..057baed4b27 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1168,6 +1168,8 @@ class VllmConfig: in ( "DeepseekV4ForCausalLM", "DeepSeekV4MTPModel", + "InklingForCausalLM", + "InklingForConditionalGeneration", "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", ) diff --git a/vllm/models/inkling/nvidia/attention.py b/vllm/models/inkling/nvidia/attention.py index 60440834cfe..fc617b89ad6 100644 --- a/vllm/models/inkling/nvidia/attention.py +++ b/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, diff --git a/vllm/models/inkling/nvidia/ops/sconv.py b/vllm/models/inkling/nvidia/ops/sconv.py index f21ff6e04cb..7ebb2a8a0a8 100644 --- a/vllm/models/inkling/nvidia/ops/sconv.py +++ b/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) diff --git a/vllm/models/inkling/nvidia/sconv_swa_attn.py b/vllm/models/inkling/nvidia/sconv_swa_attn.py index be27c5fd138..8a2b288d5ee 100644 --- a/vllm/models/inkling/nvidia/sconv_swa_attn.py +++ b/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, diff --git a/vllm/models/inkling/nvidia/short_conv.py b/vllm/models/inkling/nvidia/short_conv.py index 69d5a263749..363423bf2aa 100644 --- a/vllm/models/inkling/nvidia/short_conv.py +++ b/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 diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9898fadc25e..27fd5547257 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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, diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 854b71b69fc..18eb40640ad 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/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 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 b30712b5a5d..19919043c83 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/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(