[ROCm][MLA] Enable sparse MLA Gluon kernel from Aiter (#53492)

Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com>
Signed-off-by: Mehmet Cagri <mehmet.kaymak@amd.com>

tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py

@@ -0,0 +1,326 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="ROCm-only aiter path"
+)
+
+
+def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
+    ends = indptr.tolist()
+    return [indices[s:e].tolist() for s, e in zip(ends, ends[1:])]
+
+
+def test_triton_sparse_mla_gate(monkeypatch) -> None:
+    import vllm._aiter_ops as aiter_ops
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.platforms.rocm import on_gfx950
+
+    warnings: list[str] = []
+    monkeypatch.setattr(
+        aiter_ops.logger,
+        "warning_once",
+        lambda msg, *args, **kwargs: warnings.append(msg % args),
+    )
+
+    def enabled(aiter: bool, flag: bool) -> bool:
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER", str(int(aiter)))
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", str(int(flag)))
+        rocm_aiter_ops.refresh_env_variables()
+        return bool(rocm_aiter_ops.is_triton_sparse_mla_enabled())
+
+    try:
+        assert not enabled(aiter=True, flag=False) and not warnings
+        assert not enabled(aiter=False, flag=True)
+        assert "VLLM_ROCM_USE_AITER is off" in warnings.pop()
+        # On gfx950 the flag takes effect; elsewhere it warns and falls back.
+        assert enabled(aiter=True, flag=True) == on_gfx950()
+        if not on_gfx950():
+            assert "gfx950" in warnings.pop()
+    finally:
+        monkeypatch.undo()
+        rocm_aiter_ops.refresh_env_variables()
+
+
+@pytest.mark.parametrize(
+    "rope_dim",
+    [0, 64],
+    ids=["rope_free", "appended_rope"],  # GLM-5.3-Flash; GLM-5.1/5.2, V3.2
+)
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+@torch.inference_mode()
+def test_forward_mqa_prepares_triton_sparse_mla_inputs(
+    monkeypatch, rope_dim: int, kv_cache_dtype: str
+) -> None:
+    """What forward_mqa hands aiter: global slots built from the request-local
+    top-k rows, q quantized with the layer's scale, and the cache as stored."""
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
+        ROCMAiterMLASparseImpl,
+    )
+
+    device = torch.device("cuda")
+    gen = torch.Generator().manual_seed(0)
+    kv_lora_rank, num_heads, topk, block_size = 512, 16, 128, 16
+    head_size = kv_lora_rank + rope_dim
+    fp8 = kv_cache_dtype == "fp8"
+
+    # A decode row, then a short prefill.
+    query_lens, seq_lens = [1, 5], [40, 20]
+    num_tokens = sum(query_lens)
+    block_table = torch.randperm(8, generator=gen).view(2, 4).to(device, torch.int32)
+    topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32)
+    req_id_per_token: list[int] = []
+    rows: list[list[int]] = []
+    for req, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):
+        for pos in range(seq_len - query_len, seq_len):
+            local = torch.randperm(pos + 1, generator=gen)
+            topk_indices[len(rows), : pos + 1] = local.int()
+            slots = block_table[req, local // block_size].cpu() * block_size
+            rows.append((slots + local % block_size).tolist())
+            req_id_per_token.append(req)
+    indptr = torch.tensor([0] + [len(row) for row in rows]).cumsum(0)
+
+    impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl)
+    impl.num_heads = num_heads
+    impl.kv_lora_rank = kv_lora_rank
+    impl.qk_rope_head_dim = rope_dim
+    impl.scale = head_size**-0.5
+    impl.sinks = torch.randn(num_heads, generator=gen).to(device)
+    impl.kv_cache_dtype = kv_cache_dtype
+    impl.topk_indices_buffer = topk_indices.to(device)
+    impl.use_aiter_sparse_mla = True
+    layer = SimpleNamespace(
+        _q_scale=torch.tensor([0.02], device=device),
+        _k_scale=torch.tensor([0.01], device=device),
+    )
+    metadata = SimpleNamespace(
+        num_actual_tokens=num_tokens,
+        topk_tokens=topk,
+        req_id_per_token=torch.tensor(req_id_per_token, dtype=torch.int32).to(device),
+        block_table=block_table,
+        block_size=block_size,
+        paged_kv_indptr=indptr.to(device, torch.int32),
+        paged_kv_indices=torch.zeros(num_tokens * topk, dtype=torch.int32).to(device),
+        attn_out_dtype=torch.bfloat16,
+    )
+    cache_dtype = current_platform.fp8_dtype() if fp8 else torch.bfloat16
+    kv_cache = torch.zeros(8, block_size, head_size, device=device).to(cache_dtype)
+    # Two rows of cudagraph padding past the real tokens.
+    q = torch.randn(num_tokens + 2, num_heads, head_size, generator=gen).to(
+        device, torch.bfloat16
+    )
+    calls = []
+    monkeypatch.setattr(
+        rocm_aiter_ops, "triton_sparse_mla_fwd", lambda *a, **k: calls.append((a, k))
+    )
+
+    out, lse = impl.forward_mqa(q, kv_cache, metadata, layer)
+
+    assert len(calls) == 1
+    (q_in, kv_in, o, sm_scale, kv_indptr, kv_indices), kwargs = calls[0]
+    assert _rows_from_ragged(kv_indices, kv_indptr) == rows
+    assert q_in.shape[0] == num_tokens
+    if fp8:
+        assert q_in.dtype == cache_dtype
+        torch.testing.assert_close(
+            q_in.float() * layer._q_scale, q[:num_tokens].float(), atol=1e-3, rtol=0.07
+        )
+    else:
+        assert q_in.data_ptr() == q.data_ptr()
+    assert kv_in.shape == (8 * block_size, 1, 1, head_size)
+    assert kv_in.data_ptr() == kv_cache.data_ptr()
+    assert o.data_ptr() == out.data_ptr() and o.shape[0] == num_tokens
+    assert sm_scale == impl.scale
+    assert kwargs["kv_lora_rank"] == kv_lora_rank
+    assert kwargs["qk_rope_head_dim"] == rope_dim
+    assert kwargs["q_scale"] is layer._q_scale
+    assert kwargs["kv_scale"] is layer._k_scale
+    assert kwargs["attn_sink"] is impl.sinks
+    assert lse is None
+
+
+@pytest.mark.parametrize(
+    ("model", "compress_ratio"),
+    [("deepseek_v41", 1), ("deepseek_v41", 2), ("deepseek_v4", 4)],
+)
+@torch.inference_mode()
+def test_paged_prefill_indices_match_gathered_prefill(
+    monkeypatch, model: str, compress_ratio: int
+) -> None:
+    """The aiter prefill reads the paged caches through the SWA builder's rows
+    and build_prefill_topk_ragged_indices, and must attend exactly the keys the
+    dequantize-and-gather prefill does."""
+    import importlib
+
+    from tests.v1.attention.utils import create_vllm_config
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+        build_prefill_topk_ragged_indices,
+    )
+    from vllm.v1.kv_cache_interface import SlidingWindowMLASpec
+
+    rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm")
+    device = torch.device("cuda")
+    gen = torch.Generator().manual_seed(compress_ratio)
+    window, width, block_size, max_model_len = 16, 32, 16, 128
+    comp_block = block_size // compress_ratio
+    num_compressed = max_model_len // compress_ratio
+
+    # A decode row, then prefills with and without a prefix.
+    prefix_lens, query_lens = [20, 0, 13, 40], [1, 7, 20, 9]
+    seq_lens = [p + q for p, q in zip(prefix_lens, query_lens)]
+    num_decodes = num_decode_tokens = 1
+    num_tokens = sum(query_lens)
+    qsl = torch.tensor([0] + query_lens).cumsum(0).to(torch.int32)
+    positions = torch.cat([torch.arange(p, s) for p, s in zip(prefix_lens, seq_lens)])
+    token_to_req = torch.repeat_interleave(
+        torch.arange(len(query_lens)), torch.tensor(query_lens)
+    )
+    swa_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32)
+    comp_bt = torch.randperm(32, generator=gen)[:16].view(4, 4).to(device, torch.int32)
+
+    vllm_config = create_vllm_config(
+        model_name="facebook/opt-125m",
+        max_model_len=max_model_len,
+        block_size=block_size,
+        max_num_seqs=4,
+        max_num_batched_tokens=64,
+        hf_config_override={
+            "compress_ratios": [compress_ratio],
+            "index_topk": width,
+            "sliding_window": window,
+        },
+    )
+    swa_spec = SlidingWindowMLASpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.bfloat16,
+        sliding_window=window,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+    )
+    slots = swa_bt.cpu()[token_to_req, positions // block_size] * block_size + (
+        positions % block_size
+    )
+    common = CommonAttentionMetadata(
+        query_start_loc=qsl.to(device),
+        query_start_loc_cpu=qsl,
+        seq_lens=torch.tensor(seq_lens, dtype=torch.int32).to(device),
+        seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32),
+        num_reqs=len(query_lens),
+        num_actual_tokens=num_tokens,
+        max_query_len=max(query_lens),
+        max_seq_len=max(seq_lens),
+        block_table_tensor=swa_bt,
+        slot_mapping=slots.to(device, torch.int64),
+        positions=positions.to(device),
+        causal=True,
+    )
+
+    def build(enabled: bool):
+        monkeypatch.setattr(
+            rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled
+        )
+        return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
+            swa_spec, ["swa"], vllm_config, device
+        ).build(0, common)
+
+    assert build(enabled=False).prefill_swa_ragged_indices is None
+    md = build(enabled=True)
+    assert md.num_decode_tokens == num_decode_tokens
+
+    # Local top-k rows. Under each row's causal cap sit a failed candidate and
+    # one past the pool; past the cap, stale values neither path may read.
+    num_prefill_tokens = num_tokens - num_decode_tokens
+    topk = torch.full((num_prefill_tokens, width), -1, dtype=torch.int32)
+    for row in range(num_prefill_tokens):
+        visible = (int(positions[row + num_decode_tokens]) + 1) // compress_ratio
+        cap = min(visible, width)
+        topk[row, :cap] = torch.randperm(visible, generator=gen)[:cap].int()
+        if cap >= 2:
+            topk[row, :2] = torch.tensor([-1, num_compressed + 3])
+        topk[row, cap:] = torch.randint(
+            0, num_compressed, (width - cap,), generator=gen
+        ).int()
+    topk = topk.to(device)
+
+    topk_indices, topk_indptr = build_prefill_topk_ragged_indices(
+        topk,
+        md.token_to_req_indices,
+        md.query_start_loc,
+        md.seq_lens,
+        md.is_valid_token,
+        comp_bt,
+        block_size=comp_block,
+        compress_ratio=compress_ratio,
+        num_compressed=num_compressed,
+        token_offset=num_decode_tokens,
+    )
+    paged_keys = [
+        sorted([("s", x) for x in swa if x >= 0] + [("c", x) for x in top if x >= 0])
+        for swa, top in zip(
+            _rows_from_ragged(
+                md.prefill_swa_ragged_indices, md.prefill_swa_ragged_indptr
+            ),
+            _rows_from_ragged(topk_indices, topk_indptr),
+        )
+    ]
+
+    # The gathered path's combine, mapped from workspace offsets to cache slots.
+    M = num_compressed + window + num_tokens
+    combined, combined_lens = rocm.combine_topk_swa_indices(
+        topk,
+        md.query_start_loc[num_decodes:],
+        md.prefill_seq_lens,
+        md.prefill_gather_lens,
+        window,
+        compress_ratio,
+        width,
+        M,
+        num_compressed,
+    )
+    gather_start = (md.prefill_seq_lens - md.prefill_gather_lens).tolist()
+    swa_bt_cpu, comp_bt_cpu = swa_bt.cpu(), comp_bt.cpu()
+
+    def slot(block_table, req, pos, size):
+        return int(block_table[req, pos // size]) * size + pos % size
+
+    gathered_keys = []
+    for row, (values, n) in enumerate(zip(combined.tolist(), combined_lens.tolist())):
+        req = int(token_to_req[row + num_decode_tokens])
+        chunk_req = req - num_decodes
+        keys = []
+        for value in values[:n]:
+            if value < 0:
+                continue
+            offset = value - chunk_req * M
+            if offset < num_compressed:
+                keys.append(("c", slot(comp_bt_cpu, req, offset, comp_block)))
+            else:
+                pos = gather_start[chunk_req] + offset - num_compressed
+                keys.append(("s", slot(swa_bt_cpu, req, pos, block_size)))
+        gathered_keys.append(sorted(keys))
+
+    assert paged_keys == gathered_keys
+
+
+@pytest.mark.parametrize("model", ["deepseek_v4", "deepseek_v41"])
+def test_aiter_sparse_mla_stays_off_with_kv_connector(monkeypatch, model) -> None:
+    import importlib
+
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm")
+    monkeypatch.setattr(rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: True)
+    assert rocm._aiter_sparse_mla_enabled(has_kv_transfer=False)
+    assert not rocm._aiter_sparse_mla_enabled(has_kv_transfer=True)

vllm/_aiter_ops.py

@@ -1927,6 +1927,7 @@ class rocm_aiter_ops:
         VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
         VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
         VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
+        VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950).
         VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
         VLLM_ROCM_USE_AITER_MOE_SITUV2: SiTUv2 FlyDSL MoE activation
             dtype (a16w4 | a8w4 | a4w4).
@@ -1998,6 +1999,7 @@ class rocm_aiter_ops:
     _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
     # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
     _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
+    _TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
     _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
     _MOE_SITUV2 = _resolve_situv2_activation()
     # TODO: Consolidate under _LINEAR_ENABLED
@@ -2027,6 +2029,7 @@ class rocm_aiter_ops:
         cls._LINEAR_HIPBMM_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
         cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
         cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
+        cls._TRITON_SPARSE_MLA = envs.VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA
         cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
         cls._MOE_SITUV2 = _resolve_situv2_activation()
         _sync_aiter_situv2_moe_env()
@@ -2391,6 +2394,25 @@ class rocm_aiter_ops:
     def is_triton_rotary_embed_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED
 
+    @classmethod
+    @if_aiter_supported
+    def is_triton_sparse_mla_enabled(cls) -> bool:
+        if not cls._TRITON_SPARSE_MLA:
+            return False
+        from vllm.platforms.rocm import on_gfx950
+
+        if not cls._AITER_ENABLED:
+            reason = "VLLM_ROCM_USE_AITER is off"
+        elif not on_gfx950():
+            reason = "the kernel is gfx950-only"
+        else:
+            return True
+        logger.warning_once(
+            "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; ignoring it.",
+            reason,
+        )
+        return False
+
     @classmethod
     @if_aiter_supported
     def is_triton_gemm_enabled(cls) -> bool:
@@ -3137,6 +3159,57 @@ class rocm_aiter_ops:
     ) -> tuple[torch.Tensor, torch.Tensor]:
         return torch.ops.vllm.rocm_aiter_fused_topk(x, router_logits, top_k, gate_up)
 
