[Perf][DSv4.1] Run decoder replay layers in CUDA graphs and trim in PIECEWISE graphs (#59532)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/models/test_deepseek_v41_decoder_replay_layers.py

@@ -2,15 +2,26 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """DecoderReplayLayers runs the replay layers on the forward's replay batch."""
 
+from unittest.mock import MagicMock
+
 import pytest
 import torch
 
+from vllm.config import CompilationConfig, CUDAGraphMode
 from vllm.forward_context import (
+    BatchDescriptor,
     ForwardContext,
     get_forward_context,
     override_forward_context,
 )
-from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
+from vllm.models.deepseek_v41.decoder_replay_layers import (
+    DecoderReplayLayers,
+    ReplayBatch,
+)
+from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import (
+    DecoderReplayCudaGraphManager,
+)
+from vllm.v1.worker.gpu import cudagraph_utils
 
 pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
 
@@ -32,9 +43,9 @@ def _context(metadata):
 
 
 def _set_replay_batch(layers, rows, metadata):
-    layers.rows = torch.tensor(rows, device=DEVICE)
-    layers.forward_context = ForwardContext(
-        no_compile_layers={}, attn_metadata=metadata, slot_mapping={}
+    layers.replay_batch = ReplayBatch(
+        torch.tensor(rows, device=DEVICE),
+        ForwardContext(no_compile_layers={}, attn_metadata=metadata, slot_mapping={}),
     )
 
 
@@ -69,7 +80,7 @@ def test_run_gathers_states_and_realigns_shared_indexer_buffers():
         seen["is_padding"] = replay_context.is_padding
         return residual, pre_mix
 
-    layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates])
+    layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates], set())
     states = _states(NUM_TOKENS)
     hidden, residual = states[0], states[-1]
     full, sub = object(), object()
@@ -96,9 +107,54 @@ def test_no_replay_batch_runs_the_whole_batch():
         seen["attn_metadata"] = get_forward_context().attn_metadata
         return (hidden_states,)
 
-    layers = DecoderReplayLayers(WINDOW, run_layers, [])
+    layers = DecoderReplayLayers(WINDOW, run_layers, [], set())
     states = _states(NUM_TOKENS)
     full = object()
     with override_forward_context(_context(full)):
         (output,) = layers(*states)
     assert output is states[0] and seen["attn_metadata"] is full
+
+
+def _graph_context(num_tokens):
+    return ForwardContext(
+        no_compile_layers={},
+        attn_metadata={},
+        slot_mapping={},
+        cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
+        batch_descriptor=BatchDescriptor(num_tokens=num_tokens),
+    )
+
+
+def test_replay_graph_matches_eager(monkeypatch):
+    """The replay graph of the next captured size pads and runs the rows."""
+    monkeypatch.setattr(cudagraph_utils, "get_pp_group", MagicMock)
+    capture = torch.cuda.stream(torch.cuda.Stream())
+    monkeypatch.setattr(cudagraph_utils, "graph_capture", lambda device: capture)
+    compilation = CompilationConfig(decoder_replay_cudagraph_capture_sizes=[4])
+    cfg = MagicMock(compilation_config=compilation, speculative_config=None)
+    cfg.scheduler_config = MagicMock(max_num_seqs=2, max_num_batched_tokens=64)
+    cfg.cache_config.use_kda_recoverssm = False
+    cfg.model_config.hf_config = MagicMock(hc_mult=HC, hidden_size=HIDDEN)
+    cfg.model_config.dtype = torch.float32
+
+    def run_layers(hidden, positions, ids, pre, post, mix, residual):
+        return residual + hidden[:, None] + positions[:, None, None], pre + ids[:, None]
+
+    layers = DecoderReplayLayers(4, run_layers, [], set())
+    manager = DecoderReplayCudaGraphManager(cfg, DEVICE, layers)
+
+    def prepare(desc):
+        rows = torch.arange(desc.num_tokens)
+        layers.replay_batch = ReplayBatch(rows, _graph_context(desc.num_tokens))
+
+    manager.capture_replay_graphs(prepare)
+    states = [
+        torch.randn(12, *b.shape[1:], device=DEVICE).to(b.dtype)
+        for b in manager.input_buffers
+    ]
+    rows = torch.tensor([11, 2, 7], device=DEVICE)
+    desc = manager.dispatch(2, 3, None, 0)
+    with override_forward_context(_graph_context(desc.num_tokens)):
+        actual = manager.run(rows, states)
+    for a, e in zip(actual, run_layers(*(t[rows] for t in states))):
+        torch.testing.assert_close(a, e, rtol=0, atol=0)

