fix(minimax-m3): enable per-token FP8 KV on gfx1250

atom/model_ops/attention_mha.py

@@ -18,6 +18,7 @@
     cp_mha_gather_cache,
     dense_decode_splits,
     gluon_decode_over_limit,
+    gluon_decode_supports_arch,
     run_pa_decode,
     run_pa_fwd_asm,
 )
@@ -194,6 +195,45 @@ def _reject_per_token_scales_on_unified(self, k_scale, v_scale, phase: str):
                 "which cannot carry its per-token KV scales"
             )
 
+    def _can_use_per_token_shuffled_decode(
+        self,
+        q: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        k_scale: torch.Tensor | None,
+        v_scale: torch.Tensor | None,
+        attn_metadata,
+    ) -> bool:
+        """Whether the generic shuffled-FP8 decode contract is fully satisfied."""
+        if get_gfx() != "gfx1250" or self.use_flash_layout:
+            return False
+        if self.sinks is not None or self.alibi_slopes is not None:
+            return False
+        if self.logits_soft_cap not in (None, 0, 0.0):
+            return False
+        if (
+            k_scale is None
+            or v_scale is None
+            or k_scale.numel() <= 1
+            or v_scale.numel() <= 1
+        ):
+            return False
+        if q.ndim != 3 or k_cache.ndim != 5:
+            return False
+        if attn_metadata.max_seqlen_q != 1:
+            return False
+        if q.shape[0] != attn_metadata.context_lens.numel():
+            return False
+        if k_scale.ndim != 3 or v_scale.ndim != 3:
+            return False
+        return (
+            k_scale.shape == v_scale.shape
+            and k_scale.shape[:2] == k_cache.shape[:2]
+            and k_scale.shape[-1] == k_cache.shape[-2]
+            and k_cache.element_size() == 1
+            and v_cache.element_size() == 1
+        )
+
     def forward_impl(
         self,
         q: torch.Tensor,
@@ -463,16 +503,29 @@ def _gemma_norm_rope_then_cache(
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)
-        aiter.reshape_and_cache_with_pertoken_quant(
-            k,
-            v,
-            k_cache,
-            v_cache,
-            k_scale,
-            v_scale,
-            attn_metadata.slot_mapping,
-            asm_layout=True,
-        )
+        if self.kv_cache_dtype == "fp8":
+            aiter.reshape_and_cache_with_pertoken_quant(
+                k,
+                v,
+                k_cache,
+                v_cache,
+                k_scale,
+                v_scale,
+                attn_metadata.slot_mapping,
+                asm_layout=True,
+            )
+        else:
+            aiter.reshape_and_cache(
+                k,
+                v,
+                k_cache,
+                v_cache,
+                attn_metadata.slot_mapping,
+                kv_cache_dtype="auto",
+                k_scale=None,
+                v_scale=None,
+                asm_layout=True,
+            )
         return q, k, v
 
     def _gather_prefix_and_concat_kv(
@@ -585,6 +638,28 @@ def paged_attention_unified(
         of the reasons for it.
         """
         attn_metadata = fwd_ctx.attn_metadata
+        if self._can_use_per_token_shuffled_decode(
+            q, k_cache, v_cache, k_scale, v_scale, attn_metadata
+        ):
+            from aiter.ops.triton.attention.pa_decode_shuffle import (
+                paged_attention_decode_shuffle_swa,
+            )
+
+            o = torch.empty_like(q)
+            paged_attention_decode_shuffle_swa(
+                o,
+                q,
+                k_cache,
+                self._view_v_cache_for_pa_decode_bf16_asm(v_cache, k_cache),
+                attn_metadata.block_tables,
+                attn_metadata.context_lens,
+                self.scale,
+                k_scale,
+                v_scale,
+                self.sliding_window,
+            )
+            return o
+
         self._reject_per_token_scales_on_unified(k_scale, v_scale, "decode")
 
         if envs.ATOM_USE_UNIFIED_ATTN and self.kv_cache_dtype.startswith("fp8"):
@@ -1041,7 +1116,7 @@ def _dispatch_decode(self, max_qlen: int = 1):
         # Clamped once here so both gates see the same value; the predicate
         # clamps internally too, and a sentinel would otherwise split them.
         max_qlen = max(1, int(max_qlen))
-        over_gluon = gluon_decode_over_limit(
+        over_gluon = not gluon_decode_supports_arch() or gluon_decode_over_limit(
             max_qlen, self.num_heads, self.num_kv_heads
         )
         over_asm = (

atom/model_ops/base_attention.py

@@ -128,6 +128,14 @@ def gluon_decode_over_limit(max_qlen: int, num_heads: int, num_kv_heads: int) ->
     return max_qlen > PA_GLUON_MAX_QUERY_LEN or group_p2 > PA_GLUON_MAX_QUERY_GROUP_SIZE
 
 
+@functools.cache
+def gluon_decode_supports_arch() -> bool:
+    """Whether aiter's pa_decode_gluon supports the current GPU."""
+    from aiter.jit.utils.chip_info import get_gfx
+
+    return get_gfx() in ("gfx942", "gfx950")
+
+
 def run_pa_fwd_asm(
     q: torch.Tensor,
     k_cache: torch.Tensor,

atom/model_ops/layernorm.py

@@ -14,10 +14,12 @@
     rmsnorm2d_fwd_with_add,
 )
 from aiter.dist.communication_op import (
+    tensor_model_parallel_all_reduce,
     tensor_model_parallel_fused_allreduce_rmsnorm,
     tensor_model_parallel_fused_allreduce_rmsnorm_quant,
 )
 from aiter.dist.parallel_state import get_tensor_model_parallel_world_size
+from aiter.jit.utils.chip_info import get_gfx
 from aiter.jit.utils.torch_guard import torch_compile_guard
 from aiter.ops.gated_rmsnorm_fp8_group_quant import gated_rmsnorm_fp8_group_quant
 from aiter.ops.triton.fused_add_rmsnorm_pad import fused_add_rmsnorm_pad
@@ -969,6 +971,11 @@ def fused_allreduce_gemma_rms_norm(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """MiniMax-M3 helper for delayed TP all-reduce followed by Gemma RMSNorm."""
     if get_tensor_model_parallel_world_size() > 1:
+        if get_gfx() == "gfx1250":
+            return norm(
+                tensor_model_parallel_all_reduce(hidden_states.contiguous()),
+                residual,
+            )
         return tensor_model_parallel_fused_allreduce_rmsnorm(
             hidden_states.contiguous(),
             residual,
@@ -986,6 +993,19 @@ def fused_allreduce_gemma_rms_norm_quant(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """MiniMax-M3 helper for AR + Gemma RMSNorm + per-token FP8 quant."""
     if get_tensor_model_parallel_world_size() > 1:
+        if get_gfx() == "gfx1250":
+            normed, residual_out = norm(
+                tensor_model_parallel_all_reduce(hidden_states.contiguous()),
+                residual,
+            )
+            from aiter import get_hip_quant
+            from aiter.utility.dtypes import fp8
+
+            out_fp8, scale_out = get_hip_quant(QuantType.per_Token)(
+                normed,
+                quant_dtype=fp8,
+            )
+            return out_fp8, scale_out, residual_out
         out_fp8, residual_out, scale_out = (
             tensor_model_parallel_fused_allreduce_rmsnorm_quant(
                 hidden_states.contiguous(),

atom/model_ops/minimax_m3/sparse_attn.py

@@ -20,6 +20,7 @@
 # which couples to vLLM internals. ATOM imports triton directly everywhere else.
 import triton
 import triton.language as tl
+from aiter.jit.utils.chip_info import get_gfx
 
 from atom.model_ops.minimax_m3.index_topk import build_n_valid_column_per_row
 
@@ -244,6 +245,13 @@ def _sparse_pa_ran_on_asm(
         run_pa_fwd_asm,
     )
 
+    # gfx1250 ships BF16 sparse PA code objects, but not MiniMax's per-token
+    # FP8 GQA variants. The low-level launcher aborts the process rather than
+    # reporting an unsupported shape, so decline before applying the row
+    # profitability heuristic and use the shuffled-FP8 Triton path instead.
+    if get_gfx() == "gfx1250":
+        return False
+
     rows, group, head_dim = q.shape
     if (
         k_scale is None
@@ -1358,6 +1366,66 @@ def minimax_m3_sparse_attn_decode_asm(
     )
 
 
+def _sparse_unified_per_row(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    sparse_bt: torch.Tensor,
+    sparse_ctx: torch.Tensor,
+    sm_scale: float,
+    output: torch.Tensor,
+    k_scale: torch.Tensor | None,
+    v_scale: torch.Tensor | None,
+) -> None:
+    """Run each sparse row as a length-one sequence through unified attention."""
+    if (k_scale is not None and k_scale.numel() > 1) or (
+        v_scale is not None and v_scale.numel() > 1
+    ):
+        from aiter.ops.triton.attention.pa_decode_shuffle import (
+            paged_attention_decode_shuffle_swa,
+        )
+
+        assert k_scale is not None and v_scale is not None
+        paged_attention_decode_shuffle_swa(
+            output,
+            q,
+            k_cache,
+            v_cache,
+            sparse_bt,
+            sparse_ctx,
+            sm_scale,
+            k_scale,
+            v_scale,
+            sliding_window=-1,
+        )
+        return
+
+    from aiter.ops.triton.unified_attention import unified_attention
+
+    rows = q.shape[0]
+    unified_attention(
+        q,
+        k_cache,
+        v_cache,
+        output,
+        cu_seqlens_q=torch.arange(rows + 1, dtype=torch.int32, device=q.device),
+        seqused_k=sparse_ctx,
+        max_seqlen_q=1,
+        max_seqlen_k=sparse_bt.shape[1] * ASM_PAGE_SIZE,
+        softmax_scale=sm_scale,
+        causal=True,
+        alibi_slopes=None,
+        window_size=(-1, -1),
+        block_table=sparse_bt,
+        softcap=0,
+        q_descale=None,
+        k_descale=k_scale,
+        v_descale=v_scale,
+        sinks=None,
+        shuffled_kv_cache=True,
+    )
+
+
 @torch.no_grad()
 def _sparse_pa_per_row(
     q: torch.Tensor,  # [total_q, num_heads, head_dim==128]
@@ -1384,7 +1452,10 @@ def _sparse_pa_per_row(
     """
     from aiter.ops.triton.gluon.pa_decode_gluon import get_recommended_splits
 
-    from atom.model_ops.base_attention import run_pa_decode
+    from atom.model_ops.base_attention import (
+        gluon_decode_supports_arch,
+        run_pa_decode,
+    )
 
     assert q.shape[-1] == 128, "Paged attention here requires head_dim == 128."
     # Collapse (token, kv_head) -> row so gluon runs num_kv_heads_view == 1.
@@ -1413,6 +1484,27 @@ def _sparse_pa_per_row(
     ):
         return
 
+    if not gluon_decode_supports_arch():
+        if k_scale is not None and k_scale.numel() > 1:
+            pbs = k_scale.shape[-1]
+            fallback_k_scale = k_scale.view(nph16 * _hkv, 1, pbs)
+            fallback_v_scale = v_scale.view(nph16 * _hkv, 1, pbs)
+        else:
+            fallback_k_scale = k_scale
+            fallback_v_scale = v_scale
+        _sparse_unified_per_row(
+            q_view,
+            k_cache_view,
+            v_cache_view,
+            sparse_bt,
+            sparse_ctx,
+            sm_scale,
+            out_view,
+            fallback_k_scale,
+            fallback_v_scale,
+        )
+        return
+
     num_seqs = T * num_kv_heads
     num_kv_heads_view = 1
     query_group_size = g