diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py index 4af797a1e6..9ba2df5c8f 100644 --- a/tests/models/test_deepseek_v41_decoder_replay_layers.py +++ b/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) diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index f6efe5c516..810a1dbd9a 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/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] diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 7586245a35..0caa132d7e 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/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, diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index eefcf1f49e..afe576daa9 100644 --- a/vllm/config/compilation.py +++ b/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 diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index 869029aa74..d2d7ffd153 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/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( diff --git a/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py b/vllm/models/deepseek_v41/nvidia/decoder_replay_cudagraph.py new file mode 100644 index 0000000000..871f535584 --- /dev/null +++ b/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) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index bd91e8f12b..9c3bfa7a92 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/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 " diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 30b287db22..c09e4582ea 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/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, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b42289ef11..e42917918a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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() diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index ee3e519a6e..f03caf4dba 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/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).