+    @staticmethod
+    def triton_sparse_mla_fwd(
+        q: torch.Tensor,
+        kv_buffer: torch.Tensor,
+        o: torch.Tensor,
+        sm_scale: float,
+        kv_indptr: torch.Tensor,
+        kv_indices: torch.Tensor,
+        kv_lora_rank: int = 512,
+        qk_rope_head_dim: int = 64,
+        q_scale: torch.Tensor | None = None,
+        kv_scale: torch.Tensor | None = None,
+        attn_sink: torch.Tensor | None = None,
+        extra_kv_buffer: torch.Tensor | None = None,
+        extra_kv_indptr: torch.Tensor | None = None,
+        extra_kv_indices: torch.Tensor | None = None,
+        has_invalid: bool = True,
+    ) -> None:
+        """Sparse MLA read straight from the KV cache, for prefill and decode.
+
+        kv_indices are ragged global slot ids. With has_invalid they may hold
+        -1, which the kernel masks; a caller whose stream never holds a
+        negative slot passes False, which lets the fp8 kernel stage K straight
+        into LDS. The cache format (bf16, per-tensor fp8, or DeepSeek V4's
+        paged fp8_ds_mla) is inferred from kv_buffer and kv_scale. An fp8 q
+        must come with its q_scale and runs both dots in fp8. The extra segment
+        is DeepSeek V4's second cache: SWA window in kv_buffer, top-k
+        compressed tokens here.
+        """
+        from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd
+
+        fp8_q = q.dtype == FP8_DTYPE
+        sparse_mla_fwd(
+            q,
+            kv_buffer,
+            kv_indptr,
+            kv_indices,
+            sm_scale,
+            kv_scale=kv_scale,
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            has_invalid=has_invalid,
+            dot_precision="fp8" if fp8_q else "bf16",
+            q_scale=q_scale if fp8_q else None,
+            out=o,
+            attn_sink=attn_sink,
+            extra_kv=extra_kv_buffer,
+            extra_indptr=extra_kv_indptr,
+            extra_indices=extra_kv_indices,
+        )
+
     @staticmethod
     def mla_decode_fwd(
         q: torch.Tensor,

vllm/envs.py

@@ -148,6 +148,7 @@ if TYPE_CHECKING:
     VLLM_ROCM_AITER_MLA_DCP_VERIFY: Literal["asm", "segmented"] = "segmented"
     VLLM_ROCM_USE_AITER_MHA: bool = True
     VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
+    VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA: bool = False
     VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False
     VLLM_ROCM_USE_AITER_FP8BMM: bool = True
     VLLM_ROCM_USE_AITER_FP4BMM: bool = True
@@ -1343,6 +1344,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
     ),