tests/models/test_deepseek_v41_replay_batch.py

@@ -12,6 +12,7 @@ import torch
 
 from vllm.config import CUDAGraphMode
 from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 
@@ -57,7 +58,8 @@ def state(monkeypatch):
     cfg.scheduler_config.max_num_batched_tokens = 1024
     cfg.parallel_config.data_parallel_size = 1
     cfg.compilation_config.fast_moe_cold_start = False
-    layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None)
+    cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+    layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set())
     model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers)
     builds: list = []
 
@@ -127,13 +129,8 @@ def _prepare(state, batch, cg_mode):
         [],
         GROUPS,
     )
-    layers = state.decoder_replay_layers
-    if layers.rows is None:
-        assert layers.forward_context is None
-        return None, None
-    return SimpleNamespace(
-        rows=layers.rows, forward_context=layers.forward_context
-    ), state.builds[-1]
+    replay = state.decoder_replay_layers.replay_batch
+    return (None, None) if replay is None else (replay, state.builds[-1])
 
 
 def test_replay_batch_keeps_each_request_window(state):
@@ -191,9 +188,9 @@ def test_replay_batch_keeps_adaptive_verification_query_bound(state):
     assert replay is not None and build.batch.max_query_len == 200
 
 
-def test_only_eager_steps_trim(state):
-    """A CUDA graph keeps the layers on its whole batch; an eager step trims a
-    padded batch to its real rows."""
+def test_graph_steps_trim_only_piecewise_at_threshold(state):
+    """A CUDA graph keeps the layers on its whole batch, unless it is a PIECEWISE
+    one at the trim threshold or above; an eager step trims a batch to its rows."""
     batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512)
     for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
         assert _prepare(state, batch, cg_mode) == (None, None)
@@ -202,6 +199,10 @@ def test_only_eager_steps_trim(state):
     assert replay.rows.tolist() == REPLAY_ROWS
     assert build.cg_mode == CUDAGraphMode.NONE
     assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS)
+    state._trim_threshold = 256
+    replay, _ = _prepare(state, batch, CUDAGraphMode.PIECEWISE)
+    assert replay is not None and replay.rows.tolist() == REPLAY_ROWS
+    assert _prepare(state, batch, CUDAGraphMode.FULL) == (None, None)
 
 
 def test_whole_batch_forwards_get_no_replay_batch(state):
@@ -251,9 +252,25 @@ def test_dp_ranks_replay_together(dp_state):
 
 
 def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
-    dummy = _input_batch([512], [512], [False])
+    """An idle rank's dummy keeps one row, so peers do not wait on its replay."""
+    dummy = _input_batch([256, 256], [256, 256], [False, False])
     dp_state.other = (True, 129)
     replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE)
-    assert replay is not None and replay.rows.shape[0] == WINDOW
+    assert replay is not None and replay.rows.shape[0] == 1
     dp_metadata = replay.forward_context.dp_metadata
-    assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129]
+    assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129]
+
+
+def test_dp_ranks_pad_to_one_replay_graph(dp_state):
+    """Every rank pads its replay batch to the graph fitting the largest rank's."""
+    dp_state.replay_cudagraphs = MagicMock()
+    dp_state.replay_cudagraphs.dispatch.return_value = BatchExecutionDescriptor(
+        CUDAGraphMode.PIECEWISE, 512, None
+    )
+    dp_state.other = (True, 300)
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
+    replay, _ = _prepare(dp_state, batch, CUDAGraphMode.NONE)
+    assert dp_state.replay_cudagraphs.dispatch.call_args.args[:2] == (3, 300)
+    assert replay.forward_context.batch_descriptor.num_tokens == 512
+    dp_metadata = replay.forward_context.dp_metadata
+    assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [512, 512]

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -259,7 +259,9 @@ def _make_capture_runner(captured: bool) -> GPUModelRunner:
     """Minimal V2 runner for capture_model: fakes everything except the
     cudagraph_manager's needs_capture decision."""
     runner = GPUModelRunner.__new__(GPUModelRunner)
