[ROCm][Perf][DSv4.1] Enable decoder SWA bounded replay (#59039)

Signed-off-by: Frida Andersson <fanderss@amd.com>

.buildkite/test_areas/models_basic.yaml

@@ -82,18 +82,20 @@ steps:
       - image-build-amd
       source_file_dependencies:
       - vllm/models/deepseek_v4/
+      - vllm/models/deepseek_v41/
       - tests/models/test_deepseek_v4_vl_rocm.py
       - tests/models/test_hyv4_rocm.py
       - tests/models/transformers/test_layer_registry_e2e.py
       - tests/models/test_deepseek_v41_replay_start.py
       - tests/models/test_deepseek_v41_decoder_replay_layers.py
       - tests/models/test_deepseek_v41_replay_batch.py
+      - tests/models/test_deepseek_v41_amd_replay_forward.py
       - tests/models/test_deepseek_v4_dspark_rocm.py
       - tests/models/test_deepseek_v4_rocm_wo_a.py
       - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
       - tests/models/test_dsv41_mono_decode.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 models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.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/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
 
 - label: ":nvidia: (B200) Inkling"

tests/models/test_deepseek_v41_amd_replay_forward.py

@@ -0,0 +1,329 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""What the AMD replay forward does with a prefill longer than the window.
+
+The cut in this fixture is layer 1, so replay starts at layer 1. That is the
+same boundary as production layer 20, on a 4-layer stand-in. The final
+``mhc_post`` and the aux capture at that cut both run on the window rows.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.forward_context import ForwardContext, override_forward_context
+from vllm.models.deepseek_v41.decoder_replay_layers import ReplayBatch
+from vllm.sequence import IntermediateTensors
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+WINDOW = 128
+FULL = 300
+# Trailing window of a 300-token prefill.
+ROWS = list(range(FULL - WINDOW, FULL))
+HC, HIDDEN = 4, 8
+DEVICE = torch.device("cuda")
+# gfx942 writes the V4 record (584 B). Family-100 writes the V4.1 record (528 B).
+# The SWA layer uses block size 32, which is the size the KV-only insert checks.
+_BLOCK = 32
+_HEAD_DIM = 512
+
+
+class _RocmKvInsert:
+    """SWA cache from ``forward_kv``, checked against the KV-only insert."""
+
+    def __init__(self, num_tokens: int):
+        from vllm.models.deepseek_v41.attention import _use_v41_mxfp8_kv_record
+
+        if not hasattr(torch.ops._C, "fused_deepseek_v4_kv_rope_insert"):
+            pytest.skip("fused DeepseekV4 KV insert op is not built")
+
+        self.num_tokens = num_tokens
+        self.kv_mxfp8 = _use_v41_mxfp8_kv_record()
+        bytes_per = 528 if self.kv_mxfp8 else 584
+        num_blocks = (num_tokens + _BLOCK - 1) // _BLOCK + 1
+        cache = torch.full(
+            (num_blocks, _BLOCK, bytes_per),
+            0xA5,
+            dtype=torch.uint8,
+            device=DEVICE,
+        )
+        self.poison = cache.clone()
+        self.kv = torch.randn(
+            num_tokens, _HEAD_DIM, dtype=torch.bfloat16, device=DEVICE
+        )
+        prefix = "replay.swa"
+        slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=DEVICE)
+        self.metadata = {
+            prefix: SimpleNamespace(slot_mapping=slot_mapping, block_size=_BLOCK)
+        }
+        self.context = ForwardContext({}, self.metadata, {})
+        inv_freq = 1.0 / (
+            10000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device=DEVICE) / 64)
+        )
+        freqs = torch.outer(
+            torch.arange(num_tokens + 8, dtype=torch.float32, device=DEVICE), inv_freq
+        )
+        self.attn = SimpleNamespace(
+            kv_mxfp8=self.kv_mxfp8,
+            compressor=None,
+            indexer=None,
+            rotary_emb=SimpleNamespace(
+                cos_sin_cache=torch.cat((freqs.cos(), freqs.sin()), dim=-1)
+            ),
+            swa_cache_layer=SimpleNamespace(
+                prefix=prefix, kv_cache=cache, block_size=_BLOCK
+            ),
+            _run_parallel_input_projections=lambda hidden: (hidden, None, None),
+            _split_qkv_and_norm=lambda qr_kv: (None, None, self.kv[: qr_kv.shape[0]]),
+        )
+        self.last_rows: int | None = None
+        self.last_positions: torch.Tensor | None = None
+
+    def write(
+        self,
+        x,
+        positions,
+        input_ids,
+        pre_mix,
+        post_mix,
+        res_mix,
+        residual,
+    ):
+        from vllm.models.deepseek_v41.amd.rocm import (
+            DeepseekV41ROCMAiterMLAAttention,
+        )
+        from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+
+        assert (
+            DeepseekV41ROCMAiterMLAAttention.forward_kv
+            is DeepseekV4Attention.forward_kv
+        )
+        self.last_rows = x.shape[0]
+        self.last_positions = positions
+        DeepseekV41ROCMAiterMLAAttention.forward_kv(self.attn, positions, x)
+
+    def assert_matches_forward(self):
+        assert self.last_rows == self.num_tokens
+        assert self.last_positions is not None
+        written = self.attn.swa_cache_layer.kv_cache
+        assert not torch.equal(written, self.poison)
+        ref = self.poison.clone()
+        meta = self.metadata[self.attn.swa_cache_layer.prefix]
+        torch.ops._C.fused_deepseek_v4_kv_rope_insert(
+            self.kv,
+            ref,
+            meta.slot_mapping,
+            self.last_positions,
+            self.attn.rotary_emb.cos_sin_cache,
+            _BLOCK,
+            None,
+            self.kv_mxfp8,
+        )
+        assert torch.equal(written, ref)
+        assert torch.equal(written[-1], self.poison[-1])
+
+
+class _Layer(nn.Module):
+    def __init__(
+        self,
+        idx: int,
+        compressed: str | None = None,
+        indexer: str | None = None,
+    ):
+        super().__init__()
+        self.idx = idx
+        self.seen: list[torch.Tensor] = []
+        self.attn = SimpleNamespace(
+            swa_cache_layer=SimpleNamespace(prefix=f"layers.{idx}.swa"),
+            compressed_cache_prefix=compressed,
+            indexer=(
+                None
+                if indexer is None
+                else SimpleNamespace(k_cache=SimpleNamespace(prefix=indexer))
+            ),
+            compress_ratio=1,
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        positions,
+        input_ids,
+        pre_mix,
+        post_mix,
+        res_mix,
+        residual,
+        engram_hashes,
+        engram_mask,
+    ):
+        self.seen.append(hidden_states[:, 0, 0].detach().clone())
+        if residual is None:
+            residual = hidden_states
+        if pre_mix is None:
+            pre_mix = hidden_states[:, :, 0]
+        if post_mix is None:
+            post_mix = pre_mix
+        if res_mix is None:
+            res_mix = pre_mix
+        return hidden_states, residual, post_mix, res_mix, pre_mix
+
+
+def _model(monkeypatch):
+    import vllm.models.deepseek_v41.amd.model as amd
+
+    monkeypatch.setattr(
+        amd,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False),
+    )
+    monkeypatch.setattr(amd.EngramLayout, "from_config", lambda config: None)
+    monkeypatch.setattr(amd, "MHCPostOp", lambda: SimpleNamespace())
+
+    # Layer 0 sits before the cut, so its keys are not collected. Layer 2
+    # shares layer 1's compressed cache. Layer 3 has neither.
+    layers = [
+        _Layer(0, "layers.0.compressed", "layers.0.indexer"),
+        _Layer(1, "layers.1.compressed", "layers.1.indexer"),
+        _Layer(2, "layers.1.compressed", "layers.2.indexer"),
+        _Layer(3),
+    ]
+    # Layer 1 is the KV source, so the replay callback is this layer's write.
+    insert = _RocmKvInsert(FULL)
+    layers[1].write_kv = insert.write
+    monkeypatch.setattr(amd, "make_layers", lambda n, factory, prefix: (0, n, layers))
+
+    config = SimpleNamespace(
+        vocab_size=32,
+        hc_eps=1e-6,
+        hc_mult=HC,
+        hidden_size=HIDDEN,
+        rms_norm_eps=1e-6,
+        index_topk=4,
+        kv_source_layer_ids=[1],
+        num_hidden_layers=4,
+        sliding_window=WINDOW,
+        engram_layer_ids=(),
+        candidate_source_layer_id=-1,
+        candidate_topk_blocks=0,
+    )
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=config, dtype=torch.bfloat16),
+        quant_config=None,
+        cache_config=SimpleNamespace(swa_bounded_replay=True),
+        use_v2_model_runner=True,
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=1,
+            use_ubatching=False,
+            enable_expert_parallel=False,
+        ),
+        speculative_config=None,
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=FULL),
+    )
+    model = amd.DeepseekV4Model(vllm_config=vllm_config)
+    # Production requests aux at the first replay layer. Here that id is the cut.
+    model.aux_hidden_state_layers = (model.decoder_replay_start,)
+    model.topk_indices_buffer = model.topk_indices_buffer.to(DEVICE)
+
+    calls: list[torch.Tensor] = []
+
+    def mhc_post(hidden, residual, post_mix, res_mix):
+        calls.append(hidden[:, 0, 0].detach().clone())
+        return hidden
+
+    model.mhc_post = mhc_post
+    saved_aux: list[list[torch.Tensor]] = []
+    run_layers = model._run_layers
+
+    def wrapped(*args, **kwargs):
+        result = run_layers(*args, **kwargs)
+        # Snapshot: forward() extends the early list with the scattered aux.
+        saved_aux.append(list(result[5]))
+        return result
+
+    model._run_layers = wrapped
+    scattered: list[tuple[torch.Tensor, ...]] = []
+    replay = model.decoder_replay_layers
+    run = replay._run
+
+    def wrapped_run(*args, **kwargs):
+        result = run(*args, **kwargs)
+        scattered.append(result)
+        return result
+
+    replay._run = wrapped_run
+    return model, layers, calls, saved_aux, scattered, insert
+
+
+def _forward(model, insert: _RocmKvInsert):
+    token_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32)
+    hidden = token_ids[:, None, None].expand(FULL, HC, HIDDEN).contiguous()
+    pre_mix = torch.ones(FULL, HC, device=DEVICE)
+    rows = torch.tensor(ROWS, device=DEVICE)
+    model.decoder_replay_layers.replay_batch = ReplayBatch(
+        rows,
+        ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}),
+    )
+    positions = torch.arange(FULL, device=DEVICE)
+    with override_forward_context(insert.context):
+        out = model(
+            input_ids=positions.to(torch.int32),
+            positions=positions,
+            intermediate_tensors=IntermediateTensors(
+                {"hidden_states": hidden, "pre_mix": pre_mix}
+            ),
+        )
+    return out
+
+
+def test_amd_replay_forward_trims_to_the_window(monkeypatch):
+    model, layers, calls, saved_aux, scattered, insert = _model(monkeypatch)
+    replay_layers = model.decoder_replay_layers
+    assert model.decoder_replay_start == 1
+    assert replay_layers.first_swa_prefix == "layers.1.swa"
+    assert replay_layers.metadata_prefixes == {
+        "layers.1.swa",
+        "layers.1.compressed",
+        "layers.1.indexer",
+        "layers.2.swa",
+        "layers.2.indexer",
+        "layers.3.swa",
+    }
+    assert replay_layers.replay_batch is None
+    out = _forward(model, insert)
+    insert.assert_matches_forward()
+    replay = model.decoder_replay_layers.replay_batch
+    assert replay is not None and replay.rows.tolist() == ROWS
+
+    early = [layer.seen[0] for layer in layers[: model.decoder_replay_start]]
+    late = [layer.seen[0] for layer in layers[model.decoder_replay_start :]]
+    full_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32)
+    window_ids = full_ids[ROWS]
+    assert [t.tolist() for t in early] == [full_ids.tolist()]
+    assert [t.tolist() for t in late] == [
+        window_ids.tolist(),
+        window_ids.tolist(),
+        window_ids.tolist(),
+    ]
+
+    # Scatter writes the replay rows back and leaves the trimmed rows at zero.
+    hidden = out["hidden_states"][:, 0, 0]
+    assert torch.equal(hidden[ROWS], window_ids)
+    assert hidden[: FULL - WINDOW].eq(0).all()
+
+    # calls[0] is the cut-aux capture. calls[1] is the final mhc_post.
+    # Both run inside the replay batch, before the scatter back to FULL.
+    assert len(calls) == 2
+    assert [t.shape[0] for t in calls] == [WINDOW, WINDOW]
+    assert torch.equal(calls[0], window_ids)
+    assert torch.equal(calls[1], window_ids)
+
+    # The early loop skips this id. The replay loop does not see it either.
+    assert saved_aux[0] == []
+    assert saved_aux[1] == []
+    _hidden, _pre_mix, aux = scattered[0]
+    assert aux.shape[0] == FULL
+    assert torch.equal(aux[ROWS, 0], window_ids)
+    assert aux[: FULL - WINDOW].eq(0).all()