+    # Whether sparse MLA prefill and decode run on aiter's Triton kernel, which
+    # reads the KV cache as stored (bf16 or fp8, paged or flat). Used by the
+    # ROCM_AITER_MLA_SPARSE backend (DeepSeek V3.2, GLM-5.x) and DeepSeek
+    # V4 / V4.1. gfx950 only. By default is disabled.
+    "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA": lambda: (
+        os.getenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "False").lower()
+        in ("true", "1")
+    ),
     # Whether to use aiter rope.
     # By default is disabled.
     "VLLM_ROCM_USE_AITER_TRITON_ROPE": lambda: (

vllm/models/deepseek_v4/amd/rocm.py

@@ -35,6 +35,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
     DeepseekSparseSWAMetadataBuilder,
 )
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+    build_prefill_topk_ragged_indices,
     build_ragged_indices_from_dense,
     rocm_inv_rope_einsum,
     rocm_inverse_rope_rows_,
@@ -103,6 +104,22 @@ def _trust_dsv4_extra_cache_nan_free(
     )
 
 
+def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool:
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    if not rocm_aiter_ops.is_triton_sparse_mla_enabled():
+        return False
+    # The aiter kernel does not scrub NaNs, so like the trust check above it
+    # needs a cache only the canonical writer has touched.
+    if has_kv_transfer:
+        logger.warning_once(
+            "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV "
+            "connector; DeepSeek V4 keeps the Triton sparse attention."
+        )
+        return False
+    return True
+
+
 def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
     lengths = lengths.to(dtype=torch.int32).contiguous()
     indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