-    runner.model_state = SimpleNamespace(supports_mm_inputs=False)
+    runner.model_state = SimpleNamespace(
+        supports_mm_inputs=False, capture_inner_cudagraphs=lambda *args: None
+    )
     runner.cudagraph_manager = SimpleNamespace(
         needs_capture=lambda: captured,
         capture=lambda *args, **kwargs: None,

vllm/config/compilation.py

@@ -570,6 +570,24 @@ class CompilationConfig:
     If we limit the video count per prompt to `0`, it will also be set to `0`
     (i.e., fall back to image-only mode)."""
 
+    cudagraph_decoder_replay: bool = True
+    """Run the decoder replay layers of YOCO models (e.g. DeepSeek-V4.1) in CUDA
+    graphs of their own, and trim them in PIECEWISE graph steps. Requires
+    breakable PIECEWISE graphs; off under LoRA, prompt embeddings and non-first
+    PP ranks."""
+
+    decoder_replay_cudagraph_capture_sizes: list[int] = field(default_factory=list)
+    """Decoder replay CUDA graph sizes for YOCO models (e.g. DeepSeek-V4.1), at
+    most max_num_batched_tokens. If empty: multiples of sliding_window up to
+    min(max_cudagraph_capture_size, max_num_seqs * sliding_window), coarser past
+    16 * sliding_window. Larger replay batches run eagerly."""
+
+    decoder_replay_trim_threshold: int = Field(default=768, ge=0)
+    """For YOCO models (e.g. DeepSeek-V4.1), PIECEWISE graphs with at least this
+    many padded tokens trim the decoder replay batch; must exceed the replay
+    window (128), below which nothing trims. Independent of the replay graph
+    capture sizes; eager steps still trim."""
+
     # Inductor capture
     compile_sizes: list[int | str] | None = None
     """Sizes to compile for inductor. In addition
@@ -1033,6 +1051,11 @@ class CompilationConfig:
                 f"Invalid backend for piecewise compilation: {self.backend}"
             )
 
