[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>

.buildkite/test_areas/models_basic.yaml

@@ -66,8 +66,10 @@ steps:
   - tests/models/transformers/test_backend.py
   - tests/models/test_registry.py
   - tests/models/test_deepseek_v4_vl_rocm.py
+  - tests/models/test_deepseek_v41_decoder_replay_layers.py
+  - tests/models/test_deepseek_v41_replay_batch.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
+    - 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
   mirror:
     amd:
       label: ":amd: (MI355 DPX) Basic Models (Other)"

tests/distributed/test_engram_dp_shard.py

@@ -507,20 +507,22 @@ def _check_dummy_hash_model_forward(
                     assert torch.all(hashes == engram_ops.DEAD_ID)
                 output = engram.embed(hashes[:, 0])
             # Trailing None is previous_aux; this stub captures no aux states.
-            return output, None, None, None, None, None
+            return output, hidden, hidden, hidden, hidden, None
 
-    model = SimpleNamespace(
-        use_mega_moe=False,
-        use_sequence_parallel=False,
-        fuse_mhc_all_reduce=False,
-        engram_hash=state,
-        engram_swa_prefix="swa",
-        engram_dp_shared_memory=dp_shared_memory,
-        layers=[Decoder(engram=engram)],
-        start_layer=0,
-        end_layer=1,
-        aux_hidden_state_layers=(),
-    )
+    model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model)
+    torch.nn.Module.__init__(model)
+    model.use_mega_moe = False
+    model.use_sequence_parallel = False
+    model.fuse_mhc_all_reduce = False
+    model.engram_hash = state
+    model.engram_swa_prefix = "swa"
+    model.engram_dp_shared_memory = dp_shared_memory
+    model.layers = [Decoder(engram=engram)]
+    model.start_layer = 0
+    model.end_layer = 1
+    model.decoder_replay_start = model.end_layer
+    model.decoder_replay_layers = None
+    model.aux_hidden_state_layers = ()
     metadata = (
         None
         if dp_rank == 1
@@ -550,8 +552,7 @@ def _check_dummy_hash_model_forward(
                 dtype=torch.int32,
             ),
         ):
-            output = model_ops.DeepseekV4Model.forward(
-                model,
+            output = model(
                 torch.arange(tokens, device="cuda"),
                 torch.arange(tokens, device="cuda"),
                 intermediate_tensors=None,

tests/models/test_deepseek_v41_decoder_replay_layers.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DecoderReplayLayers runs the replay layers on the forward's replay batch."""
+
+import pytest
+import torch
+
+from vllm.forward_context import (
+    ForwardContext,
+    get_forward_context,
+    override_forward_context,
+)
+from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+DEVICE = torch.device("cuda")
+NUM_TOKENS = 401
+HC, HIDDEN = 2, 3
+# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
+WINDOW = 128
+REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
+
+
+def _context(metadata):
+    return ForwardContext(
+        no_compile_layers={},
+        attn_metadata=metadata,
+        slot_mapping={},
+        is_padding=torch.zeros(NUM_TOKENS, dtype=torch.bool, device=DEVICE),
+    )
+
+
+def _set_replay_batch(layers, rows, metadata):
+    layers.rows = torch.tensor(rows, device=DEVICE)
+    layers.forward_context = ForwardContext(
+        no_compile_layers={}, attn_metadata=metadata, slot_mapping={}
+    )
+
+
+def _states(num_tokens: int):
+    """(hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual),
+    shaped as the model's: hidden states [T, H], mixes [T, HC], residual [T, HC, H]."""
+    hidden = torch.arange(num_tokens, device=DEVICE, dtype=torch.float32)
+    hidden = hidden[:, None].expand(num_tokens, HIDDEN).contiguous()
+    mix = hidden[:, :HC].contiguous()
+    residual = hidden[:, None, :].expand(num_tokens, HC, HIDDEN).contiguous()
+    return (
+        hidden,
+        torch.arange(num_tokens, device=DEVICE),
+        None,
+        mix + 1,
+        mix + 2,
+        mix + 3,
+        residual,
+    )
+
+
+def test_run_gathers_states_and_realigns_shared_indexer_buffers():
+    topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int()
+    candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int()
+    topk_before, candidates_before = topk.clone(), candidates.clone()
+    seen = {}
+
+    def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual):
+        seen["hidden_states"] = hidden_states.clone()
+        replay_context = get_forward_context()
+        seen["attn_metadata"] = replay_context.attn_metadata
+        seen["is_padding"] = replay_context.is_padding
+        return residual, pre_mix
+
+    layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates])
+    states = _states(NUM_TOKENS)
+    hidden, residual = states[0], states[-1]
+    full, sub = object(), object()
+    _set_replay_batch(layers, REPLAY_ROWS, sub)
+    context = _context(full)
+    with override_forward_context(context):
+        outputs = layers(*states)
+        assert get_forward_context() is context
+
+    rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
+    assert torch.equal(seen["hidden_states"], hidden[rows])
+    assert seen["attn_metadata"] is sub and seen["is_padding"] is None
+    assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows])
+    assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows])
+    assert len(outputs) == 2 and outputs[0].shape == residual.shape
+    assert torch.equal(outputs[0][rows], residual[rows])
+    assert outputs[0][1:173].abs().sum() == 0
+
+
+def test_no_replay_batch_runs_the_whole_batch():
+    seen = {}
+
+    def run_layers(hidden_states, *rest):
+        seen["attn_metadata"] = get_forward_context().attn_metadata
+        return (hidden_states,)
+
+    layers = DecoderReplayLayers(WINDOW, run_layers, [])
+    states = _states(NUM_TOKENS)
+    full = object()
+    with override_forward_context(_context(full)):
+        (output,) = layers(*states)
+    assert output is states[0] and seen["attn_metadata"] is full