tests/models/test_deepseek_v41_replay_batch.py

@@ -290,6 +290,81 @@ def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state):
     assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129]
 
 
+def test_zero_prefix_window_warms_gather_and_skips_pad(monkeypatch):
+    """ROCm leaves the prefix-cache window at 0. Replay still needs replay_start
+    and a warmed gather kernel. The pad kernel stays on the prefix-cache path."""
+    import vllm.models.deepseek_v41.nvidia.model_state as model_state
+
+    calls = {"pad_warmup": 0, "gather_warmup": 0, "pad_launch": 0}
+
+    def pad_warmup(*args, **kwargs):
+        calls["pad_warmup"] += 1
+
+    def gather_warmup(*args, **kwargs):
+        calls["gather_warmup"] += 1
+
+    class _Pad:
+        warmup = staticmethod(pad_warmup)
+
+        def __getitem__(self, grid):
+            def launch(*args, **kwargs):
+                calls["pad_launch"] += 1
+
+            return launch
+
+    monkeypatch.setattr(model_state, "_pad_replayed_slots_kernel", _Pad())
+    monkeypatch.setattr(
+        model_state._gather_replay_batch_kernel, "warmup", gather_warmup
+    )
+
+    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
+    cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+    layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set())
+    model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers)
+    groups = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(
+                layer_names=["swa"],
+                kv_cache_spec=SimpleNamespace(
+                    prefix_cacheable=False, prefix_replay_tokens=0
+                ),
+            ),
+            SimpleNamespace(
+                layer_names=["mla"],
+                kv_cache_spec=SimpleNamespace(
+                    prefix_cacheable=True, prefix_replay_tokens=0
+                ),
+            ),
+        ]
+    )
+
+    def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw):
+        return {"swa": kw["model_specific_attn_metadata"]}
+
+    monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn)
+    state = DeepseekV41ModelState(cfg, model, None, DEVICE)
+    state._replay_start_np[0] = 16
+    batch = _input_batch([8], [24], [True])
+    state.prepare_attn(
+        batch,
+        CUDAGraphMode.NONE,
+        BLOCK_TABLES,
+        _slot_mappings(batch.num_tokens_after_padding),
+        [],
+        groups,
+    )
+    assert state._replay[0] == 0
+    assert calls == {"pad_warmup": 0, "gather_warmup": 1, "pad_launch": 0}
+    assert state._replay_start[:1].tolist() == [16]
+
+
 def test_dp_ranks_pad_to_one_replay_graph(dp_state):
     """Every rank pads its replay batch to the graph fitting the largest rank's."""
     dp_state.replay_cudagraphs = MagicMock()

