[ROCm][Perf][DSv4.1] Mono decode layer for DeepSeek-V4.1-Flash on MI355X

.buildkite/test_areas/models_basic.yaml

@@ -91,8 +91,9 @@ steps:
       - 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
+      - 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
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
 
 - label: ":nvidia: (B200) Inkling"

tests/models/test_dsv41_mono_decode.py

@@ -0,0 +1,261 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The DeepSeek-V4.1 mono decode layer runs only where its kernels are exact:
+decode-only steps of at most MAX_ROWS rows with causal SWA windows, eager or in
+a FULL graph. A PIECEWISE capture is replayed for mixed batches, so it must
+never record the mono launches. Layers outside the whole-layer kernels run the
+FFN launch on the same steps, from wo_b's unreduced output."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.config import CUDAGraphMode
+from vllm.models.deepseek_v41.amd import mono_decode as md
+
+M = 6
+
+
+class _Runner:
+    def __init__(self):
+        self.calls = []
+
+    def forward(self, *args, **kwargs):
+        self.calls.append((args, kwargs))
+        return ("mono outputs",)
+
+    def ffn(self, *args):
+        self.calls.append((args, {}))
+        return ("ffn outputs",)
+
+
+def _layer(comp="comp"):
+    cache = lambda: torch.zeros(4, 128 * md.RECORD, dtype=torch.uint8)  # noqa: E731
+    attn = SimpleNamespace(
+        swa_cache_layer=SimpleNamespace(prefix="swa", kv_cache=cache()),
+        compressed_cache_prefix=comp,
+        topk_indices_buffer=torch.zeros(64, 512, dtype=torch.int32),
+        compress_ratio=1,
+        _compressed_kv_cache=cache,
+    )
+    return SimpleNamespace(attn=attn)
+
+
+def _swa(**overrides):
+    fields = dict(
+        num_prefills=0,
+        num_decodes=1,
+        num_decode_tokens=M,
+        decode_swa_width=md.SWA_WIDTH,
+        decode_swa_indices=torch.zeros(64, 1, md.SWA_WIDTH, dtype=torch.int32),
+        decode_swa_lens=torch.zeros(64, dtype=torch.int32),
+        token_to_req_indices=torch.zeros(64, dtype=torch.int32),
+        slot_mapping=torch.zeros(64, dtype=torch.int64),
+        block_size=32,
+    )
+    fields.update(overrides)
+    return SimpleNamespace(**fields)
+
+
+def _inputs(rows=M):
+    return dict(
+        x=torch.zeros(rows, 5120, dtype=torch.bfloat16),
+        positions=torch.zeros(rows, dtype=torch.int64),
+        residual=torch.zeros(rows, 4, 5120, dtype=torch.bfloat16),
+        post_mix=torch.zeros(rows, 4, 1),
+        res_mix=torch.zeros(rows, 4, 4),
+        pre_mix=torch.zeros(rows, 4),
+    )
+
+
+@pytest.fixture
+def run(monkeypatch):
+    """Call the mono path for one layer under a forward context built from
+    ``mode`` and ``metadata``; returns (result, runner calls)."""
+
+    def _run(
+        mode=CUDAGraphMode.FULL, metadata=None, ffn_only=False, entry=None, **inputs
+    ):
+        if metadata is None:
+            metadata = {
+                "swa": _swa(),
+                "comp": SimpleNamespace(block_size=128, block_table=None),
+            }
+        fc = SimpleNamespace(cudagraph_runtime_mode=mode, attn_metadata=metadata)
+        runner = _Runner()
+        monkeypatch.setattr(md, "is_forward_context_available", lambda: True)
+        monkeypatch.setattr(md, "get_forward_context", lambda: fc)
+        monkeypatch.setattr(md, "_mono_runner", lambda device: runner)
+        monkeypatch.setattr(md.MonoDecodeLayer, "weights", lambda self, layer: None)
+        args = {**_inputs(), **inputs}
+        mono = md.MonoDecodeLayer(ffn_only)
+        if (entry or ("ffn" if ffn_only else "layer")) == "ffn":
+            # x stands for wo_b's unreduced output; the first layer's attention
+            # has no compressed cache
+            part, _ = args.pop("x"), args.pop("positions")
+            return mono.ffn(_layer(comp=None), part, **args), runner.calls
+        return mono(_layer(), **args), runner.calls
+
+    return _run
+
+
+@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL])
+def test_decode_step_runs_mono(run, mode):
+    out, calls = run(mode)
+    assert out == ("mono outputs",) and len(calls) == 1
+    args, kwargs = calls[0]
+    # the caches reach the kernels as [blocks, block, record] bytes
+    assert args[8].shape == (4, 32, md.RECORD) and args[8].stride() == (
+        128 * md.RECORD,
+        md.RECORD,
+        1,
+    )
+    assert kwargs["comp_cache"].shape == (4, 128, md.RECORD)
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        dict(mode=CUDAGraphMode.PIECEWISE),
+        dict(
+            metadata={
+                "swa": _swa(num_prefills=1),
+                "comp": SimpleNamespace(block_size=128),
+            }
+        ),
+        dict(
+            metadata={
+                "swa": _swa(decode_swa_width=256),
+                "comp": SimpleNamespace(block_size=128),
+            }
+        ),
+        dict(
+            metadata={
+                "swa": _swa(num_decode_tokens=M + 1),
+                "comp": SimpleNamespace(block_size=128),
+            }
+        ),
+        dict(metadata={}),  # a profile / dummy run: no attention metadata
+        dict(**_inputs(md.MAX_ROWS + 6)),
+        dict(residual=None),  # the first layer's seam broadcasts the embedding
+    ],
+    ids=[
+        "piecewise",
+        "prefill",
+        "noncausal-window",
+        "padding",
+        "no-metadata",
+        "rows",
+        "first-layer",
+    ],
+)
+def test_other_steps_take_the_original_path(run, case):
+    out, calls = run(**case)
+    assert out is None and not calls
+
+
+def _decoder_layer(engram=None, **attn):
+    attn = (
+        dict(
+            layer_id=21,
+            compressor=None,
+            indexer=None,
+            compress_ratio=1,
+            kv_mxfp8=False,
+            kv_cache_dtype="fp8_ds_mla",
+        )
+        | attn
+    )
+    ffn = SimpleNamespace(
+        shared_experts=object(),
+        gate=SimpleNamespace(tid2eid=None),
+        n_routed_experts=384,
+        n_activated_experts=6,
+        routed_scaling_factor=1.5,
+        swiglu_limit=10.0,
+        scoring_func="sqrtsoftplus",
+        renormalize=True,
+    )
+    return SimpleNamespace(
+        attn=SimpleNamespace(**attn),
+        ffn=ffn,
+        engram=engram,
+        use_sequence_parallel=False,
+        fuse_seam_norm=True,
+    )
+
+
+@pytest.mark.parametrize(
+    "layer, path",
+    [
+        (_decoder_layer(), "whole"),
+        (_decoder_layer(compress_ratio=0), "ffn"),
+        (_decoder_layer(engram=object()), "ffn"),
+        (_decoder_layer(compressor=object(), indexer=object()), "ffn"),
+        (_decoder_layer(indexer=object()), "ffn"),
+        (_decoder_layer(kv_cache_dtype="auto"), "ffn"),
+        (_decoder_layer(layer_id=40), None),
+    ],
+    ids=["standard", "first", "engram", "kv-source", "index-source", "kv", "draft"],
+)
+def test_layers_take_whole_or_ffn_launch(monkeypatch, layer, path):
+    """The whole-layer kernels serve the standard layers; every other backbone
+    layer keeps vLLM's attention and runs its FFN half as one launch."""
+    import vllm.platforms.rocm as rocm
+
+    monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1")
+    monkeypatch.setattr(rocm, "get_cdna_version", lambda: 4)
+    monkeypatch.setattr(md, "get_tensor_model_parallel_world_size", lambda: 2)
+    runner = "vllm.models.deepseek_v41.amd.mono.runner"
+    monkeypatch.setitem(sys.modules, runner, ModuleType(runner))
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40))
+    )
+    mono = md.MonoDecodeLayer.create(layer, config)
+    got = None if mono is None else "ffn" if mono.ffn_only else "whole"
+    assert got == path
+
+
+@pytest.mark.parametrize("mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL])
+def test_decode_step_runs_ffn_launch(run, mode):
+    out, calls = run(mode, ffn_only=True, metadata={"swa": _swa()})
+    assert out == ("ffn outputs",) and len(calls) == 1
+    assert calls[0][0][1].shape == (M, 5120)  # wo_b's unreduced output
+
+
+def test_ffn_layer_never_runs_whole(run):
+    out, calls = run(ffn_only=True, entry="layer")
+    assert out is None and not calls
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        dict(mode=CUDAGraphMode.PIECEWISE),
+        dict(metadata={"swa": _swa(num_prefills=1)}),
+        dict(metadata={"swa": _swa(num_decode_tokens=M + 1)}),
+        dict(metadata={}),
+        dict(**_inputs(md.MAX_ROWS + 6)),
+    ],
+    ids=["piecewise", "prefill", "padding", "no-metadata", "rows"],
+)
+def test_other_steps_reduce_wo_b_themselves(run, case):
+    """None: the layer all-reduces wo_b's partial and runs the FFN as before."""
+    case.setdefault("metadata", {"swa": _swa()})
+    out, calls = run(ffn_only=True, **case)
+    assert out is None and not calls
+
+
+@pytest.mark.parametrize("cdna", [3, 5])
+def test_opt_in_outside_the_kernels_cdna_raises(monkeypatch, cdna):
+    import vllm.platforms.rocm as rocm
+
+    monkeypatch.setenv("VLLM_ROCM_MONO_DECODE", "1")
+    monkeypatch.setattr(rocm, "get_cdna_version", lambda: cdna)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40))
+    )
+    with pytest.raises(ValueError, match="needs CDNA4"):
+        md.MonoDecodeLayer.create(_decoder_layer(), config)