+        if any(s <= 0 for s in self.decoder_replay_cudagraph_capture_sizes):
+            raise ValueError(
+                "All decoder_replay_cudagraph_capture_sizes must be positive"
+            )
+
         # Validate encoder CUDA graph configuration
         if (
             self.cudagraph_mm_encoder

vllm/models/deepseek_v41/decoder_replay_layers.py

@@ -7,25 +7,43 @@ in eager prefill steps they run on each request's last ``window`` rows only.
 ``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention
 metadata and a forward context of its own, like a microbatch;
 ``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers
-under that context and scatters the outputs back to batch rows. Steps that run
-in a CUDA graph keep the layers on the whole batch, inside the graph.
+under that context and scatters the outputs back to batch rows. FULL graphs and
+small PIECEWISE ones keep the layers on the whole batch; larger PIECEWISE graphs
+break out to the replay batch, which may run in a graph of its own
+(decoder_replay_cudagraph.py).
 """
 
 from collections.abc import Callable
+from dataclasses import dataclass
 
 import torch
 
-from vllm.forward_context import ForwardContext, override_forward_context
+from vllm.compilation.breakable_cudagraph import eager_break_during_capture
+from vllm.forward_context import (
+    ForwardContext,
+    in_piecewise_cudagraph,
+    override_forward_context,
+)
 from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput
 
 
+@dataclass
+class ReplayBatch:
+    """One forward's replay batch; ``run_graph`` is set when it runs in a graph."""
+
+    rows: torch.Tensor
+    forward_context: ForwardContext
+    run_graph: Callable[..., tuple[torch.Tensor, ...]] | None = None
+
+
 class DecoderReplayLayers:
     """Runs the replay layers on the step's replay batch.
 
     ``run_layers`` takes a batch's layer inputs and returns its per-row
     outputs. ``row_buffers`` hold per-row results the source layer's indexer
     publishes for the layers after it; they are compacted to the replay rows
-    in place.
+    in place. ``metadata_prefixes`` are the attention metadata keys the layers
+    read, so the replay batch builds only those.
     """
 
     def __init__(
@@ -33,32 +51,57 @@ class DecoderReplayLayers:
         window: int,
         run_layers: Callable[..., tuple[torch.Tensor, ...]],
         row_buffers: list[torch.Tensor],
+        metadata_prefixes: set[str],
     ) -> None:
         self.window = window
         self.run_layers = run_layers
         self.row_buffers = row_buffers
-        # The replay batch, set by the model state every step: its rows of the
-        # batch and its forward context. None runs the layers on the batch.
-        self.rows: torch.Tensor | None = None
-        self.forward_context: ForwardContext | None = None
+        self.metadata_prefixes = metadata_prefixes
+        # Set by the model state every step; None runs the layers on the batch.
+        self.replay_batch: ReplayBatch | None = None
+        self._break_outputs: list[torch.Tensor] | None = None
 
     def __call__(
         self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
     ) -> tuple[torch.Tensor, ...]:
-        rows = self.rows
-        if rows is None:
+        if self.replay_batch is not None and in_piecewise_cudagraph():
+            return self._run_in_graph_break(hidden_states, *states)
+        return self._run(hidden_states, *states)
+
+    @eager_break_during_capture
+    def _run_in_graph_break(
+        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
+    ) -> tuple[torch.Tensor, ...]:
+        """A graph break: the outputs land in fixed buffers the next segment reads."""
+        outputs = self._run(hidden_states, *states)
+        if self._break_outputs is None:
+            # Sized by the first call: the largest model graph is captured first.
+            self._break_outputs = [torch.empty_like(out) for out in outputs]
+        return tuple(
+            buf[: out.shape[0]].copy_(out)
+            for buf, out in zip(self._break_outputs, outputs)
+        )
+
+    def _run(
+        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
+    ) -> tuple[torch.Tensor, ...]:
+        replay_batch = self.replay_batch
+        if replay_batch is None:
             return self.run_layers(hidden_states, *states)
-        # A trimming step holds a prefill longer than the window, more tokens
-        # than any step whose MoE leaves its finalize to the next layer.
+        # Replay batch steps hold more than `window` tokens: no deferred MoE finalize.
         assert isinstance(hidden_states, torch.Tensor)
+        rows = replay_batch.rows
         num_rows = rows.shape[0]
         for buf in self.row_buffers:
             buf[:num_rows].copy_(buf.index_select(0, rows))
-        with override_forward_context(self.forward_context):
-            row_outputs = self.run_layers(
-                hidden_states.index_select(0, rows),
-                *(None if t is None else t.index_select(0, rows) for t in states),
-            )
+        inputs = (hidden_states, *states)
+        with override_forward_context(replay_batch.forward_context):
+            if replay_batch.run_graph is not None:
+                row_outputs = replay_batch.run_graph(rows, inputs)
+            else:
+                row_outputs = self.run_layers(
+                    *(None if t is None else t.index_select(0, rows) for t in inputs)
+                )
         # The trimmed rows' outputs stay zero; nothing reads them.
         num_tokens = hidden_states.shape[0]
         return tuple(

vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUDA graphs of the decoder replay layers on a step's replay batch."""
+
+from collections.abc import Callable
+from copy import copy
+from dataclasses import replace
+
+import torch
+
+from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.config.kernel import MEGA_MOE_BACKENDS
+from vllm.forward_context import get_forward_context, override_forward_context
+from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    BatchExecutionDescriptor,
+    CudaGraphManager,
+)
+
+
+class DecoderReplayCudaGraphManager(CudaGraphManager):
+    """Breakable PIECEWISE graphs of the replay layers keyed by the replay batch's
+    padded size. They read the layer inputs from static buffers."""
+
+    def __init__(self, vllm_config: VllmConfig, device, layers: DecoderReplayLayers):
+        compilation = copy(vllm_config.compilation_config)
+        scheduler = vllm_config.scheduler_config
+        sizes = compilation.decoder_replay_cudagraph_capture_sizes
+        if not sizes:
+            bound = min(
+                compilation.max_cudagraph_capture_size,
+                scheduler.max_num_seqs * layers.window,
+            )
+            # Replay batches hold at least `window` rows: step by window, or by 1/8
+            # of the power-of-two interval once wider (256 in [2048, 4096)).
+            sizes, s = [bound], layers.window
+            while s < bound:
+                sizes.append(s)
+                s += max(layers.window, 1 << max(s.bit_length() - 4, 0))
+        sizes = sorted(set(sizes))
+        if max(sizes) > scheduler.max_num_batched_tokens:
+            raise ValueError(
+                "decoder_replay_cudagraph_capture_sizes must not exceed "
+                "max_num_batched_tokens"
+            )
+        compilation.cudagraph_capture_sizes = sizes
+        compilation.max_cudagraph_capture_size = max(sizes)
+        replay_config = copy(vllm_config)
+        replay_config.compilation_config = compilation
+        super().__init__(replay_config, device, CUDAGraphMode.PIECEWISE, 1)
+        self.layers = layers
+        self.breakable_cg_runner: BreakableCUDAGraphWrapper = BreakableCUDAGraphWrapper(
+            layers.run_layers, replay_config
+        )
+        config, act = vllm_config.model_config.hf_config, vllm_config.model_config.dtype
+        assert isinstance(act, torch.dtype)
+        hc, h, f32 = config.hc_mult, config.hidden_size, torch.float32
+        mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS
+        ids = torch.int64 if mega_moe else torch.int32
+        # hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual
+        shapes = [(h,), (), (), (hc,), (hc, 1), (hc, hc), (hc, h)]
+        dtypes = [act, torch.int64, ids, f32, f32, f32, act]
+        n = compilation.max_cudagraph_capture_size
+        self.input_buffers = [
+            torch.zeros(n, *s, dtype=d, device=device) for s, d in zip(shapes, dtypes)
+        ]
+
+    def capture_replay_graphs(
+        self, prepare: Callable[[BatchExecutionDescriptor], None]
+    ) -> None:
+        """Capture each size on the replay batch ``prepare`` sets on the layers."""
+
+        def create_forward_fn(desc: BatchExecutionDescriptor, warmup: bool):
+            prepare(desc)
+            assert self.layers.replay_batch is not None
+            context = self.layers.replay_batch.forward_context
+            inputs = [buf[: desc.num_tokens] for buf in self.input_buffers]
+
+            def forward_fn(cg_mode: CUDAGraphMode) -> None:
+                with override_forward_context(
+                    replace(context, cudagraph_runtime_mode=cg_mode)
+                ):
+                    self.breakable_cg_runner(*inputs)
+
+            return forward_fn
+
+        self.capture(create_forward_fn, "Capturing decoder replay CUDA graphs")
+        self.layers.replay_batch = None
+
+    def run(self, rows: torch.Tensor, inputs: tuple) -> tuple[torch.Tensor, ...]:
+        """Replay the forward context's graph on ``rows`` of ``inputs``."""
+        num_rows = rows.shape[0]
+        for buf, t in zip(self.input_buffers, inputs):
+            torch.index_select(t, 0, rows, out=buf[:num_rows])
+        batch_descriptor = get_forward_context().batch_descriptor
+        assert batch_descriptor is not None
+        num_tokens = batch_descriptor.num_tokens
+        outputs = self.breakable_cg_runner(
+            *(buf[:num_tokens] for buf in self.input_buffers)
+        )
+        return tuple(out[:num_rows] for out in outputs)

vllm/models/deepseek_v41/nvidia/model.py

@@ -766,6 +766,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             and self.layers[cut].attn.swa_cache_layer.bounded_replay
         ):
             self.decoder_replay_start = cut + 1
