From 243217751901d626722609619ac71544cc846ff5 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Thu, 10 Sep 2026 10:31:05 +0000 Subject: [PATCH 1/4] [Model] Decoder-side SWA bounded replay for DeepSeek-V4.1 Layers past the last KV-source layer (21-39) own only their sliding-window KV, and decode reads just its trailing window. In prefill they now run on each request's last 128 tokens (its replay window): after layer 20 the batch is compacted to the replay windows, the mHC states are gathered, and the replay layers run under metadata built for them by private copies of the runner's builders. A trimmed request's replay_start rises to its replay window's start, so the window is floored there like a replayed prefix-cache hit. Results are scattered back so sampling and the drafter see full-batch rows. The top-k indices and candidate blocks that layer 20's indexer publishes for the layers after it are keyed by batch row, so they are realigned to the replay window as well. CUDA graphs: under a piecewise (breakable) capture the replay becomes an eager break of the model graph, so each replay re-plans from its own batch, and the replay layers replay their own breakable graphs keyed by the padded replay batch size, captured alongside the model graphs. Decode batches and full graphs never trim. The breakable capture therefore allows a nested capture while the outer one is paused in an eager break. The compacted slot mappings live in fixed buffers: the window KV insert is recorded in the graph and reads them by address. Part of --swa-bounded-replay (on by default): replay-layer window KV is unwritten outside each request's replay window, so a prefix-cache hit must replay the window, which is what the encoder-side replay does. Where the replay-layer batch cannot shrink per rank (sequence, data or prefill-context parallelism), or a drafter or Engram layer would read rows the replay layers skip, the decoder side stays off with a warning. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Yifan Qiao --- .../test_deepseek_v4_decoder_replay_layers.py | 373 ++++++++++++++++ .../v1/cudagraph/test_breakable_cudagraph.py | 38 ++ vllm/compilation/breakable_cudagraph.py | 19 +- vllm/config/cache.py | 3 +- .../deepseek_v41/decoder_replay_layers.py | 408 ++++++++++++++++++ vllm/models/deepseek_v41/nvidia/model.py | 222 ++++++++-- vllm/models/deepseek_v41/sparse_mla.py | 4 + vllm/v1/attention/backend.py | 27 ++ vllm/v1/attention/backends/mla/indexer.py | 7 + vllm/v1/attention/backends/mla/sparse_swa.py | 6 + 10 files changed, 1078 insertions(+), 29 deletions(-) create mode 100644 tests/models/test_deepseek_v4_decoder_replay_layers.py create mode 100644 vllm/models/deepseek_v41/decoder_replay_layers.py diff --git a/tests/models/test_deepseek_v4_decoder_replay_layers.py b/tests/models/test_deepseek_v4_decoder_replay_layers.py new file mode 100644 index 000000000000..117d320e94ae --- /dev/null +++ b/tests/models/test_deepseek_v4_decoder_replay_layers.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decoder-side SWA bounded replay: the replay layers see exactly each request's +replay window, with metadata that matches what the full batch would have given +those rows, except that a trimmed request's window stops at the replay window's +start.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + override_forward_context, +) +from vllm.models.deepseek_v41.decoder_replay_layers import ( + DecoderReplayLayers, + ReplayBatchBuilder, + ReplayMetadataBuilder, +) +from vllm.models.deepseek_v41.sparse_mla import DeepseekV4SparseMLAMetadataBuilder +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder +from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="metadata builders need CUDA" +) + +WINDOW = 128 +DEVICE = torch.device("cuda") + + +def _vllm_config(): + cfg = MagicMock() + cfg.model_config.max_model_len = 4096 + cfg.model_config.hf_config = SimpleNamespace(sliding_window=WINDOW, index_topk=512) + cfg.scheduler_config.max_num_batched_tokens = 2048 + cfg.scheduler_config.max_num_seqs = 16 + cfg.speculative_config = None + cfg.parallel_config.decode_context_parallel_size = 1 + cfg.parallel_config.prefill_context_parallel_size = 1 + cfg.parallel_config.cp_kv_cache_interleave_size = 1 + cfg.attention_config.resolve_indexer_kv_dtype.return_value = "fp8" + return cfg + + +def _common( + query_lens: list[int], + seq_lens: list[int], + block_size: int, + device_query_lens: list[int] | None = None, +): + num_reqs = len(query_lens) + qsl_cpu = torch.tensor([0, *torch.tensor(query_lens).cumsum(0).tolist()]).int() + qsl = torch.tensor( + [0, *torch.tensor(device_query_lens or query_lens).cumsum(0).tolist()] + ).int() + num_tokens = int(qsl_cpu[-1]) + positions = torch.cat( + [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)] + ).to(DEVICE) + block_table = ( + torch.arange(num_reqs, device=DEVICE)[:, None] * 64 + + torch.arange(64, device=DEVICE)[None, :] + ).int() + req = torch.repeat_interleave( + torch.arange(num_reqs, device=DEVICE), torch.tensor(query_lens, device=DEVICE) + ) + slot_mapping = ( + block_table[req, positions // block_size].long() * block_size + + positions % block_size + ) + return CommonAttentionMetadata( + query_start_loc=qsl.to(DEVICE), + query_start_loc_cpu=qsl_cpu, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE), + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=max(query_lens), + max_seq_len=max(seq_lens), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + positions=positions, + ) + + +def _builders(cfg): + swa_spec = SlidingWindowMLASpec( + block_size=32, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + sliding_window=WINDOW, + alignment=576, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + mla_spec = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=584, + dtype=torch.uint8, + tokens_per_state=1, + alignment=576, + ) + idx_spec = MLAAttentionSpec( + block_size=128, + num_kv_heads=1, + head_size=132, + dtype=torch.uint8, + tokens_per_state=1, + alignment=128, + ) + return ( + DeepseekSparseSWAMetadataBuilder(swa_spec, ["swa"], cfg, DEVICE), + DeepseekV4SparseMLAMetadataBuilder(mla_spec, ["mla"], cfg, DEVICE), + DeepseekV32IndexerMetadataBuilder( + idx_spec, ["idx"], cfg, DEVICE, block_table_width=64 + ), + ) + + +def _attn(topk_buffer=None, candidate_buffer=None): + source_attn = SimpleNamespace( + prefix="mla", + indexer=SimpleNamespace(k_cache=SimpleNamespace(prefix="idx")), + topk_indices_buffer=topk_buffer, + candidate_block_buffer=candidate_buffer, + ) + replay_attn = [SimpleNamespace(swa_cache_layer=SimpleNamespace(prefix="swa"))] + return source_attn, replay_attn + + +def _make_batch(common, num_batch_tokens): + builder = ReplayBatchBuilder(WINDOW, max_num_tokens=2048, max_num_reqs=16) + return builder.build(common, num_batch_tokens) + + +def _replay_layers(run_layers, topk_buffer, candidate_buffer, graph_sizes=()): + source_attn, replay_attn = _attn(topk_buffer, candidate_buffer) + return DecoderReplayLayers( + _vllm_config(), + WINDOW, + source_attn, + replay_attn, + run_layers, + list(graph_sizes), + ) + + +# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300) +QUERY_LENS = [1, 300, 100] +SEQ_LENS = [500, 300, 300] +REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)] + + +def _full_metadata(cfg): + swa_b, mla_b, idx_b = _builders(cfg) + common_swa = _common(QUERY_LENS, SEQ_LENS, block_size=32) + common_src = _common(QUERY_LENS, SEQ_LENS, block_size=128) + return { + "swa": swa_b.build(0, common_swa), + "mla": mla_b.build(0, common_src), + "idx": idx_b.build(0, common_src), + } + + +def test_batch_keeps_each_request_replay_window(): + batch = _make_batch(_common(QUERY_LENS, SEQ_LENS, block_size=32), 401) + assert batch.trims + assert batch.rows.tolist() == REPLAY_ROWS + assert batch.query_start_loc_cpu.tolist() == [0, 1, 129, 229] + assert batch.query_start_loc.tolist() == [0, 1, 129, 229] + assert batch.num_tokens == 229 + assert batch.max_query_len == WINDOW + # Only the trimmed request's window is floored, at its replay window start. + assert batch.replay_start.tolist() == [0, 300 - WINDOW, 0] + + short = _make_batch(_common([1, 100], [500, 300], block_size=32), 101) + assert not short.trims + assert short.rows.tolist() == list(range(101)) + + # Adaptive verification splits the leading verification requests on the GPU + # (CPU [2, 2] vs device [1, 3]): never trimmed, their rows stay put, and the + # device boundaries only move by the rows trimmed off the prefill. + adaptive = _make_batch( + _common([2, 2, 300], [500, 500, 300], 32, device_query_lens=[1, 3, 300]), 304 + ) + assert adaptive.rows.tolist() == [*range(4), *range(304 - WINDOW, 304)] + assert adaptive.query_start_loc.tolist() == [0, 1, 4, 4 + WINDOW] + assert adaptive.query_start_loc_cpu.tolist() == [0, 2, 4, 4 + WINDOW] + + +def test_replay_metadata_matches_full_batch_rows(): + cfg = _vllm_config() + full = _full_metadata(cfg) + batch = _make_batch(full["swa"].common, 401) + replay = ReplayMetadataBuilder(*_attn(), max_num_tokens=2048).build( + full, batch, batch.num_tokens + ) + rows = batch.rows + + swa, swa_full = replay["swa"], full["swa"] + assert swa.num_decode_tokens == 1 and swa.num_prefill_tokens == 228 + assert torch.equal(swa.slot_mapping, swa_full.slot_mapping[rows]) + assert torch.equal(swa.token_to_req_indices, swa_full.token_to_req_indices[rows]) + assert torch.equal(swa.decode_swa_indices, swa_full.decode_swa_indices) + assert torch.equal(swa.decode_swa_lens, swa_full.decode_swa_lens) + # Prefill rows are indexed past the decode token in both layouts. + prefill_rows = rows[1:] - 1 + lens, lens_full = swa.prefill_swa_lens, swa_full.prefill_swa_lens[prefill_rows] + idx = swa.prefill_swa_indices[:, 0] + idx_full = swa_full.prefill_swa_indices[prefill_rows, 0] + # Untrimmed request: identical to the full batch. + assert torch.equal(lens[WINDOW:], lens_full[WINDOW:]) + assert torch.equal(idx[WINDOW:], idx_full[WINDOW:]) + # Trimmed request: the window grows from the replay window's start, so its + # first token sees only itself and the last sees the full window as before. + assert lens[:WINDOW].tolist() == list(range(1, WINDOW + 1)) + assert torch.equal(idx[WINDOW - 1], idx_full[WINDOW - 1]) + assert idx[0, 0] == idx_full[0, WINDOW - 1] and idx[0, 1] == -1 + + mla, mla_full = replay["mla"], full["mla"] + assert torch.equal(mla.req_id_per_token, mla_full.req_id_per_token[rows]) + assert mla.query_start_loc.tolist() == [0, 1, 129, 229] + + idx_md, idx_full_md = replay["idx"].prefill.chunks[0], full["idx"].prefill.chunks[0] + assert torch.equal(idx_md.cu_seqlen_ks, idx_full_md.cu_seqlen_ks[prefill_rows]) + assert torch.equal(idx_md.cu_seqlen_ke, idx_full_md.cu_seqlen_ke[prefill_rows]) + assert torch.equal(idx_md.token_to_seq, idx_full_md.token_to_seq) + + +def test_run_gathers_states_and_realigns_shared_indexer_buffers(): + cfg = _vllm_config() + full = _full_metadata(cfg) + topk = torch.arange(401 * 4, device=DEVICE).view(401, 4).int() + candidates = torch.arange(401 * 3, device=DEVICE).view(401, 3).int() + topk_before, candidates_before = topk.clone(), candidates.clone() + seen = {} + + def run_layers(hidden_states, *rest): + seen["hidden_states"] = hidden_states.clone() + seen["replay_metadata"] = get_forward_context().attn_metadata + return (hidden_states, rest[2]) # pre_mix + + layers = _replay_layers(run_layers, topk, candidates) + hidden = torch.arange(401, device=DEVICE, dtype=torch.float32)[:, None] + states = (hidden, hidden.long(), None, hidden, hidden, hidden, hidden) + context = ForwardContext(no_compile_layers={}, attn_metadata=full, slot_mapping={}) + full_swa_indices = full["swa"].prefill_swa_indices.clone() + with override_forward_context(context): + outputs = layers(*states) + + rows = torch.tensor(REPLAY_ROWS, device=DEVICE) + assert torch.equal(seen["hidden_states"], hidden[rows]) + # The replay was built by a private builder: the runner's stays untouched. + assert seen["replay_metadata"]["swa"].num_prefill_tokens == 228 + assert torch.equal(full["swa"].prefill_swa_indices, full_swa_indices) + assert context.attn_metadata is full + assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows]) + assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows]) + # Results come back on full-batch rows; dropped rows are zero. + assert outputs[0].shape[0] == 401 + assert torch.equal(outputs[0][rows], hidden[rows]) + assert outputs[0][1:173].abs().sum() == 0 + + +def _fake_attention(x: torch.Tensor, out: torch.Tensor) -> None: + """Stands in for the eager attention break: reads the metadata in the + forward context at run time, like the real kernels do.""" + + def run() -> None: + swa = get_forward_context().attn_metadata["swa"] + # Real kernels cover the metadata's token count; padding rows stay. + n = swa.num_decode_tokens + swa.num_prefill_tokens + req = swa.token_to_req_indices[:n].to(x.dtype) + out[:n] = x[:n] + req[:, None] + + capture = BreakableCUDAGraphCapture.current() + if capture is not None and capture.capturing: + capture.add_eager(run) + else: + run() + + +def _fake_replay_layers(hidden, positions, _, pre_mix, post_mix, res_mix, residual): + # Like the window KV insert, this captured op takes the metadata's slot + # mapping by address: it must live in a buffer that is refilled per step. + slots = get_forward_context().attn_metadata["swa"].slot_mapping + x = hidden * 2 + slots[: hidden.shape[0], None].to(hidden.dtype) + out = torch.empty_like(x) + _fake_attention(x, out) + return (out + positions[:, None].to(out.dtype), pre_mix + residual) + + +def _states(seed: int): + g = torch.Generator(device=DEVICE).manual_seed(seed) + hidden = torch.randn(401, 3, device=DEVICE, generator=g) + return ( + hidden, + torch.arange(401, device=DEVICE), + None, + hidden + 1, + hidden + 2, + hidden + 3, + hidden + 4, + ) + + +def _context(metadata): + return ForwardContext( + no_compile_layers={}, + attn_metadata=metadata, + slot_mapping={}, + is_padding=torch.zeros(401, dtype=torch.bool, device=DEVICE), + ) + + +def test_late_layer_graphs_match_eager(): + """The replay runs as an eager break of the model graph and replays its own + graph; both must match the eager path on a batch the graphs never saw.""" + from vllm.platforms import current_platform + from vllm.utils.torch_utils import _current_stream_tls + + cfg = _vllm_config() + swa_b, mla_b, idx_b = _builders(cfg) + # Capture-time batch: nothing to trim, like the runner's dummy batches. + dummy = { + "swa": swa_b.build(0, _common([101, 100, 100, 100], [101, 100, 100, 100], 32)), + "mla": mla_b.build(0, _common([101, 100, 100, 100], [101, 100, 100, 100], 128)), + "idx": idx_b.build(0, _common([101, 100, 100, 100], [101, 100, 100, 100], 128)), + } + real = _full_metadata(cfg) + eager = _replay_layers(_fake_replay_layers, None, None) + graphed = _replay_layers(_fake_replay_layers, None, None, graph_sizes=[512]) + states = _states(0) + + prev_stream = getattr(_current_stream_tls, "value", None) + stream = torch.cuda.Stream() + try: + with torch.cuda.stream(stream): + with override_forward_context(_context(dummy)): + graphed(*states) # profile run: allocates the static buffers + outer = BreakableCUDAGraphCapture( + current_platform.get_global_graph_pool() + ) + with outer: + hidden_out, pre_mix_out = graphed(*states) + # Replay on a different, trimmed batch. + new_states = _states(1) + for dst, src in zip(states, new_states): + if dst is not None: + dst.copy_(src) + with override_forward_context(_context(real)): + outer.replay() + torch.accelerator.synchronize() + expected = eager(*states) + assert torch.equal(hidden_out, expected[0]) + assert torch.equal(pre_mix_out, expected[1]) + # Eager model forward, graphed replay. + outputs = graphed(*states) + assert torch.equal(outputs[0], expected[0]) + assert torch.equal(outputs[1], expected[1]) + finally: + torch.cuda.current_stream().wait_stream(stream) + _current_stream_tls.value = prev_stream diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index e792d68d89e0..5e824f52148d 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -165,6 +165,44 @@ def test_nested_capture_raises(cuda_capture_stream): pass +def test_eager_break_may_capture_its_own_graphs(cuda_capture_stream): + """An eager break runs between outer segments, so it may capture (and on + replay, replay) graphs of its own, e.g. for a differently shaped sub-batch.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + x = torch.zeros(4, device="cuda") + y = torch.zeros(4, device="cuda") + outer = BreakableCUDAGraphCapture() + inner = BreakableCUDAGraphCapture() + + def eager_step(): + if inner.num_graphs == 0: + # Capture time: the outer capture is current but paused. + assert BreakableCUDAGraphCapture.current() is outer + with inner: + y.add_(x) + assert BreakableCUDAGraphCapture.current() is outer + else: + assert BreakableCUDAGraphCapture.current() is None + inner.replay() + + with outer: + x.add_(1.0) + outer.add_eager(eager_step) + x.add_(1.0) + assert BreakableCUDAGraphCapture.current() is None + assert inner.num_graphs == 1 + + outer.replay() + torch.accelerator.synchronize() + assert x.tolist() == [2.0] * 4 + assert y.tolist() == [1.0] * 4 + outer.replay() + torch.accelerator.synchronize() + assert x.tolist() == [4.0] * 4 + assert y.tolist() == [4.0] * 4 + + def test_active_state_isolated_across_threads(cuda_capture_stream): """Verify the thread-local 'active capture' slot is per-thread. diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index ca6aa1713774..83e242982127 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -156,12 +156,22 @@ def __init__(self, pool: Any | None = None) -> None: self._num_eager_breaks: int = 0 self._current_graph: torch.cuda.CUDAGraph | None = None self._capturing: bool = False + self._outer: BreakableCUDAGraphCapture | None = None + + @property + def capturing(self) -> bool: + """Whether a graph segment is being captured right now.""" + return self._capturing # --- context manager protocol ---------------------------------------- def __enter__(self) -> BreakableCUDAGraphCapture: - if getattr(BreakableCUDAGraphCapture._tls, "active", None) is not None: + outer = BreakableCUDAGraphCapture.current() + if outer is not None and outer._capturing: raise RuntimeError("Nested BreakableCUDAGraphCapture is not supported.") + # An eager break of an outer capture may capture graphs of its own + # (see add_eager); the outer capture becomes current again on exit. + self._outer = outer BreakableCUDAGraphCapture._tls.active = self self._begin_segment() return self @@ -170,7 +180,8 @@ def __exit__(self, exc_type, exc, tb) -> None: try: self._end_segment() finally: - BreakableCUDAGraphCapture._tls.active = None + BreakableCUDAGraphCapture._tls.active = self._outer + self._outer = None # --- segment management ---------------------------------------------- @@ -332,6 +343,10 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return self.runnable(*args, **kwargs) assert batch_descriptor is not None + return self.run(batch_descriptor, *args, **kwargs) + + def run(self, batch_descriptor: BatchDescriptor, *args: Any, **kwargs: Any) -> Any: + """Capture on the first call with ``batch_descriptor``, replay after.""" entry = self.entries.get(batch_descriptor) if entry is None: entry = _BreakableEntry(batch_descriptor=batch_descriptor) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index d5f6bc8a111c..e73becaccbfe 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -228,7 +228,8 @@ class CacheConfig: swa_bounded_replay: bool = True """Keep the sliding-window KV of models that support it (DeepSeek-V4.1) out of prefix caching and rebuild it after a prefix hit by recomputing the - hit's last window. Requires model runner V2.""" + hit's last window; layers past the last KV-source layer then also prefill + only each request's trailing window. Requires model runner V2.""" kv_cache_memory_bytes: int | None = None """Size of KV Cache per GPU in bytes. By default, this is set to None diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py new file mode 100644 index 000000000000..13e1608ebdfa --- /dev/null +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decoder-side SWA bounded replay. + +Layers past the last KV-source layer (the replay layers) own nothing but their +sliding-window KV, and decode reads only the trailing ``window`` positions of +it. In prefill they therefore run on each request's last ``window`` tokens (its +replay window), attending the source's full compressed KV plus their own window, +floored at the replay window's start since earlier positions hold no window KV +for these layers. + +``ReplayBatchBuilder`` picks the rows (a ``ReplayInputBatch``), +``ReplayMetadataBuilder`` rebuilds the attention metadata for them, +``ReplayCudaGraphs`` runs the layers under graphs of their own when the model +itself runs under piecewise (breakable) graphs, and ``DecoderReplayLayers`` +ties the three into the model's forward. +""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + BreakableCUDAGraphWrapper, +) +from vllm.config import VllmConfig +from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.logger import init_logger +from vllm.utils.torch_utils import weak_ref_tensor +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool + +logger = init_logger(__name__) + +# hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual +States = tuple[torch.Tensor | None, ...] + + +@dataclass +class ReplayInputBatch: + """The batch the replay layers run on: each request's replay window rows.""" + + rows: torch.Tensor # [num_tokens] rows of the full batch + query_start_loc: torch.Tensor # [num_reqs + 1] + query_start_loc_cpu: torch.Tensor + replay_start: torch.Tensor # [num_reqs] lowest window position per request + positions: torch.Tensor | None # [num_tokens] + num_tokens: int + max_query_len: int + num_batch_tokens: int # rows of the full batch, for scattering results back + trims: bool # False when every request already fits in the window + + def gather(self, t: torch.Tensor | None) -> torch.Tensor | None: + return None if t is None else t.index_select(0, self.rows) + + def scatter(self, t: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: + """Place replay rows back into a zero-filled full-batch tensor.""" + if out is None: + out = t.new_zeros((self.num_batch_tokens, *t.shape[1:])) + else: + out.zero_() + return out.index_copy_(0, self.rows, t) + + +class ReplayBatchBuilder: + """Selects each request's replay window rows from the batch's metadata. + + The batch is laid out on the CPU each step and read by the GPU through UVA. + """ + + def __init__(self, window: int, max_num_tokens: int, max_num_reqs: int) -> None: + self.window = window + self._rows = UvaBufferPool(max_num_tokens, torch.int64) + self._dropped = UvaBufferPool(max_num_reqs + 1, torch.int32) + self._window_start = UvaBufferPool(max_num_reqs, torch.int32) + + @staticmethod + def query_lens(common: Any) -> np.ndarray: + return np.diff(common.query_start_loc_cpu[: common.num_reqs + 1].numpy()) + + def trims(self, common: Any) -> bool: + return bool((self.query_lens(common) > self.window).any()) + + def build(self, common: Any, num_batch_tokens: int) -> ReplayInputBatch: + num_reqs = common.num_reqs + query_start_loc = common.query_start_loc_cpu[: num_reqs + 1].numpy() + lens = self.query_lens(common) + trimmed = lens > self.window + keep = np.minimum(lens, self.window) + new_query_start_loc = np.zeros(num_reqs + 1, dtype=np.int32) + np.cumsum(keep, out=new_query_start_loc[1:]) + num_tokens = int(new_query_start_loc[-1]) + + # Kept row r of request i is window_start[i] + (r - kept_start[i]). + window_start = query_start_loc[1:] - keep + kept_start = new_query_start_loc[:-1] + rows = self._rows.copy_to_uva( + np.repeat(window_start - kept_start, keep) + np.arange(num_tokens) + ) + # Boundaries: the device ones minus the rows trimmed before them. Only + # prefill rows are trimmed and their CPU lengths are exact; adaptive + # verification resizes the leading decode requests on the GPU alone. + dropped_before = self._dropped.copy_to_uva( + query_start_loc - new_query_start_loc + ) + replay_query_start_loc = common.query_start_loc[: num_reqs + 1] - dropped_before + + # A trimmed request holds no replay-layer window KV below its replay + # window, on top of whatever the encoder-side replay already excludes. + seq_lens = common.seq_lens_cpu_upper_bound[:num_reqs].numpy() + replay_start = self._window_start.copy_to_uva( + np.where(trimmed, seq_lens - self.window, 0) + ) + if common.replay_start is not None: + replay_start = torch.maximum(replay_start, common.replay_start[:num_reqs]) + + return ReplayInputBatch( + rows=rows, + query_start_loc=replay_query_start_loc, + query_start_loc_cpu=torch.from_numpy(new_query_start_loc), + replay_start=replay_start, + positions=None + if common.positions is None + else common.positions.index_select(0, rows), + num_tokens=num_tokens, + max_query_len=int(keep.max()), + num_batch_tokens=num_batch_tokens, + trims=bool(trimmed.any()), + ) + + +class ReplayMetadataBuilder: + """Builds the replay layers' attention metadata for a batch. + + Private clones of the runner's builders do the building: the runner's back + the full batch's metadata with persistent buffers, and that metadata + outlives the forward (the runner hands it to the speculator). + """ + + def __init__(self, source_attn: Any, replay_attn: list[Any], max_num_tokens: int): + self.compressed_prefix = source_attn.prefix + self.indexer_prefix = source_attn.indexer.k_cache.prefix + self.swa_prefixes = [attn.swa_cache_layer.prefix for attn in replay_attn] + self._max_num_tokens = max_num_tokens + self._builders: dict[Any, Any] = {} # clones, by the runner's builder + # Compacted slot mappings, one fixed buffer per KV-cache group (keyed + # by the first layer seen in it): the window KV insert runs inside the + # replay-layer graph and reads them by address. + self._slot_mappings: dict[str, torch.Tensor] = {} + + def common(self, attn_metadata: Any) -> Any: + """The batch's CommonAttentionMetadata, as the replay layers saw it.""" + assert isinstance(attn_metadata, dict) + return attn_metadata[self.swa_prefixes[0]].common + + def build( + self, attn_metadata: Any, batch: ReplayInputBatch, num_padded: int + ) -> dict[str, Any]: + """The metadata dict the replay layers run under, ``num_padded`` rows.""" + assert isinstance(attn_metadata, dict) + replay = dict(attn_metadata) + built: dict[Any, Any] = {} + for prefix in self.swa_prefixes: + full = attn_metadata[prefix] + if full.builder not in built: + built[full.builder] = self._builder_for(full).build( + 0, self._compact(full, batch, num_padded, prefix) + ) + replay[prefix] = built[full.builder] + compressed = attn_metadata[self.compressed_prefix] + source_common = self._compact( + compressed, batch, num_padded, self.compressed_prefix + ) + replay[self.compressed_prefix] = self._builder_for(compressed).build( + 0, source_common + ) + # The indexer K cache shares the source's KV group, hence its metadata. + indexer = attn_metadata[self.indexer_prefix] + replay[self.indexer_prefix] = self._builder_for(indexer).build(0, source_common) + return replay + + def _compact( + self, full: Any, batch: ReplayInputBatch, num_padded: int, group: str + ) -> Any: + common = full.common + slot_mapping = self._slot_mappings.get(group) + if slot_mapping is None: + slot_mapping = common.slot_mapping.new_empty(self._max_num_tokens) + self._slot_mappings[group] = slot_mapping + torch.index_select( + common.slot_mapping, 0, batch.rows, out=slot_mapping[: batch.num_tokens] + ) + if num_padded > batch.num_tokens: + # Padding rows write nowhere, like the runner's own padding. + slot_mapping[batch.num_tokens : num_padded] = PAD_SLOT_ID + return common.replace_tokens( + query_start_loc=batch.query_start_loc, + query_start_loc_cpu=batch.query_start_loc_cpu, + num_actual_tokens=batch.num_tokens, + max_query_len=batch.max_query_len, + slot_mapping=slot_mapping[:num_padded], + replay_start=batch.replay_start, + positions=batch.positions, + ) + + def _builder_for(self, full: Any) -> Any: + src = full.builder + if src not in self._builders: + self._builders[src] = src.clone() + return self._builders[src] + + +class ReplayCudaGraphs: + """Breakable graphs of the replay layers, keyed by padded replay size. + + They are captured while the model graph is (the outer capture is paused in + an eager break), on inputs and into outputs at fixed addresses. + """ + + def __init__( + self, + run_layers: Callable[..., tuple[torch.Tensor, ...]], + vllm_config: VllmConfig, + sizes: list[int], + ) -> None: + self.sizes = sorted(sizes) + self.wrapper = BreakableCUDAGraphWrapper(run_layers, vllm_config) + self._inputs: list[torch.Tensor | None] | None = None + self._outputs: list[torch.Tensor] | None = None + self._is_padding: torch.Tensor | None = None + + @property + def allocated(self) -> bool: + return self._outputs is not None + + def allocate(self, states: States, outputs: tuple[torch.Tensor, ...]) -> None: + """Size the fixed input and output buffers from an eager forward.""" + size = self.sizes[-1] + self._inputs = [ + None if t is None else t.new_zeros((size, *t.shape[1:])) for t in states + ] + self._outputs = [out.new_zeros((size, *out.shape[1:])) for out in outputs] + self._is_padding = torch.zeros(size, dtype=torch.bool, device=outputs[0].device) + + def outputs(self, num_tokens: int) -> tuple[torch.Tensor, ...]: + """Fixed-address outputs for the graph segments after the replay.""" + assert self._outputs is not None + return tuple(out[:num_tokens] for out in self._outputs) + + def size_for(self, num_tokens: int) -> int | None: + """Padded size whose graph exists or can be captured now, if any.""" + size = next((s for s in self.sizes if s >= num_tokens), None) + if size is None: + return None + captured = BatchDescriptor(num_tokens=size) in self.wrapper.entries + if captured or BreakableCUDAGraphCapture.current() is not None: + return size + return None + + def padding_mask( + self, num_tokens: int, size: int, is_padding: torch.Tensor | None + ) -> torch.Tensor: + assert self._is_padding is not None + mask = self._is_padding[:size] + mask[:num_tokens] = False if is_padding is None else is_padding + mask[num_tokens:] = True + return mask + + def run( + self, size: int, batch: ReplayInputBatch, states: States + ) -> tuple[torch.Tensor, ...]: + """Run the graph of ``size`` rows on the batch's rows of ``states``.""" + assert self._inputs is not None + for buf, t in zip(self._inputs, states): + if buf is not None: + torch.index_select(t, 0, batch.rows, out=buf[: batch.num_tokens]) + inputs = [None if buf is None else buf[:size] for buf in self._inputs] + outputs = self.wrapper.run(BatchDescriptor(num_tokens=size), *inputs) + return tuple(out[: batch.num_tokens] for out in outputs) + + +class DecoderReplayLayers: + """Runs the replay layers on each request's replay window (see module doc). + + ``run_layers`` takes a batch's layer inputs (hidden states, positions, input + ids, mHC states) and returns ``(hidden_states, pre_mix, *aux_hidden_states)``. + """ + + def __init__( + self, + vllm_config: VllmConfig, + window: int, + source_attn: Any, + replay_attn: list[Any], + run_layers: Callable[..., tuple[torch.Tensor, ...]], + graph_sizes: list[int], + ) -> None: + logger.info_once( + "Decoder SWA bounded replay: layers past the last KV source prefill " + "only each request's last %d tokens.", + window, + ) + self.run_layers = run_layers + # Per-row indexer outputs the source publishes for the layers after it. + self.row_buffers = [ + buf + for buf in ( + source_attn.topk_indices_buffer, + source_attn.candidate_block_buffer, + ) + if buf is not None + ] + scheduler_config = vllm_config.scheduler_config + max_num_tokens = scheduler_config.max_num_batched_tokens + self.batch_builder = ReplayBatchBuilder( + window, max_num_tokens, scheduler_config.max_num_seqs + ) + self.metadata = ReplayMetadataBuilder(source_attn, replay_attn, max_num_tokens) + self.graphs = ( + ReplayCudaGraphs(run_layers, vllm_config, graph_sizes) + if graph_sizes + else None + ) + + def __call__(self, *states: torch.Tensor | None) -> tuple[torch.Tensor, ...]: + hidden_states = states[0] + assert hidden_states is not None + num_tokens = hidden_states.shape[0] + + outer = BreakableCUDAGraphCapture.current() + if outer is not None and outer.capturing: + # Piecewise capture of the whole model: the replay is an eager break, + # so each replay plans from its own batch. Its outputs must sit at + # fixed addresses for the graph segments after it. + assert self.graphs is not None + outputs = self.graphs.outputs(num_tokens) + weak_states = tuple(weak_ref_tensor(t) for t in states) + outer.add_eager(lambda: self._run(weak_states, outputs)) + return outputs + + attn_metadata = get_forward_context().attn_metadata + common = ( + self.metadata.common(attn_metadata) + if isinstance(attn_metadata, dict) + else None + ) + if common is not None and self.batch_builder.trims(common): + assert not torch.cuda.is_current_stream_capturing() + outputs = self._run(states) + else: + # Decode batches (including full-graph captures) and metadata-less + # warmup runs. + outputs = self.run_layers(*states) + if self.graphs is not None and not self.graphs.allocated: + # The first forward is the runner's eager profile run. + assert not torch.cuda.is_current_stream_capturing() + self.graphs.allocate(states, outputs) + return outputs + + def _run( + self, states: States, outputs: tuple[torch.Tensor, ...] | None = None + ) -> tuple[torch.Tensor, ...]: + """Run the replay layers on this batch's replay rows and scatter the results + back to full-batch rows, into ``outputs`` when given.""" + hidden_states = states[0] + assert hidden_states is not None + forward_context = get_forward_context() + batch = self.batch_builder.build( + self.metadata.common(forward_context.attn_metadata), hidden_states.shape[0] + ) + replay_outputs = self._run_replay(batch, states) + if outputs is None: + return tuple(batch.scatter(t) for t in replay_outputs) + for out, t in zip(outputs, replay_outputs): + batch.scatter(t, out=out) + return outputs + + def _run_replay( + self, batch: ReplayInputBatch, states: States + ) -> tuple[torch.Tensor, ...]: + forward_context = get_forward_context() + saved = (forward_context.attn_metadata, forward_context.is_padding) + size = self.graphs.size_for(batch.num_tokens) if self.graphs else None + try: + # A graph always runs on rebuilt metadata, so what its captured + # kernels read by address is this step's data. + if batch.trims or size is not None: + forward_context.attn_metadata = self.metadata.build( + saved[0], batch, size or batch.num_tokens + ) + if batch.trims: + for buf in self.row_buffers: + buf[: batch.num_tokens].copy_(batch.gather(buf)) + is_padding = batch.gather(saved[1]) + if size is None: + forward_context.is_padding = is_padding + return self.run_layers(*(batch.gather(t) for t in states)) + assert self.graphs is not None + forward_context.is_padding = self.graphs.padding_mask( + batch.num_tokens, size, is_padding + ) + return self.graphs.run(size, batch, states) + finally: + forward_context.attn_metadata, forward_context.is_padding = saved diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index f3530544da7b..bdae1ec6cb0c 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -9,6 +9,7 @@ import torch.nn as nn import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled from vllm.config import VllmConfig from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( @@ -65,6 +66,7 @@ make_deepseek_v4_expert_params_mapping, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -532,6 +534,34 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): prefix=f"{prefix}.layers", ) + # Decoder-side SWA bounded replay: layers past the last KV source + # prefill only each request's trailing window. + self.decoder_replay_layers: DecoderReplayLayers | None = None + self.decoder_replay_start = self.end_layer + cut = max(config.kv_source_layer_ids) + if ( + cut < self.end_layer - 1 + and self._decoder_replay_supported(vllm_config, cut) + and self.layers[cut].attn.swa_cache_layer.bounded_replay + ): + self.decoder_replay_start = cut + 1 + # Batch sizes the runner captures piecewise (breakable) graphs for. + compilation_config = vllm_config.compilation_config + graph_sizes: list[int] = [] + if ( + compilation_config.cudagraph_mode.has_piecewise_cudagraphs() + and is_breakable_cudagraph_enabled() + ): + graph_sizes = list(compilation_config.cudagraph_capture_sizes or []) + self.decoder_replay_layers = DecoderReplayLayers( + vllm_config, + config.sliding_window, + self.layers[cut].attn, + [self.layers[i].attn for i in range(cut + 1, self.end_layer)], + self._run_decoder_replay_layers, + graph_sizes, + ) + # The n-gram hash needs a slot-keyed rolling store of compressed ids # (chunked prefill / decode lookback); key it off the first local # layer's sliding-window KV cache. Only PP ranks owning an engram @@ -700,14 +730,23 @@ def forward( if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - # Every layer's post runs inside the next layer's fused pre, so aux - # hidden states are read back from there instead of recomputed. aux_hidden_by_layer: dict[int, torch.Tensor] = {} - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer( + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_by_layer, + engram_hashes, + engram_mask, + ) + late_aux: list[torch.Tensor] = [] + if self.decoder_replay_layers is not None: + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, input_ids, @@ -715,31 +754,21 @@ def forward( post_mix, res_mix, residual, - engram_hashes, - engram_mask, - capture_previous_aux=idx in self.aux_hidden_state_layers, ) - if previous_aux is not None: - # idx is the one-based id of the layer whose post this is. - if self.use_sequence_parallel: - previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] - aux_hidden_by_layer[idx] = previous_aux - if layer is not None: - # The last layer has no successor to fold its post into. - hidden_states = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix + else: + hidden_states = self._collapse( + hidden_states, + residual, + post_mix, + res_mix, + aux_hidden_by_layer, + full_num_tokens, ) - if self.end_layer in self.aux_hidden_state_layers: - final_aux = hidden_states.mean(dim=1) - if self.use_sequence_parallel: - final_aux = sp_all_gather(final_aux)[:full_num_tokens] - aux_hidden_by_layer[self.end_layer] = final_aux - aux_hidden_states = [ aux_hidden_by_layer[layer_id] for layer_id in self.aux_hidden_state_layers if layer_id in aux_hidden_by_layer - ] + ] + late_aux if not get_pp_group().is_last_rank: return IntermediateTensors( @@ -768,6 +797,147 @@ def forward( return hidden_states, aux_hidden_states return hidden_states + def _run_layers( + self, + layer_ids: range, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + aux_hidden_by_layer: dict[int, torch.Tensor], + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Every layer's post runs inside the next layer's fused pre, so aux + # hidden states are read back from there instead of recomputed. + full_num_tokens = positions.shape[0] + for idx in layer_ids: + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = ( + self.layers[idx]( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + capture_previous_aux=idx in self.aux_hidden_state_layers, + ) + ) + if previous_aux is not None: + # idx is the one-based id of the layer whose post this is. + if self.use_sequence_parallel: + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] + aux_hidden_by_layer[idx] = previous_aux + return hidden_states, residual, post_mix, res_mix, pre_mix + + def _collapse( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + aux_hidden_by_layer: dict[int, torch.Tensor], + full_num_tokens: int, + ) -> torch.Tensor: + # The last layer has no successor to fold its post into. + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) + if self.end_layer in self.aux_hidden_state_layers: + final_aux = hidden_states.mean(dim=1) + if self.use_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_by_layer[self.end_layer] = final_aux + return hidden_states + + def _run_decoder_replay_layers( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Layers past the last KV source, on whatever rows they are given.""" + aux_hidden_by_layer: dict[int, torch.Tensor] = {} + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( + range(self.decoder_replay_start, self.end_layer), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_by_layer, + ) + hidden_states = self._collapse( + hidden_states, + residual, + post_mix, + res_mix, + aux_hidden_by_layer, + positions.shape[0], + ) + return ( + hidden_states, + pre_mix, + *( + aux_hidden_by_layer[layer_id] + for layer_id in self.aux_hidden_state_layers + if layer_id in aux_hidden_by_layer + ), + ) + + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: + """Whether this rank may trim the layers after ``cut``; warns when not.""" + parallel_config = vllm_config.parallel_config + spec_config = vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + draft_hf_config = getattr(draft_config, "hf_config", None) + draft_window = getattr(draft_hf_config, "sliding_window", None) + draft_layer_types = getattr(draft_hf_config, "layer_types", None) or () + window = self.config.sliding_window + if self.start_layer > cut or self.end_layer < self.config.num_hidden_layers: + # Later ranks would run their replay layers on the scattered zero + # rows with an unfloored window. + reason = ( + "the pipeline stage holding the last KV source layer must also " + "hold every layer after it" + ) + elif ( + self.use_sequence_parallel + or parallel_config.data_parallel_size > 1 + or parallel_config.prefill_context_parallel_size > 1 + ): + reason = ( + "the replay-layer batch shrinks per rank, which sequence, data " + "and prefill-context parallelism cannot follow" + ) + elif any(i > cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits after the last KV source layer" + elif draft_config is not None and ( + draft_window is None + or draft_window > window + or any(t != "sliding_attention" for t in draft_layer_types) + ): + # The drafter derives its context KV from the replay layers' hidden + # states, which exist only for the last `sliding_window` tokens. + reason = ( + f"the drafter (sliding window {draft_window}) reads hidden states " + f"outside the target's {window}-token window" + ) + else: + return True + logger.warning_once("Decoder SWA bounded replay is off: %s.", reason) + return False + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 117b74eb41a2..fe9902e44c20 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -139,6 +139,8 @@ class DeepseekV4FlashMLAMetadata(AttentionMetadata): req_id_per_token: torch.Tensor block_size: int topk_tokens: int + common: Any = None + builder: Any = None class DeepseekV4SparseMLAMetadataBuilder( @@ -215,6 +217,8 @@ def build( req_id_per_token=req_id_per_token, block_size=self.kv_cache_spec.block_size, topk_tokens=self.topk_tokens, + common=cm, + builder=self, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 82e9fa691f59..fb9ad0a551cb 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -8,6 +8,7 @@ import numpy as np import torch +from typing_extensions import Self from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic64Sym, @@ -487,6 +488,19 @@ def naive_query_lens(self) -> torch.Tensor: def replace(self, **kwargs) -> "CommonAttentionMetadata": return replace(self, **kwargs) + def replace_tokens(self, **fields) -> "CommonAttentionMetadata": + """A copy over another set of this batch's tokens (same requests): the + fields given replace the current ones, and what was derived from the + token set is dropped.""" + return replace( + self, + logits_indices_padded=None, + num_logits_indices=None, + _num_computed_tokens_cache=None, + _token_to_req_indices_cache=None, + **fields, + ) + def compute_num_computed_tokens(self) -> torch.Tensor: """Compute num_computed_tokens on device (seq_lens - query_lens).""" if self._num_computed_tokens_cache is None: @@ -612,6 +626,19 @@ def __init__( def set_kernel_block_size(self, kernel_block_size: int) -> None: self.kernel_block_size = kernel_block_size + def clone(self, **kwargs) -> Self: + """A builder constructed like this one, with buffers of its own.""" + clone = type(self)( + self.kv_cache_spec, + self.layer_names, + self.vllm_config, + self.device, + **kwargs, + ) + if self.kernel_block_size is not None: + clone.set_kernel_block_size(self.kernel_block_size) + return clone + @classmethod def get_cudagraph_support( cls: type["AttentionMetadataBuilder"], diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 7c3580f08441..ddee86f970f8 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -5,6 +5,7 @@ import numpy as np import torch +from typing_extensions import Self import vllm.envs as envs from vllm.config import VllmConfig @@ -607,6 +608,7 @@ class DeepseekV32IndexerMetadata: decode: DeepSeekV32IndexerDecodeMetadata | None = None prefill: DeepseekV32IndexerPrefillMetadata | None = None + builder: Any = None def compute_kpool_tail_slot_mapping( @@ -764,6 +766,7 @@ def get_cudagraph_support( def __init__(self, *args, block_table_width: int, **kwargs) -> None: super().__init__(*args, **kwargs) + self.block_table_width = block_table_width scheduler_config = self.vllm_config.scheduler_config parallel_config = self.vllm_config.parallel_config self.dcp_world_size = parallel_config.decode_context_parallel_size @@ -881,6 +884,9 @@ def __init__(self, *args, block_table_width: int, **kwargs) -> None: self.indexer_decode_block_table_buffer: torch.Tensor | None = None self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens + def clone(self, **kwargs) -> Self: + return super().clone(block_table_width=self.block_table_width, **kwargs) + def _dcp_localize_decode_seq_lens( self, seq_lens: torch.Tensor, @@ -1505,6 +1511,7 @@ def build( num_prefill_tokens=num_prefill_tokens, prefill=prefill_metadata, decode=decode_metadata, + builder=self, ) return attn_metadata diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index e5e2413c822c..c38b2ee00726 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -214,6 +214,10 @@ class DeepseekSparseSWAMetadata: prefill_max_model_len: int = 0 prefill_max_num_batched_tokens: int = 0 + # What this was built from, so a model can rebuild it for a token subset. + common: Any = None + builder: Any = None + # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta # per present DeepseekV4 layer type, shared across all ~60 layers of that type # within a decode step. The first forward call of a given type triggers @@ -764,6 +768,8 @@ def build( tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A], tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A], tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A], + common=common_attn_metadata, + builder=self, **deepseek_v4_fields, # type: ignore[arg-type] ) From 3e8f381b0b52339ca63766ef0bb2f09fb4511a04 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Sat, 12 Sep 2026 08:57:28 +0000 Subject: [PATCH 2/4] [SpecDecode] Bound the DSpark drafter's context KV to its window A DSpark drafter whose layers are all sliding-window attention can only ever read the last `window` context positions of a request, so projecting and inserting context KV for the rest of a long prefill chunk is wasted work. Under the target's decoder-side SWA bounded replay those rows also carry no real hidden state (the target never computed them), so restricting the precompute to the window removes the garbage KV they produced. It rides --swa-bounded-replay. Adaptive verification sizes the leading verification requests on the GPU, but those are never trimmed, so the row selection from the CPU query boundaries still holds. The DFlash context-KV precompute gains a row-selection hook that keeps every row by default; DSpark implements it. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Yifan Qiao --- .../spec_decode/test_dspark_context_rows.py | 42 +++++++++++++ .../gpu/spec_decode/dflash/speculator.py | 24 ++++++-- .../gpu/spec_decode/dspark/speculator.py | 60 +++++++++++++++++++ 3 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 tests/v1/spec_decode/test_dspark_context_rows.py diff --git a/tests/v1/spec_decode/test_dspark_context_rows.py b/tests/v1/spec_decode/test_dspark_context_rows.py new file mode 100644 index 000000000000..c00d99871158 --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_context_rows.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A sliding-window DSpark drafter only inserts context KV for the last +``window`` scheduled tokens of each request.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="UVA buffers") + + +def _rows(window, query_lens): + speculator = SimpleNamespace( + context_window=window, _context_row_pool=UvaBufferPool(512, torch.int64) + ) + batch = SimpleNamespace( + num_reqs=len(query_lens), + query_start_loc_np=np.concatenate([[0], np.cumsum(query_lens)]).astype( + np.int32 + ), + ) + return DSparkSpeculator._context_rows(speculator, batch) + + +def test_context_rows_keep_each_request_tail(): + rows = _rows(128, [1, 300, 100]) + assert rows.tolist() == [0, *range(301 - 128, 301), *range(301, 401)] + + +@pytest.mark.parametrize( + "window, query_lens", + [(128, [6, 6, 100]), (None, [1, 300])], + ids=["all fit the window", "no window"], +) +def test_context_rows_none_when_nothing_to_skip(window, query_lens): + assert _rows(window, query_lens) is None diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a5710829e0ff..c44abe13f5c5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -297,6 +297,10 @@ def _generate_draft( num_reqs, self.num_speculative_steps ) + def _context_rows(self, input_batch: InputBatch) -> torch.Tensor | None: + """Target rows whose context KV this drafter can still read; None for all.""" + return None + @torch.inference_mode() def propose( self, @@ -411,8 +415,11 @@ def propose( # because the context shape varies per step. During dummy runs the block tables # are placeholders, so we skip the cache write to avoid clobbering real entries. # Each layer uses the context slots of its own kv-cache group. + context_hidden = self.hidden_states[:num_target_tokens] + context_positions = self.context_positions[:num_target_tokens] + context_slots: torch.Tensor | list[torch.Tensor | None] | None if dummy_run: - context_slots: torch.Tensor | list[torch.Tensor | None] | None = None + context_slots = None elif self._layer_group_idx is not None: context_slots = [ self._context_slot_mappings[gidx][:num_target_tokens] @@ -420,10 +427,19 @@ def propose( ] else: context_slots = self._context_slot_mappings[0][:num_target_tokens] + rows = self._context_rows(input_batch) + if rows is not None: + context_hidden = context_hidden.index_select(0, rows) + context_positions = context_positions.index_select(0, rows) + if isinstance(context_slots, list): + context_slots = [ + None if t is None else t.index_select(0, rows) + for t in context_slots + ] + elif context_slots is not None: + context_slots = context_slots.index_select(0, rows) self.model.precompute_and_store_context_kv( - self.hidden_states[:num_target_tokens], - self.context_positions[:num_target_tokens], - context_slots, + context_hidden, context_positions, context_slots ) batch_sync, num_batch_tokens = ( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 411aa7f0415e..6be84b8ae0a5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -25,14 +25,21 @@ from typing import Any +import numpy as np import torch from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger +from vllm.v1.kv_cache_interface import KVCacheConfig +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, InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model +from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) @@ -113,6 +120,59 @@ def load_draft_model( self.use_acceptance_estimator = False return model + def set_attn( + self, + model_state: ModelState, + kv_cache_config: KVCacheConfig, + block_tables: BlockTables, + target_input_buffers: InputBuffers, + target_attn_groups: list[list[AttentionGroup]], + ) -> None: + super().set_attn( + model_state, + kv_cache_config, + block_tables, + target_input_buffers, + target_attn_groups, + ) + # DSV4.1 decoder-side SWA bounded replay, where draft context is never + # attended beyond the fixed window at the end. Adaptive verification + # splits the leading verification requests on the GPU only; they are + # never trimmed, so the CPU row selection holds. + windows = [ + getattr(self.attn_groups[gid][0].kv_cache_spec, "sliding_window", None) + for gid in self.draft_kv_cache_group_ids + ] + self.context_window: int | None = None + if None not in windows and self.vllm_config.cache_config.swa_bounded_replay: + assert len(set(windows)) == 1, ( + "All draft KV caches must have the same sliding_window." + ) + self.context_window = windows[0] + # Rows of each request's trailing window (see _context_rows). + self._context_row_pool = UvaBufferPool(self.max_num_tokens, torch.int64) + + def _context_rows(self, input_batch: InputBatch) -> torch.Tensor | None: + """Target rows whose draft context KV can still be read, or None for all. + + With a sliding-window drafter, a request's context beyond its last + ``context_window`` scheduled tokens is never attended (and under the + target's decoder-side replay was never computed), so it is skipped. + """ + window = self.context_window + if window is None: + return None + query_start_loc = input_batch.query_start_loc_np[: input_batch.num_reqs + 1] + lens = np.diff(query_start_loc) + if lens.max() <= window: + return None + keep = np.minimum(lens, window) + # Kept row r of request i is window_start[i] + (r - kept_start[i]). + window_start = query_start_loc[1:] - keep + kept_start = np.cumsum(keep) - keep + rows = np.repeat(window_start - kept_start, keep) + np.arange(int(keep.sum())) + return self._context_row_pool.copy_to_uva(rows) + def _sample_logits( self, logits: torch.Tensor, From 9f4699a9fef9a87a79049967741035383f56f81a Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Mon, 14 Sep 2026 09:52:12 +0000 Subject: [PATCH 3/4] [SpecDecode] Pass replay_start when capturing DFlash draft graphs The DFlash capture path builds the draft attention metadata from a dummy batch without the batch's replay_start, so a DeepSeek-V4.1 DSpark drafter, whose sliding-window caches replay like the target's, failed the SWA builder's assertion at graph capture. Pass the dummy batch's zeroed replay_start, as the runner's own capture and the drafter's per-step build already do. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Yifan Qiao --- vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index 725714fcdbc0..2bbf0c1c0f4d 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -66,6 +66,7 @@ def _prepare_dflash_inputs_to_capture( kv_cache_config=kv_cache_config, for_cudagraph_capture=True, causal=causal, + replay_start=input_batch.replay_start, ) return AttentionState(attn_metadata, slot_mappings_by_layer) From 096d8e9373e250788a600727b259d18b6340a6ad Mon Sep 17 00:00:00 2001 From: zjy0516 Date: Wed, 16 Sep 2026 14:31:23 +0000 Subject: [PATCH 4/4] [Refactor] Scope decoder replay CUDA graphs with ForwardContext Give decoder replay a local ForwardContext with its own batch descriptor, attention metadata, and padding mask. Use the normal graph wrapper call instead of a separate run entry point, preserving runtime mode dispatch and leaving the parent context unchanged. Validation: - 51 tests passed across decoder replay, breakable CUDA graphs, DSpark context rows, and DeepSeek V4 SWA visibility. - Pre-commit and mypy checks passed. - Vigil TP4 GB200 GSM8K, 1319 questions, 5-shot: original 89.92%/90.45%, modified 89.61%/90.45% for first/repeated passes; no API errors. Assisted-by: OpenAI Signed-off-by: zjy0516 --- .../test_deepseek_v4_decoder_replay_layers.py | 45 ++++++++++++++--- vllm/compilation/breakable_cudagraph.py | 4 -- .../deepseek_v41/decoder_replay_layers.py | 49 +++++++++++-------- 3 files changed, 67 insertions(+), 31 deletions(-) diff --git a/tests/models/test_deepseek_v4_decoder_replay_layers.py b/tests/models/test_deepseek_v4_decoder_replay_layers.py index 117d320e94ae..ec46ab4c38cf 100644 --- a/tests/models/test_deepseek_v4_decoder_replay_layers.py +++ b/tests/models/test_deepseek_v4_decoder_replay_layers.py @@ -6,13 +6,15 @@ start.""" from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture +from vllm.config import CUDAGraphMode from vllm.forward_context import ( + BatchDescriptor, ForwardContext, get_forward_context, override_forward_context, @@ -247,16 +249,24 @@ def test_run_gathers_states_and_realigns_shared_indexer_buffers(): def run_layers(hidden_states, *rest): seen["hidden_states"] = hidden_states.clone() - seen["replay_metadata"] = get_forward_context().attn_metadata + replay_context = get_forward_context() + seen["replay_metadata"] = replay_context.attn_metadata + assert replay_context is not context + assert replay_context.batch_descriptor.num_tokens == len(REPLAY_ROWS) + assert replay_context.is_padding.shape[0] == len(REPLAY_ROWS) + assert context.attn_metadata is full return (hidden_states, rest[2]) # pre_mix layers = _replay_layers(run_layers, topk, candidates) hidden = torch.arange(401, device=DEVICE, dtype=torch.float32)[:, None] states = (hidden, hidden.long(), None, hidden, hidden, hidden, hidden) - context = ForwardContext(no_compile_layers={}, attn_metadata=full, slot_mapping={}) + context = _context(full) full_swa_indices = full["swa"].prefill_swa_indices.clone() with override_forward_context(context): outputs = layers(*states) + assert get_forward_context() is context + assert context.batch_descriptor.num_tokens == 401 + assert context.is_padding.shape[0] == 401 rows = torch.tensor(REPLAY_ROWS, device=DEVICE) assert torch.equal(seen["hidden_states"], hidden[rows]) @@ -277,7 +287,10 @@ def _fake_attention(x: torch.Tensor, out: torch.Tensor) -> None: forward context at run time, like the real kernels do.""" def run() -> None: - swa = get_forward_context().attn_metadata["swa"] + context = get_forward_context() + assert context.batch_descriptor.num_tokens == x.shape[0] + assert context.is_padding.shape[0] == x.shape[0] + swa = context.attn_metadata["swa"] # Real kernels cover the metadata's token count; padding rows stay. n = swa.num_decode_tokens + swa.num_prefill_tokens req = swa.token_to_req_indices[:n].to(x.dtype) @@ -314,11 +327,13 @@ def _states(seed: int): ) -def _context(metadata): +def _context(metadata, mode=CUDAGraphMode.NONE): return ForwardContext( no_compile_layers={}, attn_metadata=metadata, slot_mapping={}, + cudagraph_runtime_mode=mode, + batch_descriptor=BatchDescriptor(num_tokens=401), is_padding=torch.zeros(401, dtype=torch.bool, device=DEVICE), ) @@ -348,17 +363,22 @@ def test_late_layer_graphs_match_eager(): with torch.cuda.stream(stream): with override_forward_context(_context(dummy)): graphed(*states) # profile run: allocates the static buffers + with override_forward_context(_context(dummy, CUDAGraphMode.PIECEWISE)): outer = BreakableCUDAGraphCapture( current_platform.get_global_graph_pool() ) with outer: hidden_out, pre_mix_out = graphed(*states) + assert set(graphed.graphs.wrapper.entries) == { + BatchDescriptor(num_tokens=512) + } # Replay on a different, trimmed batch. new_states = _states(1) for dst, src in zip(states, new_states): if dst is not None: dst.copy_(src) - with override_forward_context(_context(real)): + context = _context(real, CUDAGraphMode.PIECEWISE) + with override_forward_context(context): outer.replay() torch.accelerator.synchronize() expected = eager(*states) @@ -368,6 +388,19 @@ def test_late_layer_graphs_match_eager(): outputs = graphed(*states) assert torch.equal(outputs[0], expected[0]) assert torch.equal(outputs[1], expected[1]) + assert get_forward_context() is context + assert context.attn_metadata is real + assert context.batch_descriptor.num_tokens == 401 + assert context.is_padding.shape[0] == 401 + # A cached subgraph must not bypass the runtime's eager dispatch. + with ( + override_forward_context(_context(real)), + patch.object(graphed.graphs.wrapper, "_replay") as replay, + ): + outputs = graphed(*states) + replay.assert_not_called() + assert torch.equal(outputs[0], expected[0]) + assert torch.equal(outputs[1], expected[1]) finally: torch.cuda.current_stream().wait_stream(stream) _current_stream_tls.value = prev_stream diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 83e242982127..5c8074204c55 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/vllm/compilation/breakable_cudagraph.py @@ -343,10 +343,6 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return self.runnable(*args, **kwargs) assert batch_descriptor is not None - return self.run(batch_descriptor, *args, **kwargs) - - def run(self, batch_descriptor: BatchDescriptor, *args: Any, **kwargs: Any) -> Any: - """Capture on the first call with ``batch_descriptor``, replay after.""" entry = self.entries.get(batch_descriptor) if entry is None: entry = _BreakableEntry(batch_descriptor=batch_descriptor) diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py index 13e1608ebdfa..7d2551e93301 100644 --- a/vllm/models/deepseek_v41/decoder_replay_layers.py +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -17,7 +17,7 @@ """ from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any import numpy as np @@ -28,7 +28,11 @@ BreakableCUDAGraphWrapper, ) from vllm.config import VllmConfig -from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.forward_context import ( + BatchDescriptor, + get_forward_context, + override_forward_context, +) from vllm.logger import init_logger from vllm.utils.torch_utils import weak_ref_tensor from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -279,7 +283,7 @@ def run( if buf is not None: torch.index_select(t, 0, batch.rows, out=buf[: batch.num_tokens]) inputs = [None if buf is None else buf[:size] for buf in self._inputs] - outputs = self.wrapper.run(BatchDescriptor(num_tokens=size), *inputs) + outputs = self.wrapper(*inputs) return tuple(out[: batch.num_tokens] for out in outputs) @@ -383,26 +387,29 @@ def _run_replay( self, batch: ReplayInputBatch, states: States ) -> tuple[torch.Tensor, ...]: forward_context = get_forward_context() - saved = (forward_context.attn_metadata, forward_context.is_padding) size = self.graphs.size_for(batch.num_tokens) if self.graphs else None - try: - # A graph always runs on rebuilt metadata, so what its captured - # kernels read by address is this step's data. - if batch.trims or size is not None: - forward_context.attn_metadata = self.metadata.build( - saved[0], batch, size or batch.num_tokens - ) - if batch.trims: - for buf in self.row_buffers: - buf[: batch.num_tokens].copy_(batch.gather(buf)) - is_padding = batch.gather(saved[1]) + attn_metadata = forward_context.attn_metadata + # A graph always runs on rebuilt metadata, so what its captured + # kernels read by address is this step's data. + if batch.trims or size is not None: + attn_metadata = self.metadata.build( + attn_metadata, batch, size or batch.num_tokens + ) + if batch.trims: + for buf in self.row_buffers: + buf[: batch.num_tokens].copy_(batch.gather(buf)) + is_padding = batch.gather(forward_context.is_padding) + if size is not None: + assert self.graphs is not None + is_padding = self.graphs.padding_mask(batch.num_tokens, size, is_padding) + replay_context = replace( + forward_context, + attn_metadata=attn_metadata, + is_padding=is_padding, + batch_descriptor=BatchDescriptor(num_tokens=size or batch.num_tokens), + ) + with override_forward_context(replay_context): if size is None: - forward_context.is_padding = is_padding return self.run_layers(*(batch.gather(t) for t in states)) assert self.graphs is not None - forward_context.is_padding = self.graphs.padding_mask( - batch.num_tokens, size, is_padding - ) return self.graphs.run(size, batch, states) - finally: - forward_context.attn_metadata, forward_context.is_padding = saved