vllm/envs.py

@@ -155,6 +155,7 @@
     VLLM_ROCM_USE_AITER_FP4BMM: bool = True
     VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
     VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False
+    VLLM_ROCM_MONO_DECODE: bool = False
     VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True
     VLLM_ROCM_USE_SKINNY_GEMM: bool = True
     VLLM_ROCM_FP8_PADDING: bool = True
@@ -1387,6 +1388,12 @@ def _resolve_rust_cli_path() -> str | None:
         os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower()
         in ("true", "1")
     ),
+    # Run the eligible layers of decode steps on the model's mono decode
+    # kernels: persistent launches a layer, the TP all-reduces in-kernel
+    # (DeepSeek-V4.1 on CDNA4). By default is disabled.
+    "VLLM_ROCM_MONO_DECODE": lambda: (
+        os.getenv("VLLM_ROCM_MONO_DECODE", "False").lower() in ("true", "1")
+    ),
     # Whether to use aiter triton kernels for gemm ops.
     # By default is enabled.
     "VLLM_ROCM_USE_AITER_TRITON_GEMM": lambda: (

vllm/models/deepseek_v41/amd/model.py

@@ -14,6 +14,7 @@
     get_pp_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
 )
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
@@ -58,6 +59,7 @@
 from vllm.models.deepseek_v4.amd.model import (
     DeepseekV4MoE as DeepseekV4MoEBase,
 )
