[ROCm][Perf][GLM-5.3-Flash] Enable AITER sparse MLA on gfx942, including FP8 KV

tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py

@@ -49,18 +49,98 @@ def enabled(aiter: bool, flag: bool) -> bool:
         rocm_aiter_ops.refresh_env_variables()
 
 
+def test_triton_sparse_mla_gate_gfx942(monkeypatch) -> None:
+    """gfx942 takes the flag only for callers that opt in, and only with an aiter
+    whose kernel lists gfx942."""
+    import vllm._aiter_ops as aiter_ops
+    import vllm.platforms.rocm as rocm
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    warnings: list[str] = []
+    monkeypatch.setattr(
+        aiter_ops.logger,
+        "warning_once",
+        lambda msg, *args, **kwargs: warnings.append(msg % args),
+    )
+    monkeypatch.setattr(rocm, "on_gfx950", lambda: False)
+    monkeypatch.setattr(rocm, "on_gfx942", lambda: True)
+    archs = {"SUPPORTED_ARCHS": ("gfx942", "gfx950")}
+    monkeypatch.setattr(
+        aiter_ops, "_aiter_sparse_mla_archs", lambda name: archs.get(name, ())
+    )
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+    monkeypatch.setenv("VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA", "1")
+    rocm_aiter_ops.refresh_env_variables()
+    try:
+        assert rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True)
+        assert not warnings
+        # DeepSeek V4 does not opt in: its packed caches are gfx950-only.
+        assert not rocm_aiter_ops.is_triton_sparse_mla_enabled()
+        assert "gfx950-only" in warnings.pop()
+        # An aiter from before ROCm/aiter#5721 has no gfx942 kernel.
+        archs["SUPPORTED_ARCHS"] = ("gfx950",)
+        assert not rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True)
+        assert "ROCm/aiter#5721" in warnings.pop()
+    finally:
+        monkeypatch.undo()
+        rocm_aiter_ops.refresh_env_variables()
+
+
+@pytest.mark.parametrize(
+    ("rope_dim", "kv_cache_dtype", "reads_fp8", "reason"),
+    [
+        (0, "auto", False, None),
+        (0, "fp8", True, None),
+        (0, "fp8", False, "ROCm/aiter#6199"),
+        (64, "auto", True, "rope-free"),
+    ],
+    ids=["bf16", "fp8", "fp8_old_aiter", "rope"],
+)
+def test_aiter_sparse_mla_gfx942_scope(
+    monkeypatch, rope_dim: int, kv_cache_dtype: str, reads_fp8: bool, reason
+) -> None:
+    """On gfx942 the backend takes aiter's kernel for rope-free sparse MLA only,
+    and with an fp8 cache only if the installed aiter reads it there."""
+    import vllm.platforms.rocm as rocm
+    from vllm._aiter_ops import rocm_aiter_ops
+    from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
+        ROCMAiterMLASparseImpl,
+    )
+
+    monkeypatch.setattr(rocm, "on_gfx942", lambda: True)
+    monkeypatch.setattr(
+        rocm_aiter_ops, "triton_sparse_mla_reads_fp8_on_gfx942", lambda: reads_fp8
+    )
+    impl = ROCMAiterMLASparseImpl.__new__(ROCMAiterMLASparseImpl)
+    impl.kv_cache_dtype = kv_cache_dtype
+    impl.qk_rope_head_dim = rope_dim
+    impl.dcp_world_size = impl.pcp_world_size = 1
+    vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.bfloat16))
+
+    got = impl._aiter_sparse_mla_unsupported_reason(vllm_config)
+    if reason is None:
+        assert got is None
+    else:
+        assert got is not None and reason in got
+
+
 @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"])