vllm/models/deepseek_v41/amd/model.py

@@ -62,6 +62,7 @@ from vllm.models.deepseek_v4.amd.model import (
 from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer
 from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention
 from vllm.models.deepseek_v41.attention import DeepseekV4Attention
+from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -421,6 +422,42 @@ class DeepseekV4DecoderLayer(nn.Module):
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix, ffn_pre
 
+    def write_kv(
+        self,
+        x: torch.Tensor,
+        positions: torch.Tensor,
+        input_ids: torch.Tensor | None,
+        pre_mix: torch.Tensor,
+        post_mix: torch.Tensor,
+        res_mix: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> None:
+        """Write the KV ``forward`` would for the replay layers' inputs, skipping
+        the rest of the layer; the replay batch then reruns its rows."""
+        # The replay batch rules out Engram and sequence parallel here.
+        assert self.engram is None and not self.use_sequence_parallel
+        fuse_attn_norm = self.fuse_seam_norm
+        residual, post_mix, res_mix, x, _attn_pre = self.mhc_pre_delayed(
+            residual,
+            self.hc_attn_fn,
+            self.hc_attn_scale,
+            self.hc_attn_base,
+            self.rms_norm_eps,
+            self.hc_eps,
+            self.hc_eps,
+            self.hc_post_alpha,
+            self.hc_sinkhorn_iters,
+            pre_mix=pre_mix,
+            sublayer_out=x,
+            post_layer_mix=post_mix,
+            comb_res_mix=res_mix,
+            norm_weight=self.attn_norm.weight if fuse_attn_norm else None,
+            norm_eps=self.attn_norm.variance_epsilon,
+        )
+        if not fuse_attn_norm:
+            x = self.attn_norm(x)
+        self.attn.forward_kv(positions, x)
+
 
 class DeepseekV4Model(nn.Module, EagleModelMixin):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -498,6 +535,45 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             prefix=f"{prefix}.layers",
         )
 
+        # Config flag is the switch: ROCm forces swa_cache_layer.bounded_replay off.
+        self.decoder_replay_layers: DecoderReplayLayers | None = None
+        self.decoder_replay_start = self.end_layer
+        cut = max(config.kv_source_layer_ids)
+        if (
+            vllm_config.cache_config.swa_bounded_replay
+            and vllm_config.use_v2_model_runner
+            and cut < self.end_layer - 1
+            and self._decoder_replay_supported(vllm_config, cut)
+        ):
+            assert self.layers[cut].attn.compress_ratio <= 1
+            # The last KV source writes every row's KV, then replays.
+            self.decoder_replay_start = cut
+            # The attention metadata keys the replay layers read.
+            metadata_prefixes: set[str] = set()
+            for layer in islice(self.layers, cut, self.end_layer):
+                attn = typing.cast(DeepseekV4DecoderLayer, layer).attn
+                metadata_prefixes.add(attn.swa_cache_layer.prefix)
+                if attn.compressed_cache_prefix is not None:
+                    metadata_prefixes.add(attn.compressed_cache_prefix)
+                if attn.indexer is not None:
+                    metadata_prefixes.add(attn.indexer.k_cache.prefix)
+            self.decoder_replay_layers = DecoderReplayLayers(
+                config.sliding_window,
+                self._run_replay_layers,
+                self.layers[cut].write_kv,
+                metadata_prefixes,
+                self.layers[cut].attn.swa_cache_layer.prefix,
+            )
+            logger.info_once(
+                "Decoder SWA bounded replay: in eager prefill steps, layer %d "
+                "writes its KV for every token, and the rest of it and layers "
+                "%d-%d run on each request's last %d tokens only.",
+                cut,
+                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
@@ -657,13 +733,33 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
         if not get_pp_group().is_first_rank:
             assert intermediate_tensors is not None
             pre_mix = intermediate_tensors["pre_mix"]
-        aux_hidden_states: list[torch.Tensor] = []
-        final_aux_recon: torch.Tensor | None = None  # avoid duplicate mhc_post call
-        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 = layer(
+        (
+            hidden_states,
+            residual,
+            post_mix,
+            res_mix,
+            pre_mix,
+            aux_hidden_states,
+            final_aux_recon,
+            ran_layers,
+        ) = self._run_layers(
+            range(self.start_layer, self.decoder_replay_start),
+            hidden_states,
+            positions,
+            input_ids,
+            pre_mix,
+            post_mix,
+            res_mix,
+            residual,
+            engram_hashes,
+            engram_mask,
+            full_num_tokens,
+        )
+        # Replay collapses inside its own batch. The early loop still sets
+        # ran_layers, so that flag alone must not post again on the full batch.
+        collapsed = False
+        if self.decoder_replay_layers is not None:
+            hidden_states, pre_mix, *late_aux = self.decoder_replay_layers(
                 hidden_states,
                 positions,
                 input_ids,
@@ -671,20 +767,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
                 post_mix,
                 res_mix,
                 residual,
-                engram_hashes,
-                engram_mask,
             )
-            if idx + 1 in self.aux_hidden_state_layers:
-                # Reconstruct the aux hidden state for draft models
-                aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix)
-                aux_hidden_state = aux_recon.mean(dim=1)
-                if self.use_sequence_parallel:
-                    aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
-                aux_hidden_states.append(aux_hidden_state)
-                final_aux_recon = aux_recon
-        if layer is not None:
-            # Reuse if the last layer was captured as an aux hidden state
-            if self.end_layer in self.aux_hidden_state_layers:
+            aux_hidden_states.extend(late_aux)
+            final_aux_recon = None
+            collapsed = True
+        if ran_layers and not collapsed:
+            if (
+                self.end_layer in self.aux_hidden_state_layers
+                and final_aux_recon is not None
+            ):
                 hidden_states = final_aux_recon
             else:
                 hidden_states = self.mhc_post(
@@ -714,6 +805,162 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             return hidden_states, aux_hidden_states
         return hidden_states
 
+    def _run_layers(
+        self,
+        layer_ids: range,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        input_ids: torch.Tensor | None,
+        pre_mix: torch.Tensor | None,
+        post_mix: torch.Tensor | None,
+        res_mix: torch.Tensor | None,
+        residual: torch.Tensor | None,
+        engram_hashes: torch.Tensor | None,
+        engram_mask: torch.Tensor | None,
+        full_num_tokens: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        list[torch.Tensor],
+        torch.Tensor | None,
+        bool,
+    ]:
+        aux_hidden_states: list[torch.Tensor] = []
+        final_aux_recon: torch.Tensor | None = None
+        ran_layers = False
+        for idx in layer_ids:
+            ran_layers = True
+            hidden_states, residual, post_mix, res_mix, pre_mix = self.layers[idx](
+                hidden_states,
+                positions,
+                input_ids,
+                pre_mix,
+                post_mix,
+                res_mix,
+                residual,
+                engram_hashes,
+                engram_mask,
+            )
+            # The aux id of the first replay layer is the post of the last
+            # full-batch layer. Capture it on the trimmed rows instead.
+            if idx + 1 in self.aux_hidden_state_layers and not (
+                self.decoder_replay_layers is not None
+                and idx + 1 == self.decoder_replay_start
+            ):
+                aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix)
+                aux_hidden_state = aux_recon.mean(dim=1)
+                if self.use_sequence_parallel:
+                    aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
+                aux_hidden_states.append(aux_hidden_state)
+                final_aux_recon = aux_recon
+        return (
+            hidden_states,
+            residual,
+            post_mix,
+            res_mix,
+            pre_mix,
+            aux_hidden_states,
+            final_aux_recon,
+            ran_layers,
+        )
+
+    def _run_replay_layers(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        input_ids: torch.Tensor | None,
+        pre_mix: torch.Tensor,
+        post_mix: torch.Tensor,
+        res_mix: torch.Tensor,
+        residual: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        """Layers past the last KV source, on the rows they are given.
+
+        The final ``mhc_post`` runs here, on this batch. The return is
+        ``(hidden_states, pre_mix, *aux)``.
+        """
+        # Post of the layer before the cut. These rows are already trimmed,
+        # and sequence parallel is unsupported while replay is on.
+        cut_aux: list[torch.Tensor] = []
+        if self.decoder_replay_start in self.aux_hidden_state_layers:
+            cut_aux.append(
+                self.mhc_post(hidden_states, residual, post_mix, res_mix).mean(dim=1)
+            )
+        (
+            hidden_states,
+            residual,
+            post_mix,
+            res_mix,
+            pre_mix,
+            aux,
+            final_aux_recon,
+            ran_layers,
+        ) = self._run_layers(
+            range(self.decoder_replay_start, self.end_layer),
+            hidden_states,
+            positions,
+            input_ids,
+            pre_mix,
+            post_mix,
+            res_mix,
+            residual,
+            None,
+            None,
+            positions.shape[0],
+        )
+        assert residual is not None and post_mix is not None
+        assert res_mix is not None and pre_mix is not None
+        if (
+            self.end_layer in self.aux_hidden_state_layers
+            and final_aux_recon is not None
+        ):
+            hidden_states = final_aux_recon
+        elif ran_layers:
+            hidden_states = self.mhc_post(hidden_states, residual, post_mix, res_mix)
+        return (hidden_states, pre_mix, *cut_aux, *aux)
+
+    def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool:
+        """Whether this rank may replay from layer ``cut`` on; 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 at or 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)
@@ -1128,6 +1375,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/amd/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,
@@ -299,6 +300,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,

vllm/models/deepseek_v41/attention.py

@@ -547,11 +547,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             swa_bounded_replay = False
         if swa_bounded_replay and current_platform.is_rocm():
             logger.warning_once(
-                "SWA bounded replay is off on ROCm (the sparse SWA metadata "
-                "builders forward replay_start, but the window clamp it relies "
-                "on lives in the FlashInfer and FlashMLA prefill kernels, so "
-                "the padded slots fault); the sliding-window cache takes part "
-                "in prefix caching instead."
+                "Prefix-cache SWA replay is off on ROCm (the sparse SWA "
+                "metadata builders forward replay_start, but the window clamp "
+                "it relies on lives in the FlashInfer and FlashMLA prefill "
+                "kernels, so the padded slots fault); the sliding-window cache "
+                "takes part in prefix caching instead."
             )
             swa_bounded_replay = False
         self.swa_cache_layer = DeepseekV4SWACache(

vllm/models/deepseek_v41/nvidia/model_state.py

@@ -337,7 +337,7 @@ class DeepseekV41ModelState(DefaultModelState):
             self._warm_up_replay_kernels(input_batch, slot_mappings)
         window, cacheable_groups = self._replay
         replay_start: torch.Tensor | None = None
-        if window:
+        if window or self.decoder_replay_layers is not None:
             num_reqs = input_batch.num_reqs
             # Decode rows sit above the hit, so only prefills carry a replay
             # start; dummy batches (captures, profiling) carry none.
@@ -349,7 +349,7 @@ class DeepseekV41ModelState(DefaultModelState):
             replay_start = self._replay_start_staging.copy_to_gpu(
                 replay_start_np, out=self._replay_start[:num_reqs]
             )
-            if replay_start_np.any():
+            if window and replay_start_np.any():
                 # The replayed tokens rebuild window KV only: their slots in the
                 # prefix-cacheable groups are padded so the cached KV stays as is.
                 _pad_replayed_slots_kernel[(num_reqs,)](
@@ -397,21 +397,21 @@ class DeepseekV41ModelState(DefaultModelState):
         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,),
-        )
+        # Pad warmup is prefix-cache SWA. A zero window must not skip the gather warmup.
+        if window:
+            _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