+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.sequence import IntermediateTensors
@@ -255,6 +257,12 @@ def __init__(
         # attn_norm / ffn_norm into its collapse, so the separate norms are
         # skipped for the seams it takes.
         self.fuse_seam_norm = HAS_AITER_MHC_FUSED_POST_PRE_DELAYED_RMS_NORM
+        # Decode steps of an eligible layer may run as the mono decode layer.
+        self.mono = MonoDecodeLayer.create(self, vllm_config)
+        if self.mono is not None and self.mono.ffn_only:
+            # A decode step hands wo_b's TP partial to the FFN launch; forward
+            # reduces it on the others.
+            self.attn.wo_b.reduce_results = False
 
     @staticmethod
     def _hc_collapse(x: torch.Tensor, pre_mix: torch.Tensor) -> torch.Tensor:
@@ -273,6 +281,10 @@ def forward(
         engram_hashes: torch.Tensor | None = None,
         engram_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.mono is not None:
+            out = self.mono(self, x, positions, residual, post_mix, res_mix, pre_mix)
+            if out is not None:
+                return out
         # Layer 0's attention seam projects the 2-D embedding with the folded
         # hc_attn_fn_broadcast instead of the 4-stream residual with hc_attn_fn.
         # The fused kernel only takes the latter, so that seam keeps the
@@ -371,6 +383,11 @@ def forward(
         x = self.attn(positions, x, None)
         if self.use_sequence_parallel:
             x = sp_reduce_scatter(x)
+        elif self.mono is not None and self.mono.ffn_only:
+            out = self.mono.ffn(self, x, residual, post_mix, res_mix, attn_pre)
+            if out is not None:
+                return out
+            x = tensor_model_parallel_all_reduce(x)
 
         residual, post_mix, res_mix, x, ffn_pre = self.mhc_pre_delayed(
             residual,