+            # The attention metadata keys the replay layers read.
+            metadata_prefixes = set()
+            for layer in islice(self.layers, cut + 1, self.end_layer):
+                attn = typing.cast(DeepseekV4DecoderLayer, layer).attn
+                metadata_prefixes.add(attn.swa_cache_layer.prefix)
+                if attn.compressed_cache_prefix is not None:
+                    metadata_prefixes.add(attn.compressed_cache_prefix)
+                if attn.indexer is not None:
+                    metadata_prefixes.add(attn.indexer.k_cache.prefix)
             self.decoder_replay_layers = DecoderReplayLayers(
                 config.sliding_window,
                 self._run_replay_layers,
@@ -774,6 +783,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                     for buf in (self.topk_indices_buffer, self.candidate_block_buffer)
                     if buf is not None
                 ],
+                metadata_prefixes,
             )
             logger.info_once(
                 "Decoder SWA bounded replay: in eager prefill steps, layers "

vllm/models/deepseek_v41/nvidia/model_state.py

@@ -8,10 +8,17 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 
+from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled
 from vllm.config import CUDAGraphMode, VllmConfig
-from vllm.distributed.parallel_state import get_dp_group
-from vllm.forward_context import DPMetadata, create_forward_context
-from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
+from vllm.distributed.parallel_state import get_dp_group, get_pp_group
+from vllm.forward_context import BatchDescriptor, DPMetadata, create_forward_context
+from vllm.models.deepseek_v41.decoder_replay_layers import (
+    DecoderReplayLayers,
+    ReplayBatch,
+)
+from vllm.models.deepseek_v41.nvidia.decoder_replay_cudagraph import (
+    DecoderReplayCudaGraphManager,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -19,8 +26,10 @@ from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.dp_utils import should_skip_dp_coordination
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
-from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
@@ -230,6 +239,31 @@ class DeepseekV41ModelState(DefaultModelState):
                 self.max_num_reqs, dtype=torch.int32, device=device
             )
             self._replay_attn_groups: list[list[AttentionGroup]] | None = None
+        layers = self.decoder_replay_layers
+        compilation_config = vllm_config.compilation_config
+        self.replay_cudagraphs: DecoderReplayCudaGraphManager | None = None
+        # PIECEWISE graphs of this many tokens or more break out to the replay batch.
+        self._trim_threshold: int | None = None
+        if (
+            layers is not None
+            and compilation_config.cudagraph_decoder_replay
+            and compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
+            and is_breakable_cudagraph_enabled()
+            and get_pp_group().is_first_rank
+            # The graphs' fixed inputs hold no LoRA mapping and no None input ids.
+            and vllm_config.lora_config is None
+            and not vllm_config.model_config.enable_prompt_embeds
+        ):
+            threshold = compilation_config.decoder_replay_trim_threshold
+            if threshold <= layers.window:
+                raise ValueError(
+                    "decoder_replay_trim_threshold must exceed the replay window "
+                    f"of {layers.window} tokens: smaller graphs have nothing to trim"
+                )
+            self.replay_cudagraphs = DecoderReplayCudaGraphManager(
+                vllm_config, device, layers
+            )
+            self._trim_threshold = threshold
 
     def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
         super().add_request(req_index, new_req_data)
@@ -409,15 +443,21 @@ class DeepseekV41ModelState(DefaultModelState):
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
         replay_start: torch.Tensor,
+        capture_desc: BatchExecutionDescriptor | None = None,
     ) -> None:
