@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4b2e1cfa7b

[Model] Decoder-side SWA bounded replay for DeepSeek-V4.1 (#58132)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Yifan Qiao committed Sep 29, 2026 at 01:25 UTC 4b2e1cfa7b8094f1a9547aee41c39e05d2ce5c97
9 files changed +1001 -60
.buildkite/test_areas/models_basic.yaml
+3 -1
@@ -66,8 +66,10 @@ steps:
66 - tests/models/transformers/test_backend.py
67 - tests/models/test_registry.py
68 - tests/models/test_deepseek_v4_vl_rocm.py
69 + - tests/models/test_deepseek_v41_decoder_replay_layers.py
70 + - tests/models/test_deepseek_v41_replay_batch.py
71 commands:
70 - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
72 + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py
73 mirror:
74 amd:
75 label: ":amd: (MI355 DPX) Basic Models (Other)"
tests/distributed/test_engram_dp_shard.py
+17 -16
@@ -507,20 +507,22 @@ def _check_dummy_hash_model_forward(
507 assert torch.all(hashes == engram_ops.DEAD_ID)
508 output = engram.embed(hashes[:, 0])
509 # Trailing None is previous_aux; this stub captures no aux states.
510 - return output, None, None, None, None, None
511 -
512 - model = SimpleNamespace(
513 - use_mega_moe=False,
514 - use_sequence_parallel=False,
515 - fuse_mhc_all_reduce=False,
516 - engram_hash=state,
517 - engram_swa_prefix="swa",
518 - engram_dp_shared_memory=dp_shared_memory,
519 - layers=[Decoder(engram=engram)],
520 - start_layer=0,
521 - end_layer=1,
522 - aux_hidden_state_layers=(),
523 - )
510 + return output, hidden, hidden, hidden, hidden, None
511 +
512 + model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model)
513 + torch.nn.Module.__init__(model)
514 + model.use_mega_moe = False
515 + model.use_sequence_parallel = False
516 + model.fuse_mhc_all_reduce = False
517 + model.engram_hash = state
518 + model.engram_swa_prefix = "swa"
519 + model.engram_dp_shared_memory = dp_shared_memory
520 + model.layers = [Decoder(engram=engram)]
521 + model.start_layer = 0
522 + model.end_layer = 1
523 + model.decoder_replay_start = model.end_layer
524 + model.decoder_replay_layers = None
525 + model.aux_hidden_state_layers = ()
526 metadata = (
527 None
528 if dp_rank == 1
@@ -550,8 +552,7 @@ def _check_dummy_hash_model_forward(
552 dtype=torch.int32,
553 ),
554 ):
553 - output = model_ops.DeepseekV4Model.forward(
554 - model,
555 + output = model(
556 torch.arange(tokens, device="cuda"),
557 torch.arange(tokens, device="cuda"),
558 intermediate_tensors=None,
tests/models/test_deepseek_v41_decoder_replay_layers.py new
+104
@@ -0,0 +1,104 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DecoderReplayLayers runs the replay layers on the forward's replay batch."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.forward_context import (
9 + ForwardContext,
10 + get_forward_context,
11 + override_forward_context,
12 +)
13 +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
14 +
15 +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
16 +
17 +DEVICE = torch.device("cuda")
18 +NUM_TOKENS = 401
19 +HC, HIDDEN = 2, 3
20 +# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
21 +WINDOW = 128
22 +REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
23 +
24 +
25 +def _context(metadata):
26 + return ForwardContext(
27 + no_compile_layers={},
28 + attn_metadata=metadata,
29 + slot_mapping={},
30 + is_padding=torch.zeros(NUM_TOKENS, dtype=torch.bool, device=DEVICE),
31 + )
32 +
33 +
34 +def _set_replay_batch(layers, rows, metadata):
35 + layers.rows = torch.tensor(rows, device=DEVICE)
36 + layers.forward_context = ForwardContext(
37 + no_compile_layers={}, attn_metadata=metadata, slot_mapping={}
38 + )
39 +
40 +
41 +def _states(num_tokens: int):
42 + """(hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual),
43 + shaped as the model's: hidden states [T, H], mixes [T, HC], residual [T, HC, H]."""
44 + hidden = torch.arange(num_tokens, device=DEVICE, dtype=torch.float32)
45 + hidden = hidden[:, None].expand(num_tokens, HIDDEN).contiguous()
46 + mix = hidden[:, :HC].contiguous()
47 + residual = hidden[:, None, :].expand(num_tokens, HC, HIDDEN).contiguous()
48 + return (
49 + hidden,
50 + torch.arange(num_tokens, device=DEVICE),
51 + None,
52 + mix + 1,
53 + mix + 2,
54 + mix + 3,
55 + residual,
56 + )
57 +
58 +
59 +def test_run_gathers_states_and_realigns_shared_indexer_buffers():
60 + topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int()
61 + candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int()
62 + topk_before, candidates_before = topk.clone(), candidates.clone()
63 + seen = {}
64 +
65 + def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual):
66 + seen["hidden_states"] = hidden_states.clone()
67 + replay_context = get_forward_context()
68 + seen["attn_metadata"] = replay_context.attn_metadata
69 + seen["is_padding"] = replay_context.is_padding
70 + return residual, pre_mix
71 +
72 + layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates])
73 + states = _states(NUM_TOKENS)
74 + hidden, residual = states[0], states[-1]
75 + full, sub = object(), object()
76 + _set_replay_batch(layers, REPLAY_ROWS, sub)
77 + context = _context(full)
78 + with override_forward_context(context):
79 + outputs = layers(*states)
80 + assert get_forward_context() is context
81 +
82 + rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
83 + assert torch.equal(seen["hidden_states"], hidden[rows])
84 + assert seen["attn_metadata"] is sub and seen["is_padding"] is None
85 + assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows])
86 + assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows])
87 + assert len(outputs) == 2 and outputs[0].shape == residual.shape
88 + assert torch.equal(outputs[0][rows], residual[rows])
89 + assert outputs[0][1:173].abs().sum() == 0
90 +
91 +
92 +def test_no_replay_batch_runs_the_whole_batch():
93 + seen = {}
94 +
95 + def run_layers(hidden_states, *rest):
96 + seen["attn_metadata"] = get_forward_context().attn_metadata
97 + return (hidden_states,)
98 +
99 + layers = DecoderReplayLayers(WINDOW, run_layers, [])
100 + states = _states(NUM_TOKENS)
101 + full = object()
102 + with override_forward_context(_context(full)):
103 + (output,) = layers(*states)
104 + assert output is states[0] and seen["attn_metadata"] is full
tests/models/test_deepseek_v41_replay_batch.py new
+259
@@ -0,0 +1,259 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""The decoder-side SWA bounded replay batch DeepseekV41ModelState prepares."""
4 +
5 +from dataclasses import replace
6 +from types import SimpleNamespace
7 +from unittest.mock import MagicMock
8 +
9 +import numpy as np
10 +import pytest
11 +import torch
12 +
13 +from vllm.config import CUDAGraphMode
14 +from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
15 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
16 +from vllm.v1.worker.gpu.model_states.default import DefaultModelState
17 +
18 +pytestmark = pytest.mark.skipif(
19 + not torch.cuda.is_available(), reason="the replay buffers live on the GPU"
20 +)
21 +
22 +WINDOW = 128
23 +DEVICE = torch.device("cuda")
24 +GROUPS = SimpleNamespace(
25 + kv_cache_groups=[
26 + SimpleNamespace(
27 + layer_names=["swa"],
28 + kv_cache_spec=SimpleNamespace(
29 + prefix_cacheable=False, prefix_replay_tokens=WINDOW
30 + ),
31 + ),
32 + SimpleNamespace(
33 + layer_names=["mla"],
34 + kv_cache_spec=SimpleNamespace(
35 + prefix_cacheable=True, prefix_replay_tokens=0
36 + ),
37 + ),
38 + ]
39 +)
40 +BLOCK_TABLES = (torch.zeros(4, 4, device=DEVICE),) * 2
41 +
42 +# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
43 +QUERY_LENS = [1, 300, 100]
44 +SEQ_LENS = [500, 300, 300]
45 +PREFILLING = [False, True, True]
46 +REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
47 +
48 +
49 +@pytest.fixture
50 +def state(monkeypatch):
51 + """A DeepseekV41ModelState whose attention builds are recorded, not run."""
52 + cfg = MagicMock()
53 + cfg.model_config.enable_prompt_embeds = False
54 + cfg.model_config.uses_mrope = False
55 + cfg.model_config.is_multimodal_model = False
56 + cfg.scheduler_config.max_num_seqs = 16
57 + cfg.scheduler_config.max_num_batched_tokens = 1024
58 + cfg.parallel_config.data_parallel_size = 1
59 + cfg.compilation_config.fast_moe_cold_start = False
60 + layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None)
61 + model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers)
62 + builds: list = []
63 +
64 + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
65 + build = SimpleNamespace(
66 + batch=input_batch,
67 + cg_mode=cg_mode,
68 + slot_mappings=slot_mappings.clone(),
69 + replay_start=kw["model_specific_attn_metadata"].replay_start,
70 + attn_metadata=object(),
71 + )
72 + builds.append(build)
73 + return {"swa": build.attn_metadata}
74 +
75 + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
76 + state = DeepseekV41ModelState(cfg, model, None, DEVICE)
77 + state.builds = builds # type: ignore[attr-defined]
78 + return state
79 +
80 +
81 +def _input_batch(
82 + query_lens: list[int],
83 + seq_lens: list[int],
84 + is_prefilling: list[bool],
85 + device_query_lens: list[int] | None = None,
86 + num_tokens_after_padding: int | None = None,
87 + max_query_len: int | None = None,
88 +) -> InputBatch:
89 + num_reqs, num_tokens = len(query_lens), sum(query_lens)
90 + num_padded = num_tokens_after_padding or num_tokens
91 + batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(16, 1024, DEVICE))
92 + return replace(
93 + batch,
94 + num_tokens=num_tokens,
95 + num_tokens_after_padding=num_padded,
96 + query_start_loc=torch.tensor(
97 + [0, *np.cumsum(device_query_lens or query_lens)],
98 + dtype=torch.int32,
99 + device=DEVICE,
100 + ),
101 + query_start_loc_np=np.array([0, *np.cumsum(query_lens)], dtype=np.int32),
102 + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE),
103 + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
104 + is_prefilling_np=np.array(is_prefilling, dtype=np.bool_),
105 + positions=torch.cat(
106 + [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)]
107 + + [torch.zeros(num_padded - num_tokens, dtype=torch.int64)]
108 + ).to(DEVICE),
109 + is_padding=torch.arange(num_padded, device=DEVICE) >= num_tokens,
110 + max_query_len=max_query_len,
111 + )
112 +
113 +
114 +def _slot_mappings(num_tokens: int) -> torch.Tensor:
115 + return torch.stack([torch.arange(num_tokens), torch.arange(num_tokens) * 10]).to(
116 + DEVICE
117 + )
118 +
119 +
120 +def _prepare(state, batch, cg_mode):
121 + """Runs prepare_attn; returns the replay batch and the replay build's inputs."""
122 + state.prepare_attn(
123 + batch,
124 + cg_mode,
125 + BLOCK_TABLES,
126 + _slot_mappings(batch.num_tokens_after_padding),
127 + [],
128 + GROUPS,
129 + )
130 + layers = state.decoder_replay_layers
131 + if layers.rows is None:
132 + assert layers.forward_context is None
133 + return None, None
134 + return SimpleNamespace(
135 + rows=layers.rows, forward_context=layers.forward_context
136 + ), state.builds[-1]
137 +
138 +
139 +def test_replay_batch_keeps_each_request_window(state):
140 + state._replay_start_np[2] = 50 # the encoder-side replay start of request 2
141 + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
142 + replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
143 + assert replay is not None
144 + rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
145 + assert torch.equal(replay.rows, rows)
146 + sub = build.batch
147 + assert build.cg_mode == CUDAGraphMode.NONE
148 + assert sub.query_start_loc.tolist() == [0, 1, 129, 229]
149 + assert sub.query_start_loc_np.tolist() == [0, 1, 129, 229]
150 + assert sub.num_tokens == sub.num_tokens_after_padding == 229
151 + assert sub.max_query_len == WINDOW
152 + # The trimmed request's window starts at its replay window; a higher
153 + # encoder-side replay start stands.
154 + assert build.replay_start.tolist() == [0, 300 - WINDOW, 50]
155 + assert torch.equal(sub.positions, batch.positions[rows])
156 + assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows])
157 + context = replay.forward_context
158 + assert context.attn_metadata == {"swa": build.attn_metadata}
159 + assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows])
160 + assert context.dp_metadata is None and context.is_padding is None
161 +
162 +
163 +def test_replay_batch_keeps_device_decode_boundaries(state):
164 + """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device
165 + [1, 3]): their rows stay and the boundaries shift only by trimmed rows."""
166 + batch = _input_batch(
167 + [2, 2, 300], [10, 10, 300], PREFILLING, device_query_lens=[1, 3, 300]
168 + )
169 + replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
170 + assert replay is not None
171 + assert replay.rows.tolist() == [0, 1, 2, 3, *range(304 - WINDOW, 304)]
172 + assert build.batch.query_start_loc.tolist() == [0, 1, 4, 4 + WINDOW]
173 +
174 +
175 +def test_prompt_logprobs_requests_keep_their_rows(state):
176 + """Every prompt row of a prompt-logprobs request is read, so it never
177 + trims; the other prefills still do."""
178 + state._keeps_rows_np[1] = True
179 + batch = _input_batch([1, 300, 300], [500, 300, 300], PREFILLING)
180 + replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
181 + assert replay is not None
182 + assert replay.rows.tolist() == [0, *range(1, 301), *range(601 - WINDOW, 601)]
183 + assert build.batch.max_query_len == 300
184 +
185 +
186 +def test_replay_batch_keeps_adaptive_verification_query_bound(state):
187 + """Adaptive verification bounds the decodes' device-side query lengths
188 + above their CPU lengths; the replay batch keeps that bound."""
189 + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, max_query_len=200)
190 + replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
191 + assert replay is not None and build.batch.max_query_len == 200
192 +
193 +
194 +def test_only_eager_steps_trim(state):
195 + """A CUDA graph keeps the layers on its whole batch; an eager step trims a
196 + padded batch to its real rows."""
197 + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512)
198 + for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
199 + assert _prepare(state, batch, cg_mode) == (None, None)
200 + replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
201 + assert replay is not None
202 + assert replay.rows.tolist() == REPLAY_ROWS
203 + assert build.cg_mode == CUDAGraphMode.NONE
204 + assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS)
205 +
206 +
207 +def test_whole_batch_forwards_get_no_replay_batch(state):
208 + short = _input_batch([100, 100], [100, 100], [True, True])
209 + for cg_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
210 + assert _prepare(state, short, cg_mode) == (None, None)
211 + assert len(state.builds) == 3 # the batch's own metadata only
212 + # Dummy batches, captures included (prepared as eager), are not prefills and
213 + # keep their rows.
214 + dummy = _input_batch(
215 + QUERY_LENS, SEQ_LENS, [False] * 3, num_tokens_after_padding=512
216 + )
217 + assert _prepare(state, dummy, CUDAGraphMode.NONE) == (None, None)
218 +
219 +
220 +@pytest.fixture
221 +def dp_state(state, monkeypatch):
222 + """`state` on DP rank 0 of 2; `state.other` sets what rank 1 reports
223 + (trims, replay tokens)."""
224 + import vllm.models.deepseek_v41.nvidia.model_state as module
225 +
226 + state.vllm_config.parallel_config.data_parallel_size = 2
227 + state.vllm_config.parallel_config.data_parallel_rank = 0
228 + state.vllm_config.parallel_config.is_moe_model = True
229 + state.other = (False, 0) # type: ignore[attr-defined]
230 +
231 + def all_reduce(tensor, group):
232 + tensor[:, 1] = torch.tensor(state.other, dtype=torch.int32)
233 +
234 + monkeypatch.setattr(module.dist, "all_reduce", all_reduce)
235 + monkeypatch.setattr(module, "get_dp_group", lambda: SimpleNamespace(cpu_group=None))
236 + return state
237 +
238 +
239 +def test_dp_ranks_replay_together(dp_state):
240 + """A trimming rank makes the others replay too, on every rank's count."""
241 + short = _input_batch([100, 100], [100, 100], [True, True])
242 + dp_state.other = (False, 200)
243 + assert _prepare(dp_state, short, CUDAGraphMode.NONE) == (None, None)
244 +
245 + dp_state.other = (True, 300) # rank 1 trims to 300 rows
246 + replay, build = _prepare(dp_state, short, CUDAGraphMode.NONE)
247 + assert replay is not None
248 + assert replay.rows.tolist() == list(range(200)) # this rank keeps every row
249 + dp_metadata = replay.forward_context.dp_metadata
250 + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [200, 300]
251 +
252 +
253 +def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
254 + dummy = _input_batch([512], [512], [False])
255 + dp_state.other = (True, 129)
256 + replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE)
257 + assert replay is not None and replay.rows.shape[0] == WINDOW
258 + dp_metadata = replay.forward_context.dp_metadata
259 + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129]
vllm/config/cache.py
+3 -1
@@ -242,7 +242,9 @@ class CacheConfig:
242 swa_bounded_replay: bool = True
243 """Keep the sliding-window KV of models that support it (DeepSeek-V4.1)
244 out of prefix caching and rebuild it after a prefix hit by recomputing the
245 - hit's last window. Requires model runner V2."""
245 + hit's last window. The layers past the last KV-source layer then also run
246 + eager prefill steps on each request's trailing window only. Requires model
247 + runner V2."""
248
249 kv_cache_memory_bytes: int | None = None
250 """Size of KV Cache per GPU in bytes. By default, this is set to None
vllm/models/deepseek_v41/decoder_replay_layers.py new
+67
@@ -0,0 +1,67 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Decoder-side SWA bounded replay: running the replay layers on their batch.
4 +
5 +The layers past the last KV-source layer own nothing but sliding-window KV, so
6 +in eager prefill steps they run on each request's last ``window`` rows only.
7 +``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention
8 +metadata and a forward context of its own, like a microbatch;
9 +``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers
10 +under that context and scatters the outputs back to batch rows. Steps that run
11 +in a CUDA graph keep the layers on the whole batch, inside the graph.
12 +"""
13 +
14 +from collections.abc import Callable
15 +
16 +import torch
17 +
18 +from vllm.forward_context import ForwardContext, override_forward_context
19 +from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput
20 +
21 +
22 +class DecoderReplayLayers:
23 + """Runs the replay layers on the step's replay batch.
24 +
25 + ``run_layers`` takes a batch's layer inputs and returns its per-row
26 + outputs. ``row_buffers`` hold per-row results the source layer's indexer
27 + publishes for the layers after it; they are compacted to the replay rows
28 + in place.
29 + """
30 +
31 + def __init__(
32 + self,
33 + window: int,
34 + run_layers: Callable[..., tuple[torch.Tensor, ...]],
35 + row_buffers: list[torch.Tensor],
36 + ) -> None:
37 + self.window = window
38 + self.run_layers = run_layers
39 + self.row_buffers = row_buffers
40 + # The replay batch, set by the model state every step: its rows of the
41 + # batch and its forward context. None runs the layers on the batch.
42 + self.rows: torch.Tensor | None = None
43 + self.forward_context: ForwardContext | None = None
44 +
45 + def __call__(
46 + self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
47 + ) -> tuple[torch.Tensor, ...]:
48 + rows = self.rows
49 + if rows is None:
50 + return self.run_layers(hidden_states, *states)
51 + # A trimming step holds a prefill longer than the window, more tokens
52 + # than any step whose MoE leaves its finalize to the next layer.
53 + assert isinstance(hidden_states, torch.Tensor)
54 + num_rows = rows.shape[0]
55 + for buf in self.row_buffers:
56 + buf[:num_rows].copy_(buf.index_select(0, rows))
57 + with override_forward_context(self.forward_context):
58 + row_outputs = self.run_layers(
59 + hidden_states.index_select(0, rows),
60 + *(None if t is None else t.index_select(0, rows) for t in states),
61 + )
62 + # The trimmed rows' outputs stay zero; nothing reads them.
63 + num_tokens = hidden_states.shape[0]
64 + return tuple(
65 + out.new_zeros((num_tokens, *out.shape[1:])).index_copy_(0, rows, out)
66 + for out in row_outputs
67 + )
vllm/models/deepseek_v41/nvidia/model.py
+222 -40
@@ -83,6 +83,7 @@ from vllm.models.deepseek_v4.nvidia.model import (
83 prepare_mega_gate_routing_metadata,
84 )
85 from vllm.models.deepseek_v41.attention import DeepseekV4Attention
86 +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
87 from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
88 DeepseekV4MegaAttnAttention,
89 )
@@ -741,6 +742,35 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
742 if successor.engram is None:
743 layer.ffn.defer_finalize()
744
745 + # Decoder-side SWA bounded replay: in eager prefill steps the layers past
746 + # the last KV source run on each request's trailing window only
747 + # (decoder_replay_layers.py).
748 + self.decoder_replay_layers: DecoderReplayLayers | None = None
749 + self.decoder_replay_start = self.end_layer
750 + cut = max(config.kv_source_layer_ids)
751 + if (
752 + cut < self.end_layer - 1
753 + and self._decoder_replay_supported(vllm_config, cut)
754 + and self.layers[cut].attn.swa_cache_layer.bounded_replay
755 + ):
756 + self.decoder_replay_start = cut + 1
757 + self.decoder_replay_layers = DecoderReplayLayers(
758 + config.sliding_window,
759 + self._run_replay_layers,
760 + [
761 + buf
762 + for buf in (self.topk_indices_buffer, self.candidate_block_buffer)
763 + if buf is not None
764 + ],
765 + )
766 + logger.info_once(
767 + "Decoder SWA bounded replay: in eager prefill steps, layers "
768 + "%d-%d run on each request's last %d tokens only.",
769 + cut + 1,
770 + self.end_layer - 1,
771 + config.sliding_window,
772 + )
773 +
774 # The n-gram hash needs a slot-keyed rolling store of compressed ids
775 # (chunked prefill / decode lookback); key it off the first local
776 # layer's sliding-window KV cache. Only PP ranks owning an engram
@@ -904,29 +934,28 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
934 hidden_states = sp_shard(hidden_states)
935 input_ids = sp_shard(input_ids)
936
907 - mega_gate_metadata = None
908 - if self.use_mega_moe:
909 - mega_gate_metadata = prepare_mega_gate_routing_metadata(
910 - input_ids,
911 - has_hash_routing=False,
912 - image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
913 - if getattr(self.config, "vision_n_layers", 0) > 0
914 - else None,
915 - )
916 -
937 residual, post_mix, res_mix = None, None, None
938 pre_mix: torch.Tensor | None = None
939 if not get_pp_group().is_first_rank:
940 assert intermediate_tensors is not None
941 pre_mix = intermediate_tensors["pre_mix"]
922 - # Every layer's post runs inside the next layer's fused pre, so aux
923 - # hidden states are read back from there instead of recomputed.
942 aux_hidden_by_layer: dict[int, torch.Tensor] = {}
925 - for idx, layer in enumerate(
926 - islice(self.layers, self.start_layer, self.end_layer),
927 - start=self.start_layer,
928 - ):
929 - hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer(
943 + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers(
944 + range(self.start_layer, self.decoder_replay_start),
945 + hidden_states,
946 + positions,
947 + input_ids,
948 + pre_mix,
949 + post_mix,
950 + res_mix,
951 + residual,
952 + aux_hidden_by_layer,
953 + engram_hashes,
954 + engram_mask,
955 + )
956 + late_aux: list[torch.Tensor] = []
957 + if self.decoder_replay_layers is not None:
958 + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers(
959 hidden_states,
960 positions,
961 input_ids,
@@ -934,36 +963,22 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
963 post_mix,
964 res_mix,
965 residual,
937 - engram_hashes,
938 - engram_mask,
939 - capture_previous_aux=idx in self.aux_hidden_state_layers,
940 - mega_gate_metadata=mega_gate_metadata,
966 )
942 - if previous_aux is not None:
943 - # idx is the one-based id of the layer whose post this is.
944 - if self.use_sequence_parallel:
945 - previous_aux = sp_all_gather(previous_aux)[:full_num_tokens]
946 - aux_hidden_by_layer[idx] = previous_aux
947 - # Without a successor boundary, the last layer finalized its own MoE.
948 - assert isinstance(hidden_states, torch.Tensor)
949 - if layer is not None:
950 - # The last layer has no successor to fold its post into.
951 - if self.fuse_mhc_all_reduce:
952 - hidden_states = tensor_model_parallel_all_reduce(hidden_states)
953 - hidden_states = mhc_post_tilelang(
954 - hidden_states, residual, post_mix, res_mix
967 + else:
968 + hidden_states = self._collapse(
969 + hidden_states,
970 + residual,
971 + post_mix,
972 + res_mix,
973 + aux_hidden_by_layer,
974 + full_num_tokens,
975 )
956 - if self.end_layer in self.aux_hidden_state_layers:
957 - final_aux = hidden_states.mean(dim=1)
958 - if self.use_sequence_parallel:
959 - final_aux = sp_all_gather(final_aux)[:full_num_tokens]
960 - aux_hidden_by_layer[self.end_layer] = final_aux
976
977 aux_hidden_states = [
978 aux_hidden_by_layer[layer_id]
979 for layer_id in self.aux_hidden_state_layers
980 if layer_id in aux_hidden_by_layer
966 - ]
981 + ] + late_aux
982
983 if not get_pp_group().is_last_rank:
984 return IntermediateTensors(
@@ -992,6 +1007,169 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1007 return hidden_states, aux_hidden_states
1008 return hidden_states
1009
1010 + def _mega_gate_metadata(
1011 + self, input_ids: torch.Tensor | None
1012 + ) -> MegaGateRoutingMetadata | None:
1013 + if not self.use_mega_moe:
1014 + return None
1015 + assert input_ids is not None
1016 + return prepare_mega_gate_routing_metadata(
1017 + input_ids,
1018 + has_hash_routing=False,
1019 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
1020 + if getattr(self.config, "vision_n_layers", 0) > 0
1021 + else None,
1022 + )
1023 +
1024 + def _run_layers(
1025 + self,
1026 + layer_ids: range,
1027 + hidden_states: torch.Tensor | MoEOutput,
1028 + positions: torch.Tensor,
1029 + input_ids: torch.Tensor | None,
1030 + pre_mix: torch.Tensor | None,
1031 + post_mix: torch.Tensor | None,
1032 + res_mix: torch.Tensor | None,
1033 + residual: torch.Tensor | None,
1034 + aux_hidden_by_layer: dict[int, torch.Tensor],
1035 + engram_hashes: torch.Tensor | None = None,
1036 + engram_mask: torch.Tensor | None = None,
1037 + ) -> tuple[
1038 + torch.Tensor | MoEOutput, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
1039 + ]:
1040 + # Every layer's post runs inside the next layer's fused pre, so aux
1041 + # hidden states are read back from there instead of recomputed.
1042 + full_num_tokens = positions.shape[0]
1043 + mega_gate_metadata = self._mega_gate_metadata(input_ids)
1044 + for idx in layer_ids:
1045 + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = (
1046 + self.layers[idx](
1047 + hidden_states,
1048 + positions,
1049 + input_ids,
1050 + pre_mix,
1051 + post_mix,
1052 + res_mix,
1053 + residual,
1054 + engram_hashes,
1055 + engram_mask,
1056 + capture_previous_aux=idx in self.aux_hidden_state_layers,
1057 + mega_gate_metadata=mega_gate_metadata,
1058 + )
1059 + )
1060 + if previous_aux is not None:
1061 + # idx is the one-based id of the layer whose post this is.
1062 + if self.use_sequence_parallel:
1063 + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens]
1064 + aux_hidden_by_layer[idx] = previous_aux
1065 + assert residual is not None and post_mix is not None
1066 + assert res_mix is not None and pre_mix is not None
1067 + return hidden_states, residual, post_mix, res_mix, pre_mix
1068 +
1069 + def _collapse(
1070 + self,
1071 + hidden_states: torch.Tensor | MoEOutput,
1072 + residual: torch.Tensor,
1073 + post_mix: torch.Tensor,
1074 + res_mix: torch.Tensor,
1075 + aux_hidden_by_layer: dict[int, torch.Tensor],
1076 + full_num_tokens: int,
1077 + ) -> torch.Tensor:
1078 + # Without a successor boundary, the last layer finalized its own MoE.
1079 + assert isinstance(hidden_states, torch.Tensor)
1080 + # The last layer has no successor to fold its post into.
1081 + if self.fuse_mhc_all_reduce:
1082 + hidden_states = tensor_model_parallel_all_reduce(hidden_states)
1083 + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
1084 + if self.end_layer in self.aux_hidden_state_layers:
1085 + final_aux = hidden_states.mean(dim=1)
1086 + if self.use_sequence_parallel:
1087 + final_aux = sp_all_gather(final_aux)[:full_num_tokens]
1088 + aux_hidden_by_layer[self.end_layer] = final_aux
1089 + return hidden_states
1090 +
1091 + def _run_replay_layers(
1092 + self,
1093 + hidden_states: torch.Tensor | MoEOutput,
1094 + positions: torch.Tensor,
1095 + input_ids: torch.Tensor | None,
1096 + pre_mix: torch.Tensor,
1097 + post_mix: torch.Tensor,
1098 + res_mix: torch.Tensor,
1099 + residual: torch.Tensor,
1100 + ) -> tuple[torch.Tensor, ...]:
1101 + """The layers past the last KV source, on whatever rows they are given;
1102 + returns their output, the last FFN's pre-mix and the aux hidden states
1103 + they capture."""
1104 + aux_hidden_by_layer: dict[int, torch.Tensor] = {}
1105 + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers(
1106 + range(self.decoder_replay_start, self.end_layer),
1107 + hidden_states,
1108 + positions,
1109 + input_ids,
1110 + pre_mix,
1111 + post_mix,
1112 + res_mix,
1113 + residual,
1114 + aux_hidden_by_layer,
1115 + )
1116 + hidden_states = self._collapse(
1117 + hidden_states,
1118 + residual,
1119 + post_mix,
1120 + res_mix,
1121 + aux_hidden_by_layer,
1122 + positions.shape[0],
1123 + )
1124 + return (
1125 + hidden_states,
1126 + pre_mix,
1127 + *(
1128 + aux_hidden_by_layer[layer_id]
1129 + for layer_id in self.aux_hidden_state_layers
1130 + if layer_id in aux_hidden_by_layer
1131 + ),
1132 + )
1133 +
1134 + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool:
1135 + """Whether this rank may trim the layers after ``cut``; warns when not."""
1136 + parallel_config = vllm_config.parallel_config
1137 + spec_config = vllm_config.speculative_config
1138 + draft_config = spec_config.draft_model_config if spec_config else None
1139 + draft_hf_config = getattr(draft_config, "hf_config", None)
1140 + draft_window = getattr(draft_hf_config, "sliding_window", None)
1141 + draft_layer_types = getattr(draft_hf_config, "layer_types", None) or ()
1142 + window = self.config.sliding_window
1143 + if self.start_layer > cut or self.end_layer < self.config.num_hidden_layers:
1144 + reason = (
1145 + "the pipeline stage holding the last KV source layer must also "
1146 + "hold every layer after it"
1147 + )
1148 + elif (
1149 + self.use_sequence_parallel
1150 + or parallel_config.prefill_context_parallel_size > 1
1151 + or parallel_config.use_ubatching
1152 + ):
1153 + reason = (
1154 + "the replay-layer batch shrinks per rank, which sequence and "
1155 + "prefill-context parallelism and microbatching cannot follow"
1156 + )
1157 + elif any(i > cut for i in getattr(self.config, "engram_layer_ids", ())):
1158 + reason = "an Engram layer sits after the last KV source layer"
1159 + elif draft_config is not None and (
1160 + draft_window is None
1161 + or draft_window > window
1162 + or any(t != "sliding_attention" for t in draft_layer_types)
1163 + ):
1164 + reason = (
1165 + f"the drafter (sliding window {draft_window}) reads hidden states "
1166 + f"outside the target's {window}-token window"
1167 + )
1168 + else:
1169 + return True
1170 + logger.warning_once("Decoder SWA bounded replay is off: %s.", reason)
1171 + return False
1172 +
1173 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1174 stacked_params_mapping = [
1175 # (param_name, shard_name, shard_id)
@@ -1397,6 +1575,10 @@ class DeepseekV41LLMForCausalLM(
1575 engram_hash = self.model.engram_hash
1576 return engram_hash.lookback_depth if engram_hash is not None else 0
1577
1578 + @property
1579 + def decoder_replay_layers(self) -> DecoderReplayLayers | None:
1580 + return self.model.decoder_replay_layers
1581 +
1582 def forward(
1583 self,
1584 input_ids: torch.Tensor,
vllm/models/deepseek_v41/nvidia/model_state.py
+321 -2
@@ -1,17 +1,24 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import replace
4 from typing import Any
5
6 import numpy as np
7 import torch
8 +import torch.distributed as dist
9 import torch.nn as nn
10
11 from vllm.config import CUDAGraphMode, VllmConfig
12 +from vllm.distributed.parallel_state import get_dp_group
13 +from vllm.forward_context import DPMetadata, create_forward_context
14 +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
15 from vllm.triton_utils import tl, triton
16 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
17 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
18 from vllm.v1.core.sched.output import NewRequestData
19 from vllm.v1.kv_cache_interface import KVCacheConfig
20 +from vllm.v1.worker.dp_utils import should_skip_dp_coordination
21 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
22 from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
23 from vllm.v1.worker.gpu.input_batch import InputBatch
24 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
@@ -81,6 +88,61 @@ def _pad_replayed_slots_kernel(
88 )
89
90
91 +@triton.jit
92 +def _gather_replay_batch_kernel(
93 + query_start_loc_ptr, # [num_reqs + 1]
94 + dropped_before_ptr, # [num_reqs + 1] rows trimmed before each boundary
95 + seq_lens_ptr, # [num_reqs]
96 + replay_start_ptr, # [num_reqs] the encoder-side replay start
97 + positions_ptr,
98 + slot_mappings_ptr, # [num_groups, num_tokens]
99 + slot_mappings_stride,
100 + rows_ptr, # out: [num_kept] batch row of every kept row
101 + kept_query_start_loc_ptr, # out: [num_reqs + 1]
102 + kept_positions_ptr, # out: [num_kept]
103 + kept_slot_mappings_ptr, # out: [num_groups, num_kept]
104 + kept_slot_mappings_stride,
105 + kept_kv_start_ptr, # out: [num_reqs]
106 + window,
107 + NUM_GROUPS: tl.constexpr,
108 + BLOCK: tl.constexpr,
109 +):
110 + """One program per request: its kept rows are the last ones before its
111 + boundary, and land at that boundary minus the rows trimmed before it."""
112 + req = tl.program_id(0)
113 + begin = tl.load(query_start_loc_ptr + req)
114 + end = tl.load(query_start_loc_ptr + req + 1)
115 + kept_begin = begin - tl.load(dropped_before_ptr + req)
116 + kept_end = end - tl.load(dropped_before_ptr + req + 1)
117 + if req == 0:
118 + tl.store(kept_query_start_loc_ptr, kept_begin)
119 + tl.store(kept_query_start_loc_ptr + req + 1, kept_end)
120 +
121 + # A trimmed request holds no replay-layer window KV below its kept rows.
122 + kv_start = tl.load(replay_start_ptr + req)
123 + if kept_end - kept_begin < end - begin:
124 + kv_start = tl.maximum(kv_start, tl.load(seq_lens_ptr + req) - window)
125 + tl.store(kept_kv_start_ptr + req, kv_start)
126 +
127 + first_row = end - (kept_end - kept_begin)
128 + for tok in range(kept_begin, kept_end, BLOCK):
129 + offs = tok + tl.arange(0, BLOCK)
130 + mask = offs < kept_end
131 + row = first_row + (offs - kept_begin)
132 + tl.store(rows_ptr + offs, row.to(tl.int64), mask=mask)
133 + tl.store(
134 + kept_positions_ptr + offs,
135 + tl.load(positions_ptr + row, mask=mask),
136 + mask=mask,
137 + )
138 + for g in tl.static_range(NUM_GROUPS):
139 + tl.store(
140 + kept_slot_mappings_ptr + g * kept_slot_mappings_stride + offs,
141 + tl.load(slot_mappings_ptr + g * slot_mappings_stride + row, mask=mask),
142 + mask=mask,
143 + )
144 +
145 +
146 class ReplayAttnMetadata(ModelSpecificAttnMetadata):
147 """Hands the batch's replay starts to the sliding-window builders."""
148
@@ -111,7 +173,10 @@ class DeepseekV41ModelState(DefaultModelState):
173 ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots
174 in the prefix-cacheable groups so the cached KV stays as is, and hands the
175 starts to the sliding-window metadata builders, whose kernels read no
114 - window KV below them.
176 + window KV below them. With the decoder side on
177 + (``model.decoder_replay_layers``) it also prepares the replay layers'
178 + batch in eager steps: each prefill's last ``window`` rows as a sub-batch
179 + with metadata and a forward context of its own, like a microbatch.
180 """
181
182 def __init__(
@@ -140,9 +205,40 @@ class DeepseekV41ModelState(DefaultModelState):
205 # the KV cache config on first use; None until then.
206 self._replay: tuple[int, torch.Tensor] | None = None
207
208 + self.decoder_replay_layers: DecoderReplayLayers | None = getattr(
209 + model, "decoder_replay_layers", None
210 + )
211 + if self.decoder_replay_layers is not None:
212 + # Requests whose every prompt row is read (prompt logprobs) never trim.
213 + self._keeps_rows_np = np.zeros(self.max_num_reqs, dtype=np.bool_)
214 + self._dropped_before = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
215 + # The replay batch's inputs: its rows of the batch, query offsets,
216 + # positions, slot mappings (sized with the KV cache groups on first
217 + # use) and the position from which each request holds replay-layer
218 + # window KV.
219 + self._kept_rows = torch.zeros(
220 + self.max_num_tokens, dtype=torch.int64, device=device
221 + )
222 + self._kept_query_start_loc = torch.zeros(
223 + self.max_num_reqs + 1, dtype=torch.int32, device=device
224 + )
225 + self._kept_positions = torch.zeros(
226 + self.max_num_tokens, dtype=torch.int64, device=device
227 + )
228 + self._kept_slot_mappings: torch.Tensor | None = None
229 + self._kept_kv_start = torch.zeros(
230 + self.max_num_reqs, dtype=torch.int32, device=device
231 + )
232 + self._replay_attn_groups: list[list[AttentionGroup]] | None = None
233 +
234 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
235 super().add_request(req_index, new_req_data)
236 self._replay_start_np[req_index] = new_req_data.replay_start
237 + if self.decoder_replay_layers is not None:
238 + params = new_req_data.sampling_params
239 + self._keeps_rows_np[req_index] = (
240 + params is not None and params.prompt_logprobs is not None
241 + )
242
243 def prepare_inputs(
244 self, input_batch: InputBatch, req_states: RequestState
@@ -196,7 +292,16 @@ class DeepseekV41ModelState(DefaultModelState):
292 device=self.device,
293 ),
294 )
295 + if self.decoder_replay_layers is not None:
296 + self._kept_slot_mappings = torch.zeros(
297 + len(specs),
298 + self.max_num_tokens,
299 + dtype=torch.int64,
300 + device=self.device,
301 + )
302 + self._warm_up_replay_kernels(input_batch, slot_mappings)
303 window, cacheable_groups = self._replay
304 + replay_start: torch.Tensor | None = None
305 if window:
306 num_reqs = input_batch.num_reqs
307 # Decode rows sit above the hit, so only prefills carry a replay
@@ -226,7 +331,7 @@ class DeepseekV41ModelState(DefaultModelState):
331 )
332 assert model_specific_attn_metadata is None
333 model_specific_attn_metadata = ReplayAttnMetadata(replay_start)
229 - return super().prepare_attn(
334 + attn_metadata = super().prepare_attn(
335 input_batch,
336 cudagraph_mode,
337 block_tables,
@@ -237,3 +342,217 @@ class DeepseekV41ModelState(DefaultModelState):
342 ubatch_idx=ubatch_idx,
343 model_specific_attn_metadata=model_specific_attn_metadata,
344 )
345 + if self.decoder_replay_layers is not None:
346 + assert replay_start is not None
347 + self._prepare_replay_batch(
348 + input_batch,
349 + cudagraph_mode,
350 + block_tables,
351 + slot_mappings,
352 + attn_groups,
353 + kv_cache_config,
354 + replay_start,
355 + )
356 + return attn_metadata
357 +
358 + def _warm_up_replay_kernels(
359 + self, input_batch: InputBatch, slot_mappings: torch.Tensor
360 + ) -> None:
361 + """Compile the replay kernels on the buffers they run on: the startup
362 + dummy batches never replay or trim, so they would not launch them."""
363 + assert self._replay is not None
364 + window, cacheable_groups = self._replay
365 + if not window:
366 + return
367 + _pad_replayed_slots_kernel.warmup(
368 + slot_mappings,
369 + slot_mappings.stride(0),
370 + cacheable_groups,
371 + input_batch.query_start_loc,
372 + input_batch.positions,
373 + self._replay_start,
374 + window,
375 + PAD_SLOT_ID,
376 + NUM_GROUPS=cacheable_groups.numel(),
377 + BLOCK=1024,
378 + grid=(1,),
379 + )
380 + if self.decoder_replay_layers is None:
381 + return
382 + assert self._kept_slot_mappings is not None
383 + _gather_replay_batch_kernel.warmup(
384 + input_batch.query_start_loc,
385 + self._dropped_before.copy_to_uva(np.zeros(1, dtype=np.int32)),
386 + input_batch.seq_lens,
387 + self._replay_start,
388 + input_batch.positions,
389 + slot_mappings,
390 + slot_mappings.stride(0),
391 + self._kept_rows,
392 + self._kept_query_start_loc,
393 + self._kept_positions,
394 + self._kept_slot_mappings,
395 + self._kept_slot_mappings.stride(0),
396 + self._kept_kv_start,
397 + self.decoder_replay_layers.window,
398 + NUM_GROUPS=slot_mappings.shape[0],
399 + BLOCK=1024,
400 + grid=(1,),
401 + )
402 +
403 + def _prepare_replay_batch(
404 + self,
405 + input_batch: InputBatch,
406 + cudagraph_mode: CUDAGraphMode,
407 + block_tables: tuple[torch.Tensor, ...],
408 + slot_mappings: torch.Tensor,
409 + attn_groups: list[list[AttentionGroup]],
410 + kv_cache_config: KVCacheConfig,
411 + replay_start: torch.Tensor,
412 + ) -> None:
413 + """Set the replay layers' batch for this forward, or none when nothing
414 + trims. Only eager steps trim: a CUDA graph keeps the layers on the whole
415 + batch, and ranks share the graph mode. Under data parallelism every rank
416 + replays if any does."""
417 + layers = self.decoder_replay_layers
418 + assert layers is not None
419 + layers.rows = layers.forward_context = None
420 + if cudagraph_mode != CUDAGraphMode.NONE:
421 + return
422 +
423 + num_reqs = input_batch.num_reqs
424 + query_lens = np.diff(input_batch.query_start_loc_np[: num_reqs + 1])
425 + kept_lens = np.where(
426 + self._keeps_rows_np[input_batch.idx_mapping_np[:num_reqs]],
427 + query_lens,
428 + np.minimum(query_lens, layers.window),
429 + )
430 + # Dummy batches (captures, an idle DP rank's) are not prefills: they
431 + # keep their rows unless a rank trims.
432 + trims = bool(
433 + (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any()
434 + )
435 + num_tokens = int(kept_lens.sum())
436 + dp_metadata = None
437 + if self.vllm_config.parallel_config.data_parallel_size > 1:
438 + trims, dp_metadata = self._agree_across_dp(trims, num_tokens)
439 + if not trims:
440 + return
441 +
442 + kept_batch, kept_slot_mappings = self._kept_input_batch(
443 + input_batch, slot_mappings, replay_start, kept_lens
444 + )
445 + attn_metadata = super().prepare_attn(
446 + kept_batch,
447 + CUDAGraphMode.NONE,
448 + block_tables,
449 + kept_slot_mappings,
450 + self._replay_groups(attn_groups),
451 + kv_cache_config,
452 + model_specific_attn_metadata=ReplayAttnMetadata(
453 + self._kept_kv_start[:num_reqs]
454 + ),
455 + )
456 + layers.rows = self._kept_rows[:num_tokens]
457 + layers.forward_context = create_forward_context(
458 + attn_metadata,
459 + self.vllm_config,
460 + dp_metadata=dp_metadata,
461 + slot_mapping=build_slot_mappings_by_layer(
462 + kept_slot_mappings, kv_cache_config
463 + ),
464 + )
465 +
466 + def _kept_input_batch(
467 + self,
468 + input_batch: InputBatch,
469 + slot_mappings: torch.Tensor,
470 + replay_start: torch.Tensor,
471 + kept_lens: np.ndarray,
472 + ) -> tuple[InputBatch, torch.Tensor]:
473 + """Build the sub-``InputBatch`` of each request's last ``kept_lens``
474 + rows, and its slot mappings.
475 +
476 + The kept rows follow the device boundaries minus the rows trimmed before
477 + them: adaptive verification resizes the decodes on the GPU alone, and
478 + decodes never trim. Like a microbatch's, the sub-batch describes the
479 + forward only; its sampling fields are carried over and must not be read.
480 + """
481 + num_reqs = input_batch.num_reqs
482 + num_tokens = int(kept_lens.sum())
483 + kept_query_start_loc_np = np.zeros(num_reqs + 1, dtype=np.int32)
484 + np.cumsum(kept_lens, out=kept_query_start_loc_np[1:])
485 + dropped_before = self._dropped_before.copy_to_uva(
486 + input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np
487 + )
488 + assert self._kept_slot_mappings is not None
489 + kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens]
490 + _gather_replay_batch_kernel[(num_reqs,)](
491 + input_batch.query_start_loc,
492 + dropped_before,
493 + input_batch.seq_lens,
494 + replay_start,
495 + input_batch.positions,
496 + slot_mappings,
497 + slot_mappings.stride(0),
498 + self._kept_rows,
499 + self._kept_query_start_loc,
500 + self._kept_positions,
501 + kept_slot_mappings,
502 + kept_slot_mappings.stride(0),
503 + self._kept_kv_start,
504 + self.decoder_replay_layers.window, # type: ignore[union-attr]
505 + NUM_GROUPS=slot_mappings.shape[0],
506 + BLOCK=1024,
507 + )
508 + # Adaptive verification may lengthen the decodes on the device, up to
509 + # the batch's bound.
510 + max_query_len = max(int(kept_lens.max()), input_batch.max_query_len or 0)
511 + kept_batch = replace(
512 + input_batch,
513 + num_tokens=num_tokens,
514 + num_tokens_after_padding=num_tokens,
515 + query_start_loc=self._kept_query_start_loc[: num_reqs + 1],
516 + query_start_loc_np=kept_query_start_loc_np,
517 + max_query_len=max_query_len,
518 + positions=self._kept_positions[:num_tokens],
519 + fast_prefill=None,
520 + )
521 + return kept_batch, kept_slot_mappings
522 +
523 + def _agree_across_dp(
524 + self, trims: bool, num_tokens: int
525 + ) -> tuple[bool, DPMetadata | None]:
526 + """Whether any rank trims and, then, the replay layers' DP metadata from
527 + every rank's replay token count."""
528 + parallel_config = self.vllm_config.parallel_config
529 + agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32)
530 + agreed[:, parallel_config.data_parallel_rank] = torch.tensor(
531 + [trims, num_tokens]
532 + )
533 + if not should_skip_dp_coordination():
534 + dist.all_reduce(agreed, group=get_dp_group().cpu_group)
535 + if not agreed[0].any():
536 + return False, None
537 + return True, DPMetadata.make(parallel_config, num_tokens, agreed[1])
538 +
539 + def _replay_groups(
540 + self, attn_groups: list[list[AttentionGroup]]
541 + ) -> list[list[AttentionGroup]]:
542 + """The attention groups with metadata builders of their own, like each
543 + microbatch's: a builder keeps the metadata it built, and the runner's
544 + hold the batch's."""
545 + if self._replay_attn_groups is None:
546 + self._replay_attn_groups = []
547 + for groups in attn_groups:
548 + replay_groups = []
549 + for group in groups:
550 + replay_group = replace(group, metadata_builders=[])
551 + replay_group.create_metadata_builders(
552 + self.vllm_config,
553 + self.device,
554 + group.metadata_builders[0].kernel_block_size,
555 + )
556 + replay_groups.append(replay_group)
557 + self._replay_attn_groups.append(replay_groups)
558 + return self._replay_attn_groups
vllm/models/deepseek_v41/nvidia/vl_model.py
+5
@@ -51,6 +51,7 @@ from ..common.mm_preprocess import (
51 DeepseekV4VLProcessingInfo,
52 )
53 from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
54 +from ..decoder_replay_layers import DecoderReplayLayers
55 from .model import (
56 DeepseekV41LLMForCausalLM,
57 _linear_scale_param_name,
@@ -281,6 +282,10 @@ class DeepseekV41ForCausalLM(
282 def token_lookback_depth(self) -> int:
283 return self.language_model.token_lookback_depth
284
285 + @property
286 + def decoder_replay_layers(self) -> DecoderReplayLayers | None:
287 + return self.language_model.decoder_replay_layers
288 +
289 def forward(
290 self,
291 input_ids: torch.Tensor,