tests/models/test_deepseek_v41_replay_batch.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The decoder-side SWA bounded replay batch DeepseekV41ModelState prepares."""
+
+from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="the replay buffers live on the GPU"
+)
+
+WINDOW = 128
+DEVICE = torch.device("cuda")
+GROUPS = SimpleNamespace(
+    kv_cache_groups=[
+        SimpleNamespace(
+            layer_names=["swa"],
+            kv_cache_spec=SimpleNamespace(
+                prefix_cacheable=False, prefix_replay_tokens=WINDOW
+            ),
+        ),
+        SimpleNamespace(
+            layer_names=["mla"],
+            kv_cache_spec=SimpleNamespace(
+                prefix_cacheable=True, prefix_replay_tokens=0
+            ),
+        ),
+    ]
+)
+BLOCK_TABLES = (torch.zeros(4, 4, device=DEVICE),) * 2
+
+# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300)
+QUERY_LENS = [1, 300, 100]
+SEQ_LENS = [500, 300, 300]
+PREFILLING = [False, True, True]
+REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)]
+
+
+@pytest.fixture
+def state(monkeypatch):
+    """A DeepseekV41ModelState whose attention builds are recorded, not run."""
+    cfg = MagicMock()
+    cfg.model_config.enable_prompt_embeds = False
+    cfg.model_config.uses_mrope = False
+    cfg.model_config.is_multimodal_model = False
+    cfg.scheduler_config.max_num_seqs = 16
+    cfg.scheduler_config.max_num_batched_tokens = 1024
+    cfg.parallel_config.data_parallel_size = 1
+    cfg.compilation_config.fast_moe_cold_start = False
+    layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None)
+    model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers)
+    builds: list = []
+
+    def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
+        build = SimpleNamespace(
+            batch=input_batch,
+            cg_mode=cg_mode,
+            slot_mappings=slot_mappings.clone(),
+            replay_start=kw["model_specific_attn_metadata"].replay_start,
+            attn_metadata=object(),
+        )
+        builds.append(build)
+        return {"swa": build.attn_metadata}
+
+    monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
+    state = DeepseekV41ModelState(cfg, model, None, DEVICE)
+    state.builds = builds  # type: ignore[attr-defined]
+    return state
+
+
+def _input_batch(
+    query_lens: list[int],
+    seq_lens: list[int],
+    is_prefilling: list[bool],
+    device_query_lens: list[int] | None = None,
+    num_tokens_after_padding: int | None = None,
+    max_query_len: int | None = None,
+) -> InputBatch:
+    num_reqs, num_tokens = len(query_lens), sum(query_lens)
+    num_padded = num_tokens_after_padding or num_tokens
+    batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(16, 1024, DEVICE))
+    return replace(
+        batch,
+        num_tokens=num_tokens,
+        num_tokens_after_padding=num_padded,
+        query_start_loc=torch.tensor(
+            [0, *np.cumsum(device_query_lens or query_lens)],
+            dtype=torch.int32,
+            device=DEVICE,
+        ),
+        query_start_loc_np=np.array([0, *np.cumsum(query_lens)], dtype=np.int32),
+        seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE),
+        seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
+        is_prefilling_np=np.array(is_prefilling, dtype=np.bool_),
+        positions=torch.cat(
+            [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)]
+            + [torch.zeros(num_padded - num_tokens, dtype=torch.int64)]
+        ).to(DEVICE),
+        is_padding=torch.arange(num_padded, device=DEVICE) >= num_tokens,
+        max_query_len=max_query_len,
+    )
+
+
+def _slot_mappings(num_tokens: int) -> torch.Tensor:
+    return torch.stack([torch.arange(num_tokens), torch.arange(num_tokens) * 10]).to(
+        DEVICE
+    )
+
+
+def _prepare(state, batch, cg_mode):
+    """Runs prepare_attn; returns the replay batch and the replay build's inputs."""
+    state.prepare_attn(
+        batch,
+        cg_mode,
+        BLOCK_TABLES,
+        _slot_mappings(batch.num_tokens_after_padding),
+        [],
+        GROUPS,
+    )
+    layers = state.decoder_replay_layers
+    if layers.rows is None:
+        assert layers.forward_context is None
+        return None, None
+    return SimpleNamespace(
+        rows=layers.rows, forward_context=layers.forward_context
+    ), state.builds[-1]
+
+
+def test_replay_batch_keeps_each_request_window(state):
+    state._replay_start_np[2] = 50  # the encoder-side replay start of request 2
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
+    replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
+    assert replay is not None
+    rows = torch.tensor(REPLAY_ROWS, device=DEVICE)
+    assert torch.equal(replay.rows, rows)
+    sub = build.batch
+    assert build.cg_mode == CUDAGraphMode.NONE
+    assert sub.query_start_loc.tolist() == [0, 1, 129, 229]
+    assert sub.query_start_loc_np.tolist() == [0, 1, 129, 229]
+    assert sub.num_tokens == sub.num_tokens_after_padding == 229
+    assert sub.max_query_len == WINDOW
+    # The trimmed request's window starts at its replay window; a higher
+    # encoder-side replay start stands.
+    assert build.replay_start.tolist() == [0, 300 - WINDOW, 50]
+    assert torch.equal(sub.positions, batch.positions[rows])
+    assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows])
+    context = replay.forward_context
+    assert context.attn_metadata == {"swa": build.attn_metadata}
+    assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows])
+    assert context.dp_metadata is None and context.is_padding is None
+
+
+def test_replay_batch_keeps_device_decode_boundaries(state):
+    """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device
+    [1, 3]): their rows stay and the boundaries shift only by trimmed rows."""
+    batch = _input_batch(
+        [2, 2, 300], [10, 10, 300], PREFILLING, device_query_lens=[1, 3, 300]
+    )
+    replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
+    assert replay is not None
+    assert replay.rows.tolist() == [0, 1, 2, 3, *range(304 - WINDOW, 304)]
+    assert build.batch.query_start_loc.tolist() == [0, 1, 4, 4 + WINDOW]
+
+
+def test_prompt_logprobs_requests_keep_their_rows(state):
+    """Every prompt row of a prompt-logprobs request is read, so it never
+    trims; the other prefills still do."""
+    state._keeps_rows_np[1] = True
+    batch = _input_batch([1, 300, 300], [500, 300, 300], PREFILLING)
+    replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
+    assert replay is not None
+    assert replay.rows.tolist() == [0, *range(1, 301), *range(601 - WINDOW, 601)]
+    assert build.batch.max_query_len == 300
+
+
+def test_replay_batch_keeps_adaptive_verification_query_bound(state):
+    """Adaptive verification bounds the decodes' device-side query lengths
+    above their CPU lengths; the replay batch keeps that bound."""
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, max_query_len=200)
+    replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
+    assert replay is not None and build.batch.max_query_len == 200
+
+
+def test_only_eager_steps_trim(state):
+    """A CUDA graph keeps the layers on its whole batch; an eager step trims a
+    padded batch to its real rows."""
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512)
+    for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
+        assert _prepare(state, batch, cg_mode) == (None, None)
+    replay, build = _prepare(state, batch, CUDAGraphMode.NONE)
+    assert replay is not None
+    assert replay.rows.tolist() == REPLAY_ROWS
+    assert build.cg_mode == CUDAGraphMode.NONE
+    assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS)
+
+
+def test_whole_batch_forwards_get_no_replay_batch(state):
+    short = _input_batch([100, 100], [100, 100], [True, True])
+    for cg_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL):
+        assert _prepare(state, short, cg_mode) == (None, None)
+    assert len(state.builds) == 3  # the batch's own metadata only
+    # Dummy batches, captures included (prepared as eager), are not prefills and
+    # keep their rows.
+    dummy = _input_batch(
+        QUERY_LENS, SEQ_LENS, [False] * 3, num_tokens_after_padding=512
+    )
+    assert _prepare(state, dummy, CUDAGraphMode.NONE) == (None, None)
+
+
+@pytest.fixture
+def dp_state(state, monkeypatch):
+    """`state` on DP rank 0 of 2; `state.other` sets what rank 1 reports
+    (trims, replay tokens)."""
+    import vllm.models.deepseek_v41.nvidia.model_state as module
+
+    state.vllm_config.parallel_config.data_parallel_size = 2
+    state.vllm_config.parallel_config.data_parallel_rank = 0
+    state.vllm_config.parallel_config.is_moe_model = True
+    state.other = (False, 0)  # type: ignore[attr-defined]
+
+    def all_reduce(tensor, group):
+        tensor[:, 1] = torch.tensor(state.other, dtype=torch.int32)
+
+    monkeypatch.setattr(module.dist, "all_reduce", all_reduce)
+    monkeypatch.setattr(module, "get_dp_group", lambda: SimpleNamespace(cpu_group=None))
+    return state
+
+
+def test_dp_ranks_replay_together(dp_state):
+    """A trimming rank makes the others replay too, on every rank's count."""
+    short = _input_batch([100, 100], [100, 100], [True, True])
+    dp_state.other = (False, 200)
+    assert _prepare(dp_state, short, CUDAGraphMode.NONE) == (None, None)
+
+    dp_state.other = (True, 300)  # rank 1 trims to 300 rows
+    replay, build = _prepare(dp_state, short, CUDAGraphMode.NONE)
+    assert replay is not None
+    assert replay.rows.tolist() == list(range(200))  # this rank keeps every row
+    dp_metadata = replay.forward_context.dp_metadata
+    assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [200, 300]
+
+
+def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
+    dummy = _input_batch([512], [512], [False])
+    dp_state.other = (True, 129)
+    replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE)
+    assert replay is not None and replay.rows.shape[0] == WINDOW
+    dp_metadata = replay.forward_context.dp_metadata
+    assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129]