+@pytest.mark.parametrize(
+    ("kv_cache_dtype", "quantizes_q"),
+    [("auto", False), ("fp8", False), ("fp8", True)],
+    ids=["bf16", "fp8", "fp8_bf16_q"],  # fp8_bf16_q: gfx942's kernel
+)
 @torch.inference_mode()
 def test_forward_mqa_prepares_triton_sparse_mla_inputs(
-    monkeypatch, rope_dim: int, kv_cache_dtype: str
+    monkeypatch, rope_dim: int, kv_cache_dtype: str, quantizes_q: bool
 ) -> 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."""
+    top-k rows, q quantized with the layer's scale (or, where the kernel
+    quantizes q itself, q as is with fp8 dots), 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,
@@ -99,6 +179,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs(
     impl.owns_indexer = True
     impl.index_epoch = 0
     impl.use_aiter_sparse_mla = True
+    impl.aiter_sparse_mla_quantizes_q = quantizes_q
     layer = SimpleNamespace(
         _q_scale=torch.tensor([0.02], device=device),
         _k_scale=torch.tensor([0.01], device=device),
@@ -130,7 +211,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs(
     (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:
+    if fp8 and not quantizes_q:
         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
@@ -146,6 +227,7 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs(
     assert kwargs["q_scale"] is layer._q_scale
     assert kwargs["kv_scale"] is layer._k_scale
     assert kwargs["attn_sink"] is impl.sinks
+    assert kwargs["dot_precision"] == ("fp8" if quantizes_q else None)
     assert lse is None
 
 

vllm/_aiter_ops.py

@@ -673,6 +673,20 @@ def _rocm_aiter_gather_kv_b_proj_fake(
     pass
 
 
+@functools.cache
+def _aiter_sparse_mla_archs(name: str) -> tuple[str, ...]:
+    """The archs the installed aiter's sparse_mla lists under name, or ().
+
+    gfx942 joins SUPPORTED_ARCHS in ROCm/aiter#5721 and FP8_SCALAR_ARCHS, the
+    archs that read a per-tensor fp8 cache, in ROCm/aiter#6199.
+    """
+    try:
+        from aiter.ops.triton.attention import sparse_mla
+    except ImportError:
+        return ()
+    return tuple(getattr(sparse_mla, name, ()))
+
+
 def _rocm_aiter_mla_decode_fwd_impl(
     q: torch.Tensor,
     kv_buffer: torch.Tensor,
@@ -1984,7 +1998,8 @@ 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_TRITON_SPARSE_MLA: Controls Triton sparse MLA (gfx950,
+            and rope-free sparse MLA on gfx942).
         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).
@@ -2455,15 +2470,28 @@ def is_triton_rotary_embed_enabled(cls) -> bool:
 
     @classmethod
     @if_aiter_supported
-    def is_triton_sparse_mla_enabled(cls) -> bool:
+    def is_triton_sparse_mla_enabled(cls, gfx942_ok: bool = False) -> bool:
+        """gfx942_ok: the caller passes only flat bf16 or per-tensor fp8 caches,
+        which aiter's kernel also runs on gfx942. DeepSeek V4's packed caches
+        are gfx950-only."""
         if not cls._TRITON_SPARSE_MLA:
             return False
-        from vllm.platforms.rocm import on_gfx950
+        from vllm.platforms.rocm import on_gfx942, on_gfx950
 
         if not cls._AITER_ENABLED:
             reason = "VLLM_ROCM_USE_AITER is off"
-        elif not on_gfx950():
-            reason = "the kernel is gfx950-only"
+        elif on_gfx950():
+            return True
+        elif not (gfx942_ok and on_gfx942()):
+            reason = (
+                "the kernel is gfx942/gfx950-only"
+                if gfx942_ok
+                else "the kernel is gfx950-only"
+            )
+        elif "gfx942" not in _aiter_sparse_mla_archs("SUPPORTED_ARCHS"):
+            reason = (
+                "the installed aiter has no gfx942 sparse MLA kernel (ROCm/aiter#5721)"
+            )
         else:
             return True
         logger.warning_once(
@@ -3451,6 +3479,11 @@ def fused_topk(
     ) -> 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_reads_fp8_on_gfx942() -> bool:
+        """Whether aiter's sparse MLA kernel reads an fp8 cache on gfx942."""
+        return "gfx942" in _aiter_sparse_mla_archs("FP8_SCALAR_ARCHS")
+
     @staticmethod
     def triton_sparse_mla_fwd(
         q: torch.Tensor,
@@ -3468,6 +3501,7 @@ def triton_sparse_mla_fwd(
         extra_kv_indptr: torch.Tensor | None = None,
         extra_kv_indices: torch.Tensor | None = None,
         has_invalid: bool = True,
+        dot_precision: str | None = None,
     ) -> None:
         """Sparse MLA read straight from the KV cache, for prefill and decode.
 
@@ -3476,13 +3510,17 @@ def triton_sparse_mla_fwd(
         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
+        must come with its q_scale and runs both dots in fp8. dot_precision
+        overrides that choice: gfx942's kernel takes no fp8 q, so with an fp8
+        cache it gets a bf16 q and "fp8", and quantizes q itself. 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
+        if dot_precision is None:
+            dot_precision = "fp8" if fp8_q else "bf16"
         sparse_mla_fwd(
             q,
             kv_buffer,
@@ -3493,7 +3531,7 @@ def triton_sparse_mla_fwd(
             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",
+            dot_precision=dot_precision,
             q_scale=q_scale if fp8_q else None,
             out=o,
             attn_sink=attn_sink,

vllm/envs.py

@@ -1371,7 +1371,9 @@ def _resolve_rust_cli_path() -> str | None:
     # 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.
+    # V4 / V4.1. gfx950 only, except that the ROCM_AITER_MLA_SPARSE backend
+    # also uses it on gfx942 for rope-free sparse MLA (GLM-5.3-Flash) when the
+    # installed aiter has a gfx942 kernel. 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")

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

@@ -780,6 +780,9 @@ def log2sumexp2(a: torch.Tensor, dim: int) -> torch.Tensor:
 class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]):
     supports_dcp = False
     use_aiter_sparse_mla = False
+    # gfx942's aiter kernel takes no fp8 q. With an fp8 cache it gets q in the
+    # model dtype and quantizes it itself, under fp8 dots.
+    aiter_sparse_mla_quantizes_q = False
 
     def __init__(
         self,
@@ -845,10 +848,15 @@ def __init__(
         )
         self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
 
-        if rocm_aiter_ops.is_triton_sparse_mla_enabled():
+        if rocm_aiter_ops.is_triton_sparse_mla_enabled(gfx942_ok=True):
             reason = self._aiter_sparse_mla_unsupported_reason(vllm_config)
             if reason is None:
+                from vllm.platforms.rocm import on_gfx942
+
                 self.use_aiter_sparse_mla = True
+                self.aiter_sparse_mla_quantizes_q = (
+                    on_gfx942() and kv_cache_dtype.startswith("fp8")
+                )
             else:
                 logger.warning_once(
                     "VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA is set, but %s; "
@@ -868,6 +876,19 @@ def _aiter_sparse_mla_unsupported_reason(
         # 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"
+        from vllm.platforms.rocm import on_gfx942
+
+        if on_gfx942():
+            if self.qk_rope_head_dim != 0:
+                return "on gfx942 it covers rope-free sparse MLA only"
+            if (
+                cache_dtype.startswith("fp8")
+                and not rocm_aiter_ops.triton_sparse_mla_reads_fp8_on_gfx942()
+            ):
+                return (
+                    "the installed aiter cannot read an fp8 KV cache on gfx942 "
+                    "(ROCm/aiter#6199)"
+                )
         return None
 
     def record_logical_topk_ready(self) -> None:
@@ -1148,6 +1169,7 @@ def _forward_mla_aiter(
             q_scale=layer._q_scale,
             kv_scale=layer._k_scale,
             attn_sink=self.sinks,
+            dot_precision="fp8" if self.aiter_sparse_mla_quantizes_q else None,
             # 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,
@@ -1165,9 +1187,10 @@ def forward_mqa(
         # MQA 576/512 approach for both prefill and decode
 
         fp8_attention = self.kv_cache_dtype.startswith("fp8")
+        quantize_q = fp8_attention and not self.aiter_sparse_mla_quantizes_q
         if isinstance(q, tuple):
             ql_nope, q_pe = q
-            if fp8_attention:
+            if quantize_q:
                 q = layer._decode_concat_quant_fp8_op(  # type: ignore[attr-defined]
                     ql_nope, q_pe, layer._q_scale
                 )
@@ -1215,7 +1238,7 @@ def forward_mqa(
         # write the latent and rope to kv cache
         if fp8_attention:
             kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(current_platform.fp8_dtype())
-            if q.dtype != current_platform.fp8_dtype():
+            if quantize_q and q.dtype != current_platform.fp8_dtype():
                 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)