-        """Set the replay layers' batch for this forward, or none when nothing
-        trims. Only eager steps trim: a CUDA graph keeps the layers on the whole
-        batch, and ranks share the graph mode. Under data parallelism every rank
-        replays if any does."""
+        """Set the replay layers' batch for this forward, or none. Eager steps trim;
+        captures and PIECEWISE graphs at the trim threshold or above break out to a
+        replay batch, trimmed or not; FULL and smaller PIECEWISE graphs keep the
+        whole batch. Under DP every rank replays if any does."""
         layers = self.decoder_replay_layers
         assert layers is not None
-        layers.rows = layers.forward_context = None
-        if cudagraph_mode != CUDAGraphMode.NONE:
+        layers.replay_batch = None
+        force_replay_batch = capture_desc is not None or (
+            cudagraph_mode == CUDAGraphMode.PIECEWISE
+            and self._trim_threshold is not None
+            and input_batch.num_tokens_after_padding >= self._trim_threshold
+        )
+        if cudagraph_mode != CUDAGraphMode.NONE and not force_replay_batch:
             return
 
         num_reqs = input_batch.num_reqs
@@ -427,20 +467,47 @@ class DeepseekV41ModelState(DefaultModelState):
             query_lens,
             np.minimum(query_lens, layers.window),
         )
+        # A batch of context-free non-prefill rows is a dummy (an idle DP rank's):
+        # nothing reads it and its idx_mapping hits stale slots, so keep one row.
+        seq_lens = input_batch.seq_lens_cpu_upper_bound[:num_reqs].numpy()
+        is_dummy = not input_batch.is_prefilling_np[:num_reqs].any() and bool(
+            (seq_lens == query_lens).all()
+        )
+        if is_dummy and capture_desc is None:
+            kept_lens = np.zeros_like(kept_lens)
+            kept_lens[0] = 1
         # Dummy batches (captures, an idle DP rank's) are not prefills: they
         # keep their rows unless a rank trims.