@@ -446,6 +463,9 @@ class DeepseekV4ROCMAiterMLASparseMetadata(DeepseekV4FlashMLAMetadata):
 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
     decode_swa_ragged_indices: torch.Tensor | None = None
     decode_swa_ragged_indptr: torch.Tensor | None = None
+    # Built only for the aiter prefill, which reads the SWA cache in place.
+    prefill_swa_ragged_indices: torch.Tensor | None = None
+    prefill_swa_ragged_indptr: torch.Tensor | None = None
 
 
 class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
@@ -565,6 +585,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
             dtype=torch.int32,
             device=self.device,
         )
+        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(
+            self.vllm_config.kv_transfer_config is not None
+        )
 
     def build(
         self,
@@ -580,6 +603,23 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
             replay_start=replay_start,
         )
 
+        prefill_ragged_indices = None
+        prefill_ragged_indptr = None
+        if (
+            self._use_aiter_sparse_mla
+            and base.num_prefill_tokens > 0
+            and base.prefill_swa_indices is not None
+            and base.prefill_swa_lens is not None
+        ):
+            # Paged-coordinate rows (image visibility included), so every
+            # layer's prefill shares one build.
+            prefill_ragged_indices, prefill_ragged_indptr = (
+                build_ragged_indices_from_dense(
+                    base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1),
+                    base.prefill_swa_lens,
+                )
+            )
+
         ragged_indices = None
         ragged_indptr = None
         if (
@@ -606,6 +646,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild
             **vars(base),
             decode_swa_ragged_indices=ragged_indices,
             decode_swa_ragged_indptr=ragged_indptr,
+            prefill_swa_ragged_indices=prefill_ragged_indices,
+            prefill_swa_ragged_indptr=prefill_ragged_indptr,
         )
 
 
@@ -623,6 +665,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
     """ROCm sparse MLA attention layer for DeepSeek V4."""
 
     backend_cls = DeepseekV4ROCMAiterMLASparseBackend
+    _use_aiter_sparse_mla = False
 
     def __init__(self, *args, **kwargs):
         vllm_config = args[0] if args else kwargs["vllm_config"]
@@ -648,6 +691,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
         else:
             # Disable indexer inner overlap.
             self.indexer.aux_stream = None
+        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)
 
     def _run_sequential_pipeline(
         self,
@@ -1190,18 +1234,20 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
         if attn_metadata is None:
             # Warmup dummy run: no real metadata. Reserve the same bf16
             # gather workspace _forward_prefill would; the dequantize / topk
-            # / sparse_fwd kernels are skipped this step.
-            swa_only = self.compress_ratio <= 1
-            N = (
-                0
-                if swa_only
-                else (self.max_model_len + self.compress_ratio - 1)
-                // self.compress_ratio
-            )
-            M = N + self.window_size + self.max_num_batched_tokens
-            current_workspace_manager().get_simultaneous(
-                ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
-            )
+            # / sparse_fwd kernels are skipped this step. The aiter prefill
+            # reads the caches in place and never asks for it.
+            if not self._use_aiter_sparse_mla:
+                swa_only = self.compress_ratio <= 1
+                N = (
+                    0
+                    if swa_only
+                    else (self.max_model_len + self.compress_ratio - 1)
+                    // self.compress_ratio
+                )
+                M = N + self.window_size + self.max_num_batched_tokens
+                current_workspace_manager().get_simultaneous(
+                    ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
+                )
             output.zero_()
             return
 
@@ -1312,6 +1358,22 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
                 topk_ragged_indices = attn_metadata.c128a_decode_topk_ragged_indices
                 topk_ragged_indptr = attn_metadata.c128a_decode_topk_ragged_indptr
 
+        if self._use_aiter_sparse_mla:
+            assert swa_metadata.decode_swa_ragged_indices is not None
+            assert swa_metadata.decode_swa_ragged_indptr is not None
+            self._aiter_sparse_mla(
+                q,
+                output,
+                self.swa_cache_layer.kv_cache,
+                swa_metadata.decode_swa_ragged_indices,
+                swa_metadata.decode_swa_ragged_indptr,
+                kv_cache,
+                topk_ragged_indices,
+                topk_ragged_indptr,
+            )
+            # The aiter kernel leaves the inverse RoPE to the caller.
+            return 0
+
         return rocm_sparse_attn_decode(
             q=q,
             kv_cache=kv_cache,
@@ -1351,6 +1413,17 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
         attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
         swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
     ) -> None:
+        if self._use_aiter_sparse_mla:
+            self._forward_prefill_aiter(
+                q=q,
+                compressed_k_cache=compressed_k_cache,
+                swa_k_cache=swa_k_cache,
+                output=output,
+                attn_metadata=attn_metadata,
+                swa_metadata=swa_metadata,
+            )
+            return
+
         swa_only = attn_metadata is None
 
         num_prefills = swa_metadata.num_prefills
@@ -1476,3 +1549,90 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
                 attn_sink=self.attn_sink,
                 output=output[query_start:query_end],
             )