vllm/config/cache.py

@@ -242,7 +242,9 @@ 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. The layers past the last KV-source layer then also run
+    eager prefill steps on each request's trailing window only. 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

vllm/models/deepseek_v41/decoder_replay_layers.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Decoder-side SWA bounded replay: running the replay layers on their batch.
+
+The layers past the last KV-source layer own nothing but sliding-window KV, so
+in eager prefill steps they run on each request's last ``window`` rows only.
+``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention
+metadata and a forward context of its own, like a microbatch;
+``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers
+under that context and scatters the outputs back to batch rows. Steps that run
+in a CUDA graph keep the layers on the whole batch, inside the graph.
+"""
+
+from collections.abc import Callable
+
+import torch
+
+from vllm.forward_context import ForwardContext, override_forward_context
+from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput
+
+
+class DecoderReplayLayers:
+    """Runs the replay layers on the step's replay batch.
+
+    ``run_layers`` takes a batch's layer inputs and returns its per-row
+    outputs. ``row_buffers`` hold per-row results the source layer's indexer
+    publishes for the layers after it; they are compacted to the replay rows
+    in place.
+    """
+
+    def __init__(
+        self,
+        window: int,
+        run_layers: Callable[..., tuple[torch.Tensor, ...]],
+        row_buffers: list[torch.Tensor],
+    ) -> None:
+        self.window = window
+        self.run_layers = run_layers
+        self.row_buffers = row_buffers
+        # The replay batch, set by the model state every step: its rows of the
+        # batch and its forward context. None runs the layers on the batch.
+        self.rows: torch.Tensor | None = None
+        self.forward_context: ForwardContext | None = None
+
+    def __call__(
+        self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None
+    ) -> tuple[torch.Tensor, ...]:
+        rows = self.rows
+        if rows is None:
+            return self.run_layers(hidden_states, *states)
+        # A trimming step holds a prefill longer than the window, more tokens
+        # than any step whose MoE leaves its finalize to the next layer.
+        assert isinstance(hidden_states, torch.Tensor)
+        num_rows = rows.shape[0]
+        for buf in self.row_buffers:
+            buf[:num_rows].copy_(buf.index_select(0, rows))
+        with override_forward_context(self.forward_context):
+            row_outputs = self.run_layers(
+                hidden_states.index_select(0, rows),
+                *(None if t is None else t.index_select(0, rows) for t in states),
+            )
+        # The trimmed rows' outputs stay zero; nothing reads them.
+        num_tokens = hidden_states.shape[0]
+        return tuple(
+            out.new_zeros((num_tokens, *out.shape[1:])).index_copy_(0, rows, out)
+            for out in row_outputs
+        )

vllm/models/deepseek_v41/nvidia/model.py

@@ -83,6 +83,7 @@ from vllm.models.deepseek_v4.nvidia.model import (
     prepare_mega_gate_routing_metadata,
 )
 from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
 from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
     DeepseekV4MegaAttnAttention,
 )
@@ -741,6 +742,35 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 if successor.engram is None:
                     layer.ffn.defer_finalize()
 
+        # Decoder-side SWA bounded replay: in eager prefill steps the layers past
+        # the last KV source run on each request's trailing window only
+        # (decoder_replay_layers.py).
+        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
+            self.decoder_replay_layers = DecoderReplayLayers(
+                config.sliding_window,
+                self._run_replay_layers,
+                [
+                    buf
+                    for buf in (self.topk_indices_buffer, self.candidate_block_buffer)
+                    if buf is not None
+                ],
+            )
+            logger.info_once(
+                "Decoder SWA bounded replay: in eager prefill steps, layers "
+                "%d-%d run on each request's last %d tokens only.",
+                cut + 1,
+                self.end_layer - 1,
+                config.sliding_window,
+            )
+
         # 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
@@ -904,29 +934,28 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             hidden_states = sp_shard(hidden_states)
             input_ids = sp_shard(input_ids)
 
-        mega_gate_metadata = None
-        if self.use_mega_moe:
-            mega_gate_metadata = prepare_mega_gate_routing_metadata(
-                input_ids,
-                has_hash_routing=False,
-                image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
-                if getattr(self.config, "vision_n_layers", 0) > 0
-                else None,
-            )
-
         residual, post_mix, res_mix = None, None, None
         pre_mix: torch.Tensor | None = None
         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,
@@ -934,36 +963,22 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 post_mix,
                 res_mix,
                 residual,
-                engram_hashes,
-                engram_mask,
-                capture_previous_aux=idx in self.aux_hidden_state_layers,
-                mega_gate_metadata=mega_gate_metadata,
             )
-            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
-        # Without a successor boundary, the last layer finalized its own MoE.
-        assert isinstance(hidden_states, torch.Tensor)
-        if layer is not None:
-            # The last layer has no successor to fold its post into.
-            if self.fuse_mhc_all_reduce:
-                hidden_states = tensor_model_parallel_all_reduce(hidden_states)
-            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(
@@ -992,6 +1007,169 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             return hidden_states, aux_hidden_states
         return hidden_states
 
+    def _mega_gate_metadata(
+        self, input_ids: torch.Tensor | None
+    ) -> MegaGateRoutingMetadata | None:
+        if not self.use_mega_moe:
+            return None
+        assert input_ids is not None
+        return prepare_mega_gate_routing_metadata(
+            input_ids,
+            has_hash_routing=False,
+            image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID
+            if getattr(self.config, "vision_n_layers", 0) > 0
+            else None,
+        )
+
+    def _run_layers(
+        self,
+        layer_ids: range,
+        hidden_states: torch.Tensor | MoEOutput,
+        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 | MoEOutput, 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]
+        mega_gate_metadata = self._mega_gate_metadata(input_ids)
+        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,
+                    mega_gate_metadata=mega_gate_metadata,
+                )
+            )
+            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
+        assert residual is not None and post_mix is not None
+        assert res_mix is not None and pre_mix is not None
+        return hidden_states, residual, post_mix, res_mix, pre_mix
+
+    def _collapse(
+        self,
+        hidden_states: torch.Tensor | MoEOutput,
+        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:
+        # Without a successor boundary, the last layer finalized its own MoE.
+        assert isinstance(hidden_states, torch.Tensor)
+        # The last layer has no successor to fold its post into.
+        if self.fuse_mhc_all_reduce:
+            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+        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_replay_layers(
+        self,
+        hidden_states: torch.Tensor | MoEOutput,
+        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, ...]:
+        """The layers past the last KV source, on whatever rows they are given;
+        returns their output, the last FFN's pre-mix and the aux hidden states
+        they capture."""
+        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:
+            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.prefill_context_parallel_size > 1
+            or parallel_config.use_ubatching
+        ):
+            reason = (
+                "the replay-layer batch shrinks per rank, which sequence and "
+                "prefill-context parallelism and microbatching 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)
+        ):
+            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)
@@ -1397,6 +1575,10 @@ class DeepseekV41LLMForCausalLM(
         engram_hash = self.model.engram_hash
         return engram_hash.lookback_depth if engram_hash is not None else 0
 
+    @property
+    def decoder_replay_layers(self) -> DecoderReplayLayers | None:
+        return self.model.decoder_replay_layers
+
     def forward(
         self,
         input_ids: torch.Tensor,

vllm/models/deepseek_v41/nvidia/model_state.py

@@ -1,17 +1,24 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import replace
 from typing import Any
 
 import numpy as np
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 
 from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed.parallel_state import get_dp_group
+from vllm.forward_context import DPMetadata, create_forward_context
+from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.dp_utils import should_skip_dp_coordination
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
 from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
@@ -81,6 +88,61 @@ def _pad_replayed_slots_kernel(
             )
 
 
+@triton.jit
+def _gather_replay_batch_kernel(
+    query_start_loc_ptr,  # [num_reqs + 1]
+    dropped_before_ptr,  # [num_reqs + 1] rows trimmed before each boundary
+    seq_lens_ptr,  # [num_reqs]
+    replay_start_ptr,  # [num_reqs] the encoder-side replay start
+    positions_ptr,
+    slot_mappings_ptr,  # [num_groups, num_tokens]
+    slot_mappings_stride,
+    rows_ptr,  # out: [num_kept] batch row of every kept row
+    kept_query_start_loc_ptr,  # out: [num_reqs + 1]
+    kept_positions_ptr,  # out: [num_kept]
+    kept_slot_mappings_ptr,  # out: [num_groups, num_kept]
+    kept_slot_mappings_stride,
+    kept_kv_start_ptr,  # out: [num_reqs]
+    window,
+    NUM_GROUPS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """One program per request: its kept rows are the last ones before its
+    boundary, and land at that boundary minus the rows trimmed before it."""
+    req = tl.program_id(0)
+    begin = tl.load(query_start_loc_ptr + req)
+    end = tl.load(query_start_loc_ptr + req + 1)
+    kept_begin = begin - tl.load(dropped_before_ptr + req)
+    kept_end = end - tl.load(dropped_before_ptr + req + 1)
+    if req == 0:
+        tl.store(kept_query_start_loc_ptr, kept_begin)
+    tl.store(kept_query_start_loc_ptr + req + 1, kept_end)
+
+    # A trimmed request holds no replay-layer window KV below its kept rows.
+    kv_start = tl.load(replay_start_ptr + req)
+    if kept_end - kept_begin < end - begin:
+        kv_start = tl.maximum(kv_start, tl.load(seq_lens_ptr + req) - window)
+    tl.store(kept_kv_start_ptr + req, kv_start)
+
+    first_row = end - (kept_end - kept_begin)
+    for tok in range(kept_begin, kept_end, BLOCK):
+        offs = tok + tl.arange(0, BLOCK)
+        mask = offs < kept_end
+        row = first_row + (offs - kept_begin)
+        tl.store(rows_ptr + offs, row.to(tl.int64), mask=mask)
+        tl.store(
+            kept_positions_ptr + offs,
+            tl.load(positions_ptr + row, mask=mask),
+            mask=mask,
+        )
+        for g in tl.static_range(NUM_GROUPS):
+            tl.store(
+                kept_slot_mappings_ptr + g * kept_slot_mappings_stride + offs,
+                tl.load(slot_mappings_ptr + g * slot_mappings_stride + row, mask=mask),
+                mask=mask,
+            )
+
+
 class ReplayAttnMetadata(ModelSpecificAttnMetadata):
     """Hands the batch's replay starts to the sliding-window builders."""
 
