[ROCm][Perf][DSv4.1] Build the SWA decode metadata in one launch

tests/models/test_dsv41_mono_decode.py

@@ -8,6 +8,7 @@
 unreduced output."""
 
 import sys
+from collections import defaultdict
 from types import ModuleType, SimpleNamespace
 
 import pytest
@@ -732,6 +733,64 @@ def test_rocm_attention_takes_the_gfx942_indexer_with_the_flag(
     assert rocm.DeepseekV41Gfx942Indexer.attn_cls is Gfx942SparseAttnIndexer
 
 
+@pytest.mark.parametrize("monokernel", [True, False])
+@pytest.mark.parametrize("cdna", [3, 4])
+def test_swa_decode_meta_runs_on_cdna3_and_cdna4_with_the_flag(
+    monkeypatch, cdna, monokernel
+):
+    """swa_decode_meta's kernel is plain Triton, so with --monokernel a decode
+    batch gets its SWA metadata from the one launch on CDNA4 (gfx950) as on
+    CDNA3 (gfx942). Without the flag, the builder keeps vLLM's build."""
+    import vllm.platforms.rocm as rocm_platform
+    from vllm.models.deepseek_v41.amd import ops, rocm
+    from vllm.models.deepseek_v41.amd.ops import swa_decode_meta
+
+    monkeypatch.setattr(ops, "current_platform", SimpleNamespace(is_rocm=lambda: True))
+    monkeypatch.setattr(rocm_platform, "get_cdna_version", lambda: cdna)
+    monkeypatch.setattr(ops, "_monokernel", ops._monokernel)
+    ops.record_monokernel(_config(monokernel=monokernel))
+    launches = []
+    monkeypatch.setattr(
+        swa_decode_meta, "launch", lambda *args, noncausal: launches.append(noncausal)
+    )
+    ints = lambda *shape: torch.zeros(shape, dtype=torch.int32)  # noqa: E731
+    builder = SimpleNamespace(
+        decode_threshold=M,
+        window_size=md.SWA_WIDTH,
+        block_size=128,
+        decode_swa_indices=ints(64, 1, md.SWA_WIDTH),
+        decode_swa_lens=ints(64),
+        is_valid_token=torch.zeros(64, dtype=torch.bool),
+        token_to_req_indices=ints(64),
+        no_replay_start=ints(2),
+        decode_swa_ragged_indices_buffer=ints(64 * md.SWA_WIDTH),
+        decode_swa_ragged_indptr_buffer=ints(65),
+        build_tile_scheduler=lambda rows: defaultdict(lambda: None),
+        _build_deepseek_v4_metadata=lambda *args: {},
+    )
+    query_start_loc = torch.tensor([0, M], dtype=torch.int32)
+    cam = SimpleNamespace(
+        causal=True,
+        max_query_len=M,
+        num_reqs=1,
+        num_actual_tokens=M,
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc,
+        seq_lens=torch.tensor([1000], dtype=torch.int32),
+        seq_lens_cpu_upper_bound=torch.tensor([1000], dtype=torch.int32),
+        slot_mapping=torch.arange(M),
+        block_table_tensor=ints(1, 8),
+        token_to_req_indices=lambda buffer: buffer[:M],
+    )
+    out = swa_decode_meta.build(
+        builder, 0, cam, False, None, rocm.DeepseekV41ROCMAiterSparseSWAMetadata
+    )
+    # No later test may see this test's cached gate.
+    ops.record_monokernel(_config(monokernel=False))
+    assert launches == ([False] if monokernel else [])
+    assert (out is not None) is monokernel
+
+
 def test_dspark_speculator_asks_the_draft_model_first():
     """A draft model that samples a step itself returns True, and vLLM's
     loop then does not run. The stub speculator has none of the loop's

vllm/models/deepseek_v41/amd/ops/__init__.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""gfx942 (MI300X, MI325X) decode helpers for DeepSeek-V4.1.
+"""gfx942 (MI300X, MI325X) decode helpers for DeepSeek-V4.1. swa_decode_meta
+also runs on gfx950 (MI355X).
 
 ``topk`` holds the sparse indexer's top-512 hooks, built from
 ``topk512_gfx942.cu``, and ``cand_logits`` the indexer logits of the DSpark
 candidate positions. ``sparse_indexer`` is the indexer layer that runs them on
 decode steps. ``topk_pack`` and ``swa_decode_meta`` build the decode
 attention metadata, ``fused_markov`` is DSpark's Markov sampling in one
 launch, and ``draft_lm_head`` keeps the DSpark draft's 5-row logits GEMM off