+
+    def _forward_prefill_aiter(
+        self,
+        q: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: DeepseekV4ROCMAiterMLASparseMetadata | None,
+        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+    ) -> None:
+        """Prefill straight off the paged caches, with no bf16 gather."""
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        assert swa_metadata.prefill_swa_ragged_indices is not None
+        assert swa_metadata.prefill_swa_ragged_indptr is not None
+
+        topk_ragged_indices = None
+        topk_ragged_indptr = None
+        if attn_metadata is None:
+            compressed_k_cache = None
+        else:
+            assert compressed_k_cache is not None
+            assert swa_metadata.token_to_req_indices is not None
+            assert swa_metadata.query_start_loc is not None
+            assert swa_metadata.seq_lens is not None
+            assert swa_metadata.is_valid_token is not None
+            if self.compress_ratio == 4:
+                assert self.topk_indices_buffer is not None
+                topk_indices = self.topk_indices_buffer[
+                    num_decode_tokens : num_decode_tokens + num_prefill_tokens
+                ]
+            else:
+                topk_indices = attn_metadata.c128a_prefill_topk_indices
+            assert topk_indices is not None
+            topk_ragged_indices, topk_ragged_indptr = build_prefill_topk_ragged_indices(
+                topk_indices,
+                swa_metadata.token_to_req_indices,
+                swa_metadata.query_start_loc,
+                swa_metadata.seq_lens,
+                swa_metadata.is_valid_token,
+                attn_metadata.block_table,
+                block_size=attn_metadata.block_size // self.compress_ratio,
+                compress_ratio=self.compress_ratio,
+                num_compressed=(self.max_model_len + self.compress_ratio - 1)
+                // self.compress_ratio,
+                token_offset=num_decode_tokens,
+                num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
+            )
+
+        self._aiter_sparse_mla(
+            q[:num_prefill_tokens],
+            output[:num_prefill_tokens],
+            swa_k_cache,
+            swa_metadata.prefill_swa_ragged_indices,
+            swa_metadata.prefill_swa_ragged_indptr,
+            compressed_k_cache,
+            topk_ragged_indices,
+            topk_ragged_indptr,
+        )
+
+    def _aiter_sparse_mla(
+        self,
+        q: torch.Tensor,
+        output: torch.Tensor,
+        swa_k_cache: torch.Tensor,
+        swa_indices: torch.Tensor,
+        swa_indptr: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        topk_indices: torch.Tensor | None,
+        topk_indptr: torch.Tensor | None,
+    ) -> None:
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        # Both fp8 caches are read in place: the SWA window as the main segment,
+        # the top-k compressed tokens as the extra one.
+        rocm_aiter_ops.triton_sparse_mla_fwd(
+            q,
+            swa_k_cache,
+            output,
+            self.scale,
+            swa_indptr,
+            swa_indices,
+            attn_sink=self.attn_sink[: q.shape[1]],
+            extra_kv_buffer=compressed_k_cache,
+            extra_kv_indptr=topk_indptr,
+            extra_kv_indices=topk_indices,
+        )

vllm/models/deepseek_v41/amd/rocm.py

@@ -38,6 +38,7 @@ from vllm.v1.attention.backends.mla.sparse_swa import (
     DeepseekSparseSWAMetadata,
 )
 from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
+    build_prefill_topk_ragged_indices,
     build_ragged_indices_from_dense,
     rocm_inv_rope_einsum,
     rocm_inverse_rope_mxfp8_rows,
@@ -64,6 +65,22 @@ def _trust_dsv4_extra_cache_nan_free(
     )
 
 
+def _aiter_sparse_mla_enabled(has_kv_transfer: bool) -> bool:
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    if not rocm_aiter_ops.is_triton_sparse_mla_enabled():
+        return False
+    # The aiter kernel does not scrub NaNs, so like the trust check above it
+    # needs a cache only the canonical writer has touched.
+    if has_kv_transfer:
+        logger.warning_once(
+            "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is ignored with a KV "
+            "connector; DeepSeek V4 keeps the Triton sparse attention."
+        )
+        return False
+    return True
+
+
 def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
     lengths = lengths.to(dtype=torch.int32).contiguous()
     indptr = torch.zeros(lengths.shape[0] + 1, dtype=torch.int32, device=lengths.device)
@@ -408,6 +425,9 @@ _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
 class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
     decode_swa_ragged_indices: torch.Tensor | None = None
     decode_swa_ragged_indptr: torch.Tensor | None = None
+    # Built only for the aiter prefill, which reads the SWA cache in place.
+    prefill_swa_ragged_indices: torch.Tensor | None = None
+    prefill_swa_ragged_indptr: torch.Tensor | None = None
 
 
 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
@@ -432,6 +452,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
             dtype=torch.int32,
             device=self.device,
         )
+        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(
+            self.vllm_config.kv_transfer_config is not None
+        )
 
     def build(
         self,
@@ -447,6 +470,22 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
             replay_start=replay_start,
         )
 
+        prefill_ragged_indices = None
+        prefill_ragged_indptr = None
+        if (
+            self._use_aiter_sparse_mla
+            and base.num_prefill_tokens > 0
+            and base.prefill_swa_indices is not None
+            and base.prefill_swa_lens is not None
+        ):
+            # Paged-coordinate rows, so every layer's prefill shares one build.
+            prefill_ragged_indices, prefill_ragged_indptr = (
+                build_ragged_indices_from_dense(
+                    base.prefill_swa_indices.reshape(base.num_prefill_tokens, -1),
+                    base.prefill_swa_lens,
+                )
+            )
+
         ragged_indices = None
         ragged_indptr = None
         if (
@@ -473,6 +512,8 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
             **vars(base),
             decode_swa_ragged_indices=ragged_indices,
             decode_swa_ragged_indptr=ragged_indptr,
+            prefill_swa_ragged_indices=prefill_ragged_indices,
+            prefill_swa_ragged_indptr=prefill_ragged_indptr,
         )
 
 
@@ -497,6 +538,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
 
     backend_cls = DeepseekV4ROCMAiterMLASparseBackend
     swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
+    _use_aiter_sparse_mla = False
 
     def __init__(self, *args, **kwargs):
         vllm_config = args[0] if args else kwargs["vllm_config"]
@@ -518,6 +560,9 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
         # same thing. The source owns one cache per compress ratio, refreshed
         # when it runs; consumers below it read through.
         self._topk_ragged_cache: dict[int, _TopkRagged] = {}