-        trims = bool(
+        needs_replay_batch = force_replay_batch or bool(
             (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any()
         )
         num_tokens = int(kept_lens.sum())
-        dp_metadata = None
+        dp_tokens = None
         if self.vllm_config.parallel_config.data_parallel_size > 1:
-            trims, dp_metadata = self._agree_across_dp(trims, num_tokens)
-        if not trims:
+            needs_replay_batch, dp_tokens = self._agree_across_dp(
+                needs_replay_batch, num_tokens
+            )
+        if not needs_replay_batch:
             return
 
+        cudagraphs, run_graph = self.replay_cudagraphs, None
+        num_padded, cg_mode = num_tokens, CUDAGraphMode.NONE
+        if cudagraphs is not None:
+            # Every rank runs the graph that fits the largest rank's batch.
+            dp_max_tokens = num_tokens if dp_tokens is None else int(dp_tokens.max())
+            desc = capture_desc or cudagraphs.dispatch(num_reqs, dp_max_tokens, None, 0)
+            if desc.cg_mode == CUDAGraphMode.PIECEWISE:
+                num_padded, cg_mode = desc.num_tokens, desc.cg_mode
+                run_graph = cudagraphs.run
+                if dp_tokens is not None:
+                    dp_tokens.fill_(num_padded)
+        dp_metadata = None
+        if dp_tokens is not None:
+            dp_metadata = DPMetadata.make(
+                self.vllm_config.parallel_config, num_padded, dp_tokens
+            )
         kept_batch, kept_slot_mappings = self._kept_input_batch(
-            input_batch, slot_mappings, replay_start, kept_lens
+            input_batch, slot_mappings, replay_start, kept_lens, num_padded
         )
         attn_metadata = super().prepare_attn(
             kept_batch,
@@ -453,22 +520,55 @@ class DeepseekV41ModelState(DefaultModelState):
                 self._kept_kv_start[:num_reqs]
             ),
         )
-        layers.rows = self._kept_rows[:num_tokens]
-        layers.forward_context = create_forward_context(
-            attn_metadata,
-            self.vllm_config,
-            dp_metadata=dp_metadata,
-            slot_mapping=build_slot_mappings_by_layer(
-                kept_slot_mappings, kv_cache_config
+        layers.replay_batch = ReplayBatch(
+            self._kept_rows[:num_tokens],
+            create_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                dp_metadata=dp_metadata,
+                cudagraph_runtime_mode=cg_mode,
+                batch_descriptor=BatchDescriptor(num_tokens=num_padded),
+                slot_mapping=build_slot_mappings_by_layer(
+                    kept_slot_mappings, kv_cache_config
+                ),
             ),
+            run_graph,
         )
 
+    def capture_inner_cudagraphs(
+        self,
+        input_buffers: InputBuffers,
+        block_tables: BlockTables,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        def prepare(desc: BatchExecutionDescriptor) -> None:
+            # Capture reuses the runtime path on a dummy batch, so the graph reads
+            # the buffers a real step fills.
+            num_reqs = min(desc.num_tokens, self.max_num_reqs)
+            self._prepare_replay_batch(
+                InputBatch.make_dummy(num_reqs, desc.num_tokens, input_buffers),
+                CUDAGraphMode.NONE,
+                block_tables.get_dummy_block_tables(num_reqs),
+                block_tables.get_dummy_slot_mappings(desc.num_tokens),
+                attn_groups,
+                kv_cache_config,
+                self._replay_start[:num_reqs].zero_(),
+                capture_desc=desc,
+            )
+
+        if self.replay_cudagraphs is not None:
+            # profile_only captures use pre-KV-cache groups: drop the cache.
+            self._replay_attn_groups = None
+            self.replay_cudagraphs.capture_replay_graphs(prepare)
+
     def _kept_input_batch(
         self,
         input_batch: InputBatch,
         slot_mappings: torch.Tensor,
         replay_start: torch.Tensor,
         kept_lens: np.ndarray,
+        num_padded: int,
     ) -> tuple[InputBatch, torch.Tensor]:
         """Build the sub-``InputBatch`` of each request's last ``kept_lens``
         rows, and its slot mappings.
@@ -486,7 +586,8 @@ class DeepseekV41ModelState(DefaultModelState):
             input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np
         )
         assert self._kept_slot_mappings is not None
-        kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens]
+        kept_slot_mappings = self._kept_slot_mappings[:, :num_padded]
+        kept_slot_mappings[:, num_tokens:].fill_(PAD_SLOT_ID)
         _gather_replay_batch_kernel[(num_reqs,)](
             input_batch.query_start_loc,
             dropped_before,
@@ -511,42 +612,45 @@ class DeepseekV41ModelState(DefaultModelState):
         kept_batch = replace(
             input_batch,
             num_tokens=num_tokens,
-            num_tokens_after_padding=num_tokens,
+            num_tokens_after_padding=num_padded,
             query_start_loc=self._kept_query_start_loc[: num_reqs + 1],
             query_start_loc_np=kept_query_start_loc_np,
             max_query_len=max_query_len,
-            positions=self._kept_positions[:num_tokens],
+            positions=self._kept_positions[:num_padded],
             fast_prefill=None,
         )
         return kept_batch, kept_slot_mappings
 
     def _agree_across_dp(
-        self, trims: bool, num_tokens: int
-    ) -> tuple[bool, DPMetadata | None]:
-        """Whether any rank trims and, then, the replay layers' DP metadata from
-        every rank's replay token count."""
+        self, needs_replay_batch: bool, num_tokens: int
+    ) -> tuple[bool, torch.Tensor | None]:
+        """Whether any rank needs a replay batch and, if so, each rank's token count."""
         parallel_config = self.vllm_config.parallel_config
         agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32)
         agreed[:, parallel_config.data_parallel_rank] = torch.tensor(
-            [trims, num_tokens]
+            [needs_replay_batch, num_tokens]
         )
         if not should_skip_dp_coordination():
             dist.all_reduce(agreed, group=get_dp_group().cpu_group)
         if not agreed[0].any():
             return False, None
-        return True, DPMetadata.make(parallel_config, num_tokens, agreed[1])
+        return True, agreed[1]
 
     def _replay_groups(
         self, attn_groups: list[list[AttentionGroup]]
     ) -> list[list[AttentionGroup]]:
         """The attention groups with metadata builders of their own, like each
         microbatch's: a builder keeps the metadata it built, and the runner's
-        hold the batch's."""
+        hold the batch's. Only the groups the replay layers read are built."""
         if self._replay_attn_groups is None:
+            assert self.decoder_replay_layers is not None
+            prefixes = self.decoder_replay_layers.metadata_prefixes
             self._replay_attn_groups = []
             for groups in attn_groups:
                 replay_groups = []
                 for group in groups:
+                    if prefixes.isdisjoint(group.layer_names):
+                        continue
                     replay_group = replace(group, metadata_builders=[])
                     replay_group.create_metadata_builders(
                         self.vllm_config,

vllm/v1/worker/gpu/model_runner.py

@@ -1077,6 +1077,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                             self.lora_config, self
                         ),
                     )
+                    self.model_state.capture_inner_cudagraphs(
+                        input_buffers,
+                        self.block_tables,
+                        self.attn_groups,
+                        self.kv_cache_config,
+                    )
                     if self.speculator is not None:
                         with use_workspace_lane(self._draft_workspace_lane):
                             self.speculator.capture()

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

@@ -17,7 +17,8 @@ from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
-from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
 from vllm.v1.worker.gpu.states import RequestState
@@ -126,6 +127,16 @@ class ModelState(ABC):
     def apply_staged_writes(self) -> None:
         return None
 
+    def capture_inner_cudagraphs(
+        self,
+        input_buffers: InputBuffers,
+        block_tables: BlockTables,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+    ) -> None:
+        """Capture the CUDA graphs this state runs inside the model's forward."""
+        return None
+
     def get_additional_cg_support(self) -> tuple[AttentionCGSupport, str | None]:
         """Cudagraph support of attention groups this ModelState builds outside
         ``init_attn_backend`` (e.g. encoder-only layers).