@@ -111,7 +173,10 @@ class DeepseekV41ModelState(DefaultModelState):
     ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots
     in the prefix-cacheable groups so the cached KV stays as is, and hands the
     starts to the sliding-window metadata builders, whose kernels read no
-    window KV below them.
+    window KV below them. With the decoder side on
+    (``model.decoder_replay_layers``) it also prepares the replay layers'
+    batch in eager steps: each prefill's last ``window`` rows as a sub-batch
+    with metadata and a forward context of its own, like a microbatch.
     """
 
     def __init__(
@@ -140,9 +205,40 @@ class DeepseekV41ModelState(DefaultModelState):
         # the KV cache config on first use; None until then.
         self._replay: tuple[int, torch.Tensor] | None = None
 
+        self.decoder_replay_layers: DecoderReplayLayers | None = getattr(
+            model, "decoder_replay_layers", None
+        )
+        if self.decoder_replay_layers is not None:
+            # Requests whose every prompt row is read (prompt logprobs) never trim.
+            self._keeps_rows_np = np.zeros(self.max_num_reqs, dtype=np.bool_)
+            self._dropped_before = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
+            # The replay batch's inputs: its rows of the batch, query offsets,
+            # positions, slot mappings (sized with the KV cache groups on first
+            # use) and the position from which each request holds replay-layer
+            # window KV.
+            self._kept_rows = torch.zeros(
+                self.max_num_tokens, dtype=torch.int64, device=device
+            )
+            self._kept_query_start_loc = torch.zeros(
+                self.max_num_reqs + 1, dtype=torch.int32, device=device
+            )
+            self._kept_positions = torch.zeros(
+                self.max_num_tokens, dtype=torch.int64, device=device
+            )
+            self._kept_slot_mappings: torch.Tensor | None = None
+            self._kept_kv_start = torch.zeros(
+                self.max_num_reqs, dtype=torch.int32, device=device
+            )
+            self._replay_attn_groups: list[list[AttentionGroup]] | None = None
+
     def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
         super().add_request(req_index, new_req_data)
         self._replay_start_np[req_index] = new_req_data.replay_start
+        if self.decoder_replay_layers is not None:
+            params = new_req_data.sampling_params
+            self._keeps_rows_np[req_index] = (
+                params is not None and params.prompt_logprobs is not None
+            )
 
     def prepare_inputs(
         self, input_batch: InputBatch, req_states: RequestState
@@ -196,7 +292,16 @@ class DeepseekV41ModelState(DefaultModelState):
                     device=self.device,
                 ),
             )
+            if self.decoder_replay_layers is not None:
+                self._kept_slot_mappings = torch.zeros(
+                    len(specs),
+                    self.max_num_tokens,
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+            self._warm_up_replay_kernels(input_batch, slot_mappings)
         window, cacheable_groups = self._replay
+        replay_start: torch.Tensor | None = None
         if window:
             num_reqs = input_batch.num_reqs
             # Decode rows sit above the hit, so only prefills carry a replay
@@ -226,7 +331,7 @@ class DeepseekV41ModelState(DefaultModelState):
                 )
             assert model_specific_attn_metadata is None
             model_specific_attn_metadata = ReplayAttnMetadata(replay_start)
-        return super().prepare_attn(
+        attn_metadata = super().prepare_attn(
             input_batch,
             cudagraph_mode,
             block_tables,
@@ -237,3 +342,217 @@ class DeepseekV41ModelState(DefaultModelState):
             ubatch_idx=ubatch_idx,
             model_specific_attn_metadata=model_specific_attn_metadata,
         )
+        if self.decoder_replay_layers is not None:
+            assert replay_start is not None
+            self._prepare_replay_batch(
+                input_batch,
+                cudagraph_mode,
+                block_tables,
+                slot_mappings,
+                attn_groups,
+                kv_cache_config,
+                replay_start,
+            )
+        return attn_metadata
+
+    def _warm_up_replay_kernels(
+        self, input_batch: InputBatch, slot_mappings: torch.Tensor
+    ) -> None:
+        """Compile the replay kernels on the buffers they run on: the startup
+        dummy batches never replay or trim, so they would not launch them."""
+        assert self._replay is not None
+        window, cacheable_groups = self._replay
+        if not window:
+            return
+        _pad_replayed_slots_kernel.warmup(
+            slot_mappings,
+            slot_mappings.stride(0),
+            cacheable_groups,
+            input_batch.query_start_loc,
+            input_batch.positions,
+            self._replay_start,
+            window,
+            PAD_SLOT_ID,
+            NUM_GROUPS=cacheable_groups.numel(),
+            BLOCK=1024,
+            grid=(1,),
+        )
+        if self.decoder_replay_layers is None:
+            return
+        assert self._kept_slot_mappings is not None
+        _gather_replay_batch_kernel.warmup(
+            input_batch.query_start_loc,
+            self._dropped_before.copy_to_uva(np.zeros(1, dtype=np.int32)),
+            input_batch.seq_lens,
+            self._replay_start,
+            input_batch.positions,
+            slot_mappings,
+            slot_mappings.stride(0),
+            self._kept_rows,
+            self._kept_query_start_loc,
+            self._kept_positions,
+            self._kept_slot_mappings,
+            self._kept_slot_mappings.stride(0),
+            self._kept_kv_start,
+            self.decoder_replay_layers.window,
+            NUM_GROUPS=slot_mappings.shape[0],
+            BLOCK=1024,
+            grid=(1,),
+        )
+
+    def _prepare_replay_batch(
+        self,
+        input_batch: InputBatch,
+        cudagraph_mode: CUDAGraphMode,
+        block_tables: tuple[torch.Tensor, ...],
+        slot_mappings: torch.Tensor,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        replay_start: torch.Tensor,
+    ) -> None:
+        """Set the replay layers' batch for this forward, or none when nothing
+        trims. Only eager steps trim: a CUDA graph keeps the layers on the whole
+        batch, and ranks share the graph mode. Under data parallelism every rank
+        replays if any does."""
+        layers = self.decoder_replay_layers
+        assert layers is not None
+        layers.rows = layers.forward_context = None
+        if cudagraph_mode != CUDAGraphMode.NONE:
+            return
+
+        num_reqs = input_batch.num_reqs
+        query_lens = np.diff(input_batch.query_start_loc_np[: num_reqs + 1])
+        kept_lens = np.where(
+            self._keeps_rows_np[input_batch.idx_mapping_np[:num_reqs]],
+            query_lens,
+            np.minimum(query_lens, layers.window),
+        )
+        # Dummy batches (captures, an idle DP rank's) are not prefills: they
+        # keep their rows unless a rank trims.
+        trims = bool(
+            (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any()
+        )
+        num_tokens = int(kept_lens.sum())
+        dp_metadata = None
+        if self.vllm_config.parallel_config.data_parallel_size > 1:
+            trims, dp_metadata = self._agree_across_dp(trims, num_tokens)
+        if not trims:
+            return
+
+        kept_batch, kept_slot_mappings = self._kept_input_batch(
+            input_batch, slot_mappings, replay_start, kept_lens
+        )
+        attn_metadata = super().prepare_attn(
+            kept_batch,
+            CUDAGraphMode.NONE,
+            block_tables,
+            kept_slot_mappings,
+            self._replay_groups(attn_groups),
+            kv_cache_config,
+            model_specific_attn_metadata=ReplayAttnMetadata(
+                self._kept_kv_start[:num_reqs]
+            ),
+        )
+        layers.rows = self._kept_rows[:num_tokens]
+        layers.forward_context = create_forward_context(
+            attn_metadata,
+            self.vllm_config,
+            dp_metadata=dp_metadata,
+            slot_mapping=build_slot_mappings_by_layer(
+                kept_slot_mappings, kv_cache_config
+            ),
+        )
+
+    def _kept_input_batch(
+        self,
+        input_batch: InputBatch,
+        slot_mappings: torch.Tensor,
+        replay_start: torch.Tensor,
+        kept_lens: np.ndarray,
+    ) -> tuple[InputBatch, torch.Tensor]:
+        """Build the sub-``InputBatch`` of each request's last ``kept_lens``
+        rows, and its slot mappings.
+
+        The kept rows follow the device boundaries minus the rows trimmed before
+        them: adaptive verification resizes the decodes on the GPU alone, and
+        decodes never trim. Like a microbatch's, the sub-batch describes the
+        forward only; its sampling fields are carried over and must not be read.
+        """
+        num_reqs = input_batch.num_reqs
+        num_tokens = int(kept_lens.sum())
+        kept_query_start_loc_np = np.zeros(num_reqs + 1, dtype=np.int32)
+        np.cumsum(kept_lens, out=kept_query_start_loc_np[1:])
+        dropped_before = self._dropped_before.copy_to_uva(
+            input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np
+        )
+        assert self._kept_slot_mappings is not None
+        kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens]
+        _gather_replay_batch_kernel[(num_reqs,)](
+            input_batch.query_start_loc,
+            dropped_before,
+            input_batch.seq_lens,
+            replay_start,
+            input_batch.positions,
+            slot_mappings,
+            slot_mappings.stride(0),
+            self._kept_rows,
+            self._kept_query_start_loc,
+            self._kept_positions,
+            kept_slot_mappings,
+            kept_slot_mappings.stride(0),
+            self._kept_kv_start,
+            self.decoder_replay_layers.window,  # type: ignore[union-attr]
+            NUM_GROUPS=slot_mappings.shape[0],
+            BLOCK=1024,
+        )
+        # Adaptive verification may lengthen the decodes on the device, up to
+        # the batch's bound.
+        max_query_len = max(int(kept_lens.max()), input_batch.max_query_len or 0)
+        kept_batch = replace(
+            input_batch,
+            num_tokens=num_tokens,
+            num_tokens_after_padding=num_tokens,
+            query_start_loc=self._kept_query_start_loc[: num_reqs + 1],
+            query_start_loc_np=kept_query_start_loc_np,
+            max_query_len=max_query_len,
+            positions=self._kept_positions[:num_tokens],
+            fast_prefill=None,
+        )
+        return kept_batch, kept_slot_mappings
+
+    def _agree_across_dp(
+        self, trims: bool, num_tokens: int
+    ) -> tuple[bool, DPMetadata | None]:
+        """Whether any rank trims and, then, the replay layers' DP metadata from
+        every rank's replay token count."""
+        parallel_config = self.vllm_config.parallel_config
+        agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32)
+        agreed[:, parallel_config.data_parallel_rank] = torch.tensor(
+            [trims, num_tokens]
+        )
+        if not should_skip_dp_coordination():
+            dist.all_reduce(agreed, group=get_dp_group().cpu_group)
+        if not agreed[0].any():
+            return False, None
+        return True, DPMetadata.make(parallel_config, num_tokens, agreed[1])
+
+    def _replay_groups(
+        self, attn_groups: list[list[AttentionGroup]]
+    ) -> list[list[AttentionGroup]]:
+        """The attention groups with metadata builders of their own, like each
+        microbatch's: a builder keeps the metadata it built, and the runner's
+        hold the batch's."""
+        if self._replay_attn_groups is None:
+            self._replay_attn_groups = []
+            for groups in attn_groups:
+                replay_groups = []
+                for group in groups:
+                    replay_group = replace(group, metadata_builders=[])
+                    replay_group.create_metadata_builders(
+                        self.vllm_config,
+                        self.device,
+                        group.metadata_builders[0].kernel_block_size,
+                    )
+                    replay_groups.append(replay_group)
+                self._replay_attn_groups.append(replay_groups)
+        return self._replay_attn_groups

vllm/models/deepseek_v41/nvidia/vl_model.py

@@ -51,6 +51,7 @@ from ..common.mm_preprocess import (
     DeepseekV4VLProcessingInfo,
 )
 from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin
+from ..decoder_replay_layers import DecoderReplayLayers
 from .model import (
     DeepseekV41LLMForCausalLM,
     _linear_scale_param_name,
@@ -281,6 +282,10 @@ class DeepseekV41ForCausalLM(
     def token_lookback_depth(self) -> int:
         return self.language_model.token_lookback_depth
 
+    @property
+    def decoder_replay_layers(self) -> DecoderReplayLayers | None:
+        return self.language_model.decoder_replay_layers
+
     def forward(
         self,
         input_ids: torch.Tensor,