+        self._prefill_topk_ragged_cache: dict[
+            str, tuple[torch.Tensor, torch.Tensor]
+        ] = {}
         self._index_source_prefix: str | None = None
         if self.compress_ratio > 0:
             assert self.index_source_layer_id is not None
@@ -530,6 +575,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                     "this rank; PP splits inside a v4.1 index-sharing group "
                     "are not supported."
                 )
+        self._use_aiter_sparse_mla = _aiter_sparse_mla_enabled(self._has_kv_transfer)
         if self.compressor is None and self.indexer is None:
             # Dense layers have nothing to overlap; keep the base serial path.
             self.aux_stream_list = None
@@ -915,9 +961,10 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
         attn_metadata = forward_context.attn_metadata
 
         if attn_metadata is None:
-            # Warmup dummy run: no real metadata. Reserve the same bf16
-            # gather workspace _forward_prefill would; the dequantize / topk
-            # / sparse_fwd kernels are skipped this step.
+            # Warmup dummy run: no real metadata. Reserve the workspace
+            # _forward_prefill would; the dequantize / topk / sparse_fwd kernels
+            # are skipped this step. The aiter prefill reads the caches in place,
+            # so it asks only for the bf16 rows of the MXFP8 output.
             swa_only = self.compress_ratio == 0
             N = (
                 0
@@ -926,9 +973,14 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                 // self.compress_ratio
             )
             M = N + self.window_size + self.max_num_batched_tokens
-            current_workspace_manager().get_simultaneous(
-                *self._prefill_workspace_shapes(M, self.max_num_batched_tokens, q)
+            shapes = self._prefill_workspace_shapes(
+                M,
+                self.max_num_batched_tokens,
+                q,
+                gather=not self._use_aiter_sparse_mla,
             )
+            if shapes:
+                current_workspace_manager().get_simultaneous(*shapes)
             if mxfp8_out is None:
                 output.zero_()
             else:
@@ -1087,6 +1139,38 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                 num_decode_tokens=num_decode_tokens,
             )
 