-wvSplitK. All of them run only when ``enabled()`` is true.
+wvSplitK. All of them run only when ``enabled()`` is true, except
+``swa_decode_meta``, which runs when ``swa_decode_meta_enabled()`` is true.
 """
 
 import functools
@@ -28,15 +30,29 @@ def record_monokernel(vllm_config) -> None:
     global _monokernel
     _monokernel = bool(vllm_config.kernel_config.monokernel)
     enabled.cache_clear()
+    swa_decode_meta_enabled.cache_clear()
 
 
 @functools.cache
 def enabled() -> bool:
     """True with --monokernel on a gfx942 GPU. Every helper keeps the result
     of vLLM's op it replaces, so each one returns False for a call it cannot
-    take and vLLM then runs its own op. gfx950 always keeps vLLM's ops here."""
+    take and vLLM then runs its own op. gfx950 keeps vLLM's ops here, except
+    for swa_decode_meta (swa_decode_meta_enabled)."""
     if not _monokernel or not current_platform.is_rocm():
         return False
     from vllm.platforms.rocm import on_gfx942
 
     return on_gfx942()
+
+
+@functools.cache
+def swa_decode_meta_enabled() -> bool:
+    """True with --monokernel on a CDNA3 (gfx942) or CDNA4 (gfx950) GPU.
+    swa_decode_meta is plain Triton and writes the buffers that the attention
+    of both reads."""
+    if not _monokernel or not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import get_cdna_version
+
+    return get_cdna_version() in (3, 4)

vllm/models/deepseek_v41/amd/ops/swa_decode_meta.py

@@ -22,16 +22,17 @@
 ComputeDSparkNoncausalSWAIndicesKernel in place of the window kernel. That
 is 3 more builds of 8 launches after every target step.
 
-On gfx942 with --monokernel (``ops.enabled()``), decode
-batches of up to MAX_ROWS rows, causal or the draft block, use it. Other
-batches take the rest of DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build.
+On gfx942 and gfx950 with --monokernel
+(``ops.swa_decode_meta_enabled()``), decode batches of up to MAX_ROWS rows,
+causal or the draft block, use it. Other batches take the rest of
+DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build.
 tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py compares it with the vLLM ops
 that it replaces.
 """
 
 import torch
 
-from vllm.models.deepseek_v41.amd.ops import enabled
+from vllm.models.deepseek_v41.amd.ops import swa_decode_meta_enabled
 from vllm.triton_utils import tl, triton
 
 # One program holds all rows of the batch in registers, as a ROWS x WIDTH
@@ -182,7 +183,7 @@ def build(builder, common_prefix_len, cam, fast_build, replay_start, metadata_cl
     """Return the ROCm SWA metadata of a decode batch, causal or the DSpark
     draft block, built with one launch, or None when the batch needs the rest
     of DeepseekV41ROCMAiterSparseSWAMetadataBuilder.build."""
-    if not enabled():
+    if not swa_decode_meta_enabled():
         return None
     from vllm.v1.attention.backends.mla import sparse_swa
     from vllm.v1.attention.backends.utils import split_decodes_and_prefills

vllm/models/deepseek_v41/amd/rocm.py

@@ -380,7 +380,7 @@ def build(
         fast_build: bool = False,
         replay_start: torch.Tensor | None = None,
     ) -> DeepseekV41ROCMAiterSparseSWAMetadata:
-        # On gfx942 with --monokernel, a decode batch gets this
+        # On gfx942 and gfx950 with --monokernel, a decode batch gets this
         # metadata from one launch instead of 8 launches per SWA group
         # (swa_decode_meta.py). Other batches return None there and take
         # the build below.