+        if self._use_aiter_sparse_mla:
+            assert swa_metadata.decode_swa_ragged_indices is not None
+            assert swa_metadata.decode_swa_ragged_indptr is not None
+            out = output
+            if output_mxfp8 is not None:
+                # The aiter kernel has no MXFP8 epilogue: write bf16 rows, then
+                # rotate and quantize them as the decode reduce would have.
+                out = self._mxfp8_bf16_rows(q.shape[0], q)
+            assert out is not None
+            self._aiter_sparse_mla(
+                q,
+                out,
+                self.swa_cache_layer.kv_cache,
+                swa_metadata.decode_swa_ragged_indices,
+                swa_metadata.decode_swa_ragged_indptr,
+                kv_cache,
+                topk_ragged_indices,
+                topk_ragged_indptr,
+            )
+            if output_mxfp8 is not None:
+                rocm_inverse_rope_mxfp8_rows(
+                    out,
+                    positions,
+                    self.rotary_emb.cos_sin_cache,
+                    self.rope_head_dim,
+                    output_mxfp8[0],
+                    output_mxfp8[1],
+                )
+                return q.shape[0]
+            # The aiter kernel leaves the inverse RoPE to the caller.
+            return 0
+
         return rocm_sparse_attn_decode(
             q=q,
             kv_cache=kv_cache,
@@ -1132,6 +1216,19 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
         With ``mxfp8_out`` the rows go to a bf16 workspace first and are then
         inverse-RoPE'd and quantized into it in one pass.
         """
+        if self._use_aiter_sparse_mla:
+            self._forward_prefill_aiter(
+                q=q,
+                positions=positions,
+                compressed_k_cache=compressed_k_cache,
+                swa_k_cache=swa_k_cache,
+                output=output,
+                attn_metadata=attn_metadata,
+                swa_metadata=swa_metadata,
+                mxfp8_out=mxfp8_out,
+            )
+            return
+
         swa_only = attn_metadata is None
 
         num_prefills = swa_metadata.num_prefills
@@ -1249,15 +1346,17 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
             )
 
     def _prefill_workspace_shapes(
-        self, M: int, num_prefill_tokens: int, q: torch.Tensor
+        self, M: int, num_prefill_tokens: int, q: torch.Tensor, gather: bool = True
     ) -> list[tuple[tuple[int, ...], torch.dtype]]:
         """The prefill gather buffer, plus bf16 rows for the MXFP8 output.
 
         One ``get_simultaneous`` call: separate calls alias the same memory.
+        ``gather=False`` leaves out the gather buffer, which the aiter kernel,
+        reading the caches in place, never needs.
         """
-        shapes: list[tuple[tuple[int, ...], torch.dtype]] = [
-            ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16)
-        ]
+        shapes: list[tuple[tuple[int, ...], torch.dtype]] = []
+        if gather:
+            shapes.append(((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16))
         if _ON_GFX950:
             shapes.append(
                 (
@@ -1266,3 +1365,136 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
                 )
             )
         return shapes
+
+    def _prefill_topk_ragged(
+        self,
+        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+        attn_metadata: DeepseekV4FlashMLAMetadata,
+        compressed_k_cache: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prefill counterpart of ``_decode_topk_ragged``: layers that share an
+        index source and a compressed cache reuse one build per step."""
+        assert swa_metadata.token_to_req_indices is not None
+        assert swa_metadata.query_start_loc is not None
+        assert swa_metadata.seq_lens is not None
+        assert swa_metadata.is_valid_token is not None
+        assert self.topk_indices_buffer is not None
+        assert self._index_source_prefix is not None
+        assert self.compressed_cache_prefix is not None
+
+        source = self._static_forward_context[self._index_source_prefix]
+        if source is self:
+            source._prefill_topk_ragged_cache = {}
+        cached = source._prefill_topk_ragged_cache.get(self.compressed_cache_prefix)
+        if cached is not None:
+            return cached
+
+        num_decode_tokens = swa_metadata.num_decode_tokens
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        built = build_prefill_topk_ragged_indices(
+            self.topk_indices_buffer[
+                num_decode_tokens : num_decode_tokens + num_prefill_tokens
+            ],
+            swa_metadata.token_to_req_indices,
+            swa_metadata.query_start_loc,
+            swa_metadata.seq_lens,
+            swa_metadata.is_valid_token,
+            attn_metadata.block_table,
+            block_size=attn_metadata.block_size // self.compress_ratio,
+            compress_ratio=self.compress_ratio,
+            num_compressed=(self.max_model_len + self.compress_ratio - 1)
+            // self.compress_ratio,
+            token_offset=num_decode_tokens,
+            num_rows=compressed_k_cache.shape[0] * compressed_k_cache.shape[1],
+        )
+        source._prefill_topk_ragged_cache[self.compressed_cache_prefix] = built
+        return built
+
+    def _forward_prefill_aiter(
+        self,
+        q: torch.Tensor,
+        positions: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        swa_k_cache: torch.Tensor,
+        output: torch.Tensor | None,
+        attn_metadata: DeepseekV4FlashMLAMetadata | None,
+        swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata,
+        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> None:
+        """Prefill straight off the paged caches, with no bf16 gather.
+
+        With ``mxfp8_out`` the kernel writes bf16 rows to the workspace, which
+        are then inverse-RoPE'd and quantized into it as ``_forward_prefill``
+        does.
+        """
+        num_prefill_tokens = swa_metadata.num_prefill_tokens
+        assert swa_metadata.prefill_swa_ragged_indices is not None
+        assert swa_metadata.prefill_swa_ragged_indptr is not None
+
+        topk_indices = None
+        topk_indptr = None
+        if attn_metadata is None:
+            compressed_k_cache = None
+        else:
+            assert compressed_k_cache is not None
+            topk_indices, topk_indptr = self._prefill_topk_ragged(
+                swa_metadata, attn_metadata, compressed_k_cache
+            )
+
+        if mxfp8_out is not None:
+            output = self._mxfp8_bf16_rows(num_prefill_tokens, q)
+        assert output is not None
+        self._aiter_sparse_mla(
+            q[:num_prefill_tokens],
+            output[:num_prefill_tokens],
+            swa_k_cache,
+            swa_metadata.prefill_swa_ragged_indices,
+            swa_metadata.prefill_swa_ragged_indptr,
+            compressed_k_cache,
+            topk_indices,
+            topk_indptr,
+        )
+        if mxfp8_out is not None:
+            rocm_inverse_rope_mxfp8_rows(
+                output[:num_prefill_tokens],
+                positions[:num_prefill_tokens],
+                self.rotary_emb.cos_sin_cache,
+                self.rope_head_dim,
+                mxfp8_out[0][:num_prefill_tokens],
+                mxfp8_out[1][:num_prefill_tokens],
+            )
+
+    def _mxfp8_bf16_rows(self, num_tokens: int, q: torch.Tensor) -> torch.Tensor:
+        """bf16 attention rows for the aiter kernel to write when the layer's
+        output is MXFP8, taken from the workspace the warmup reserved."""
+        shapes = self._prefill_workspace_shapes(0, num_tokens, q, gather=False)
+        assert shapes, "an MXFP8 attention output is gfx950-only"
+        return current_workspace_manager().get_simultaneous(*shapes)[0]
+
+    def _aiter_sparse_mla(
+        self,
+        q: torch.Tensor,
+        output: torch.Tensor,
+        swa_k_cache: torch.Tensor,
+        swa_indices: torch.Tensor,
+        swa_indptr: torch.Tensor,
+        compressed_k_cache: torch.Tensor | None,
+        topk_indices: torch.Tensor | None,
+        topk_indptr: torch.Tensor | None,
+    ) -> None:
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        # Both fp8 caches are read in place: the SWA window as the main segment,
+        # the top-k compressed tokens as the extra one.
+        rocm_aiter_ops.triton_sparse_mla_fwd(
+            q,
+            swa_k_cache,
+            output,
+            self.scale,
+            swa_indptr,
+            swa_indices,
+            attn_sink=self.attn_sink[: q.shape[1]],
+            extra_kv_buffer=compressed_k_cache,
+            extra_kv_indptr=topk_indptr,
+            extra_kv_indices=topk_indices,
+        )

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -396,9 +396,11 @@ class ROCMAiterMLASparseMetadataBuilder(
         attention_context = vllm_config.compilation_config.static_forward_context
         # Sink decode must use AITER's nonpersistent path. In particular,
         # gfx942 has no persistent+LSE kernel, and its metadata heuristic
-        # terminates for HY-V4's TP1 H64 shape.
+        # terminates for HY-V4's TP1 H64 shape. The Triton sparse MLA kernel
+        # never reads the persistent metadata either.
         self._use_persistent_metadata = all(
             getattr(attention_context[name].impl, "sinks", None) is None
+            and not getattr(attention_context[name].impl, "use_aiter_sparse_mla", False)
             for name in layer_names
         )
         # Bounds the KV-split heuristic (see `_sparse_decode_max_split`).
@@ -716,6 +718,7 @@ class ROCMAiterMLASparseImpl(
     is_sparse = True
     supports_dense_mha_prefill = False
     supports_dcp = False
+    use_aiter_sparse_mla = False
 
     def __init__(
         self,
@@ -765,6 +768,32 @@ class ROCMAiterMLASparseImpl(
         (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
             (q_concat_shape, vllm_config.model_config.dtype),
         )
+        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
+
+        if rocm_aiter_ops.is_triton_sparse_mla_enabled():
+            reason = self._aiter_sparse_mla_unsupported_reason(vllm_config)
+            if reason is None:
+                self.use_aiter_sparse_mla = True
+            else:
+                logger.warning_once(
+                    "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; "
+                    "using the default sparse MLA path instead.",
+                    reason,
+                )
+
+    def _aiter_sparse_mla_unsupported_reason(
+        self, vllm_config: VllmConfig
+    ) -> str | None:
+        model_dtype = vllm_config.model_config.dtype
+        if model_dtype != torch.bfloat16:
+            return f"the model dtype is {model_dtype}, not bfloat16"
+        cache_dtype = self.kv_cache_dtype
+        if cache_dtype not in ("auto", "bfloat16", "fp8", "fp8_e4m3"):
+            return f"kv-cache dtype {cache_dtype!r} is neither bfloat16 nor e4m3 fp8"
+        # Context parallelism merges each rank's partials with a per-token LSE.
+        if self.dcp_world_size > 1 or self.pcp_world_size > 1:
+            return "context parallelism is not supported"
+        return None
 
     def record_logical_topk_ready(self) -> None:
         # This impl shares the top-k indices buffer via SharedTopkIndicesBuffer
@@ -975,6 +1004,43 @@ class ROCMAiterMLASparseImpl(
 
         return output, lse
 
+    def _forward_mla_aiter(
+        self,
+        layer: AttentionLayer,
+        q: torch.Tensor,  # [sq, heads, d_qk], not head-padded
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: ROCMAiterMLASparseMetadata,
+    ) -> torch.Tensor:
+        """_forward_mla on aiter's Triton sparse MLA kernel.
+
+        Reads the same index stream, but needs no q head padding and no
+        persistent metadata, and its launch depends on shapes only, so it is
+        CUDA-graph capturable.
+        """
+        num_actual_tokens = attn_metadata.num_actual_tokens
+        output = torch.empty(
+            [q.shape[0], self.num_heads, self.kv_lora_rank],
+            dtype=attn_metadata.attn_out_dtype,
+            device=q.device,
+        )
+        rocm_aiter_ops.triton_sparse_mla_fwd(
+            q[:num_actual_tokens],
+            kv_c_and_k_pe_cache.view(-1, 1, 1, q.shape[-1]),
+            output[:num_actual_tokens],
+            self.scale,
+            attn_metadata.paged_kv_indptr,
+            attn_metadata.paged_kv_indices,
+            kv_lora_rank=self.kv_lora_rank,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            q_scale=layer._q_scale,
+            kv_scale=layer._k_scale,
+            attn_sink=self.sinks,
+            # triton_convert_req_index_to_global_index writes 0, never -1,
+            # for an invalid top-k entry, so no slot in the stream is negative.
+            has_invalid=False,
+        )
+        return output
+
     def forward_mqa(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -1026,6 +1092,11 @@ class ROCMAiterMLASparseImpl(
                 original_q_shape = q.shape
                 q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
                 q = q.view(original_q_shape)
+        if self.use_aiter_sparse_mla:
+            output = self._forward_mla_aiter(
+                layer, q, kv_c_and_k_pe_cache, attn_metadata
+            )
+            return output, None
         mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
         return self._forward_mla(
             layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -4306,3 +4306,99 @@ def rocm_sparse_attn_decode(
     if direct_out is None:
         output.copy_(attn_out.to(output.dtype))
     return output.shape[0] if inv_rope_positions is not None else 0
+
+
+@triton.jit(do_not_specialize=["token_offset"])
+def _prefill_topk_global_slots_kernel(
+    out_ptr,
+    lens_ptr,
+    topk_ptr,
+    topk_stride,
+    token_to_req_ptr,
+    query_start_loc_ptr,
+    seq_lens_ptr,
+    is_valid_token_ptr,
+    block_table_ptr,
+    block_table_stride,
+    token_offset,
+    num_compressed,
+    TOPK: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+):
+    row = tl.program_id(0)
+    token = row + token_offset
+    req = tl.load(token_to_req_ptr + token)
+    query_start = tl.load(query_start_loc_ptr + req)
+    query_len = tl.load(query_start_loc_ptr + req + 1) - query_start
+    pos = tl.load(seq_lens_ptr + req) - query_len + token - query_start
+    # Only (pos + 1) // ratio compressed tokens exist yet; the rest of the row
+    # can hold stale indices.
+    row_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOPK)
+    row_len = tl.where(tl.load(is_valid_token_ptr + token), row_len, 0)
+
+    for start in range(0, TOPK, BLOCK_W):
+        cols = start + tl.arange(0, BLOCK_W)
+        live = cols < row_len
+        local = tl.load(topk_ptr + row * topk_stride + cols, mask=live, other=-1)
+        valid = live & (local >= 0) & (local < num_compressed)
+        local = tl.where(valid, local, 0)
+        block = tl.load(
+            block_table_ptr + req * block_table_stride + local // BLOCK_SIZE,
+            mask=valid,
+            other=0,
+        )
+        slot = tl.where(valid, block * BLOCK_SIZE + local % BLOCK_SIZE, -1)
+        tl.store(out_ptr + row * TOPK + cols, slot, mask=cols < TOPK)
+    tl.store(lens_ptr + row, row_len)
+
+
+def build_prefill_topk_ragged_indices(
+    topk_indices: torch.Tensor,
+    token_to_req_indices: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    is_valid_token: torch.Tensor,
+    block_table: torch.Tensor,
+    block_size: int,
+    compress_ratio: int,
+    num_compressed: int,
+    token_offset: int,
+    num_rows: int = -1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Map prefill top-k rows to a ragged stream of compressed-cache slots.
+
+    ``topk_indices`` holds local compressed positions for the prefill tokens,
+    which sit at ``token_offset`` in the batch; ``token_to_req_indices``,
+    ``query_start_loc``, ``seq_lens`` and ``block_table`` are batch-wide.
+    ``block_size`` is the compressed cache's, i.e. already divided by the ratio.
+    """
+    topk_indices = topk_indices.reshape(topk_indices.shape[0], -1)
+    num_tokens, width = topk_indices.shape
+    dense = torch.empty(
+        (num_tokens, width), dtype=torch.int32, device=topk_indices.device
+    )
+    lens = torch.empty(num_tokens, dtype=torch.int32, device=topk_indices.device)
+    if num_tokens > 0 and width > 0:
+        _prefill_topk_global_slots_kernel[(num_tokens,)](
+            dense,
+            lens,
+            topk_indices,
+            topk_indices.stride(0),
+            token_to_req_indices,
+            query_start_loc,
+            seq_lens,
+            is_valid_token,
+            block_table,
+            block_table.stride(0),
+            token_offset,
+            num_compressed,
+            TOPK=width,
+            COMPRESS_RATIO=compress_ratio,
+            BLOCK_SIZE=block_size,
+            BLOCK_W=min(triton.next_power_of_2(width), 1024),
+        )
+    else:
+        lens.zero_()
+    return build_ragged_indices_from_dense(dense, lens, num_rows=num_rows)