feat(attention): FlyDSL paged decode for MiniMax-M3 on gfx1250

atom/model_ops/attention_mha.py

@@ -9,6 +9,7 @@
 from aiter.jit.utils.chip_info import get_gfx
 from aiter.ops.mha import _flash_attn_varlen_forward
 from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
+from aiter.ops.triton.gluon.pa_decode_gluon import pa_decode_gluon_supports_arch
 from aiter.ops.triton.unified_attention import unified_attention
 from torch import nn
 
@@ -17,6 +18,7 @@
     PA_ASM_MAX_QUERY_GROUP_SIZE,
     cp_mha_gather_cache,
     dense_decode_splits,
+    flydsl_decode_seqs,
     gluon_decode_over_limit,
     run_pa_decode,
     run_pa_fwd_asm,
@@ -434,8 +436,9 @@ def _gemma_norm_rope_then_cache(
         the separate kernels at 16384 tokens, Qwen3.5 TP4 on MI355X), so
         above ATOM_QK_NORM_ROPE_CACHE_FUSION_MAX_TOKENS run them separately:
         Gemma q/k norm in one launch, the Qwen3.5 MRoPE Triton kernel (or the
-        module's own rotary when that declines), then the per-token FP8 cache
-        write. Same SHUFFLE layout and scales the fused kernel writes.
+        module's own rotary when that declines), then the cache write: per-token
+        FP8 for an FP8 cache, unquantized otherwise. Same SHUFFLE layout and
+        scales the fused kernel writes.
         """
         from atom.model_ops.layernorm import DualRMSNorm
         from atom.model_ops.triton_mrope import try_mrope_qk_fused
@@ -463,16 +466,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(
@@ -574,6 +590,51 @@ def _view_v_cache_for_pa_decode_bf16_asm(
         x = int(k_cache.shape[-1])
         return v_cache.view(n, nh, block_size // x, head_dim, x)
 
+    def _shuffle_v_for_flydsl(self, v_cache, k_cache):
+        # FlyDSL on gfx1250 reads a 4D V as the plain layout; this cache is
+        # SHUFFLE, so pass its 5D view. gluon takes the 4D view as it always has.
+        if pa_decode_gluon_supports_arch() or v_cache.dim() != 4 or k_cache.dim() != 5:
+            return v_cache
+        return self._view_v_cache_for_pa_decode_bf16_asm(v_cache, k_cache)
+
+    def _flydsl_takes_decode(self, q, k_cache, v_cache, attn_metadata) -> bool:
+        """Without gluon decode (gfx1250): whether FlyDSL accepts this call.
+
+        ATOM_USE_UNIFIED_ATTN keeps the call on unified_attention.
+        """
+        if (
+            envs.ATOM_USE_UNIFIED_ATTN
+            or pa_decode_gluon_supports_arch()
+            or self.sliding_window > 0
+        ):
+            return False
+        num_seqs = attn_metadata.context_lens.shape[0]
+        return bool(
+            flydsl_decode_seqs(
+                output=q,
+                q=q,
+                k_cache=k_cache,
+                v_cache=self._shuffle_v_for_flydsl(v_cache, k_cache),
+                block_tables=attn_metadata.block_tables,
+                context_lens=attn_metadata.context_lens,
+                max_seqlen_q=attn_metadata.max_seqlen_q,
+                max_context_partition_num=dense_decode_splits(
+                    num_seqs, k_cache.shape[1]
+                ),
+                context_partition_size=256,
+                compute_type=(
+                    torch.bfloat16
+                    if self.kv_cache_dtype == "bf16"
+                    else aiter.dtypes.fp8
+                ),
+                q_scale=None,
+                alibi_slopes=None,
+                sinks=self.sinks,
+                sliding_window=self.sliding_window,
+                ps=True,
+            )
+        )
+
     @mark_trace(prefix="paged_attention_unified", torch_compile=False)
     def paged_attention_unified(
         self, q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx: ForwardContext
@@ -585,6 +646,10 @@ def paged_attention_unified(
         of the reasons for it.
         """
         attn_metadata = fwd_ctx.attn_metadata
+        if self._flydsl_takes_decode(q, k_cache, v_cache, attn_metadata):
+            return self.paged_attention_triton(
+                q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx
+            )
         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"):
@@ -622,8 +687,10 @@ def paged_attention_unified(
     def paged_attention_triton(
         self, q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx: ForwardContext
     ):
-        """Decode through the gluon paged kernel. Callers must have cleared
-        gluon_decode_over_limit -- _dispatch_decode is where that happens."""
+        """Decode through run_pa_decode: gluon, or FlyDSL where ATOM_PA_FLYDSL
+        routes it (always FlyDSL on gfx1250). Callers must have cleared
+        gluon_decode_over_limit -- _dispatch_decode is where that happens --
+        or, without gluon, _flydsl_takes_decode."""
 
         attn_metadata = fwd_ctx.attn_metadata
 
@@ -674,7 +741,7 @@ def paged_attention_triton(
             output=o,
             q=q,
             k_cache=k_cache,
-            v_cache=v_cache,
+            v_cache=self._shuffle_v_for_flydsl(v_cache, k_cache),
             context_lens=attn_metadata.context_lens,
             block_tables=attn_metadata.block_tables,
             softmax_scale=self.scale,
@@ -1037,11 +1104,16 @@ def _dispatch_decode(self, max_qlen: int = 1):
         checked here rather than inside its runner: past its envelope
         get_heuristic_kernel silently re-runs with mtp=1, a kernel built for
         another query length, and computes a wrong answer instead of refusing.
+
+        One exception: without gluon decode (gfx1250) this picks unified, and
+        paged_attention_unified hands FlyDSL the calls it accepts. That check
+        needs the cache and metadata, which this method does not see.
         """
         # 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(
+        # No gluon decode (gfx1250): unified, which hands FlyDSL what it takes.
+        over_gluon = not pa_decode_gluon_supports_arch() or gluon_decode_over_limit(
             max_qlen, self.num_heads, self.num_kv_heads
         )
         over_asm = (

atom/model_ops/base_attention.py

@@ -3,13 +3,17 @@
 
 # from flash_attn import flash_attn_with_kvcache
 import functools
+import importlib
 import logging
 import os
 from abc import ABC, abstractmethod
 
 import torch
 import triton
 import triton.language as tl
+from aiter.ops.triton.gluon.pa_decode_gluon import (
+    pa_decode_gluon_supports_arch,
+)
 from torch import nn
 
 from atom.config import get_current_atom_config
@@ -175,7 +179,7 @@ def run_pa_fwd_asm(
 _FLYDSL_PA_MAX_PARTITIONS = 256
 _FLYDSL_PA_TILE = 256
 _FLYDSL_PA_BLOCK_SIZES = (16, 64, 128)
-_FLYDSL_PA_ARCHS = ("gfx942", "gfx950")
+_FLYDSL_PA_ARCHS = ("gfx942", "gfx950", "gfx1250")
 _flydsl_pa_routed: set[tuple] = set()
 _flydsl_plan_refused: set[tuple] = set()
 
@@ -219,7 +223,9 @@ def _flydsl_explicit_plan_api() -> bool:
     import inspect
 
     try:
-        from aiter.ops.flydsl import pa_decode as mod
+        # By module path: aiter.ops.flydsl lazily exports a `pa_decode`
+        # function, which `from aiter.ops.flydsl import pa_decode` can return.
+        mod = importlib.import_module("aiter.ops.flydsl.pa_decode")
     except ImportError:
         ok = False
     else:
@@ -234,7 +240,7 @@ def _flydsl_explicit_plan_api() -> bool:
         )
     if not ok:
         logger.warning(
-            "ATOM_PA_FLYDSL=1 but this aiter predates the #5809 FlyDSL paged "
+            "FlyDSL paged decode is enabled but this aiter predates the #5809 "
             "decode API; using gluon."
         )
     return ok
@@ -282,9 +288,16 @@ def _flydsl_pa_decode_num_seqs(
         return None
     if sliding_window > 0 or not ps:
         return None
-    if compute_type is not aiter.dtypes.fp8 or k_cache.dtype is not aiter.dtypes.fp8:
+    # Native BF16 KV uses 8-element vectors and no scales; FP8 uses 16.
+    is_bf16_kv = compute_type is torch.bfloat16 and k_cache.dtype is torch.bfloat16
+    is_fp8_kv = compute_type is aiter.dtypes.fp8 and k_cache.dtype is aiter.dtypes.fp8
+    if not (is_fp8_kv or is_bf16_kv):
+        return None
+    # BF16 KV stays on gluon where gluon exists.
+    if is_bf16_kv and pa_decode_gluon_supports_arch():
         return None
-    if k_cache.dim() != 5 or k_cache.shape[-1] != 16:
+    x = 8 if is_bf16_kv else 16
+    if k_cache.dim() != 5 or k_cache.shape[-1] != x:
         return None
     # block_size is dim -2 of the page-16 cache. --block-size 256/1024 reaches
     # this site whenever use_triton_attn is set, which the published recipe does.
@@ -301,9 +314,9 @@ def _flydsl_pa_decode_num_seqs(
     head_dim = q.shape[-1]
     if not (head_dim == 64 or (head_dim % 128 == 0 and head_dim <= 1024)):
         return None
-    # The cache encodes head_dim as num_hgroups * 16; aiter rejects a q whose
+    # The cache encodes head_dim as num_hgroups * x; aiter rejects a q whose
     # own head_dim disagrees, and nothing upstream forces the two to match.
-    if k_cache.shape[2] * 16 != head_dim:
+    if k_cache.shape[2] * x != head_dim:
         return None
     if q.dtype not in (torch.bfloat16, torch.float16):
         return None
@@ -342,7 +355,37 @@ def _flydsl_pa_decode_num_seqs(
 
 
 _FLYDSL_PLAN_MAX_BATCH = 4096
+# aiter's planner keeps one LDS int32 per sequence plus this much scan storage.
+_FLYDSL_PLAN_SCAN_LDS_BYTES = 512
 _FLYDSL_PLAN_SCRATCH: dict[tuple, tuple] = {}
+
+
+@functools.lru_cache(maxsize=1)
+def _flydsl_plan_max_batch() -> int:
+    """Most sequences one FlyDSL plan takes on this GPU.
+
+    gfx1250 only, where sparse decode has no gluon to fall back to: past
+    _FLYDSL_PLAN_MAX_BATCH a plan sets max_batch_size (aiter #6462), bounded by
+    the planner's LDS: the 16-byte-rounded task_ends array plus scan storage
+    must fit one block (81792 sequences on gfx1250).
+    """
+    import inspect
+
+    from aiter.jit.utils.chip_info import get_gfx_runtime
+
+    planner = importlib.import_module("aiter.ops.flydsl.pa_decode").plan_pa_decode
+    if (
+        get_gfx_runtime() != "gfx1250"
+        or "max_batch_size" not in inspect.signature(planner).parameters
+    ):
+        return _FLYDSL_PLAN_MAX_BATCH
+    lds = torch.cuda.get_device_properties(
+        torch.cuda.current_device()
+    ).shared_memory_per_block
+    by_lds = (lds - _FLYDSL_PLAN_SCAN_LDS_BYTES) // 16 * 16 // 4
+    return max(_FLYDSL_PLAN_MAX_BATCH, by_lds)
+
+
 _FLYDSL_RUNTIME_PLANS: dict[tuple, object] = {}
 
 
@@ -490,6 +533,9 @@ def _flydsl_prepare_explicit_plan(
         budget = (
             {} if workgroup_budget is None else {"workgroup_budget": workgroup_budget}
         )
+        # A refresh keeps the plan's max_batch_size, so set it at creation.
+        if num_seqs > _FLYDSL_PLAN_MAX_BATCH:
+            budget["max_batch_size"] = num_seqs
         plan = plan_pa_decode(
             lengths,
             num_kv_heads,
@@ -518,6 +564,25 @@ def _flydsl_prepare_explicit_plan(
     return plan
 
 
+def flydsl_decode_seqs(*, max_context_length=None, **kwargs):
+    """Sequence count FlyDSL's paged decode takes a run_pa_decode call with.
+
+    False when FlyDSL is off, None when the call is outside FlyDSL's domain.
+    """
+    flydsl_seqs = envs.ATOM_PA_FLYDSL and _flydsl_pa_decode_num_seqs(**kwargs)
+    # aiter's planner rejects more sequences than one plan holds.
+    if flydsl_seqs and flydsl_seqs > _flydsl_plan_max_batch():
+        flydsl_seqs = None
+    # Only on the FlyDSL path; a bound below the query window goes to gluon.
+    if flydsl_seqs:
+        context_bound = _flydsl_context_bound(
+            max_context_length, kwargs["block_tables"], kwargs["k_cache"]
+        )
+        if context_bound < kwargs["max_seqlen_q"]:
+            flydsl_seqs = None
+    return flydsl_seqs
+
+
 def run_pa_decode(
     output: torch.Tensor,
     q: torch.Tensor,
@@ -552,12 +617,12 @@ def run_pa_decode(
     it dispatched only gluon until ``ATOM_PA_FLYDSL`` arrived, and four call
     sites -- including the vLLM and SGLang bridges -- import it.
 
-    gluon unless ``ATOM_PA_FLYDSL=1``, and then only where FlyDSL's domain
-    covers the call: ``_flydsl_pa_decode_num_seqs`` mirrors the kernel's own
-    validation so an unsupported shape falls back here instead of raising
+    gluon unless ``ATOM_PA_FLYDSL=1``, and then only where FlyDSL's
+    domain covers the call: ``_flydsl_pa_decode_num_seqs`` mirrors the kernel's
+    own validation so an unsupported shape falls back here instead of raising
     inside aiter.
     """
-    flydsl_seqs = envs.ATOM_PA_FLYDSL and _flydsl_pa_decode_num_seqs(
+    flydsl_seqs = flydsl_decode_seqs(
         output=output,
         q=q,
         k_cache=k_cache,
@@ -573,15 +638,8 @@ def run_pa_decode(
         sinks=sinks,
         sliding_window=sliding_window,
         ps=ps,
+        max_context_length=max_context_length,
     )
-    # aiter's planner rejects more than _FLYDSL_PLAN_MAX_BATCH sequences.
-    if flydsl_seqs and flydsl_seqs > _FLYDSL_PLAN_MAX_BATCH:
-        flydsl_seqs = None
-    # Only on the FlyDSL path; a bound below the query window goes to gluon.
-    if flydsl_seqs:
-        context_bound = _flydsl_context_bound(max_context_length, block_tables, k_cache)
-        if context_bound < max_seqlen_q:
-            flydsl_seqs = None
     # Inside the guard, not before it: this runs 63 times per decode step and
     # attention is an eager region of the token graphs, so replay does not
     # elide it. A deployment that never enables FlyDSL should pay nothing here.
@@ -616,6 +674,7 @@ def run_pa_decode(
         from aiter.ops.flydsl.pa_decode import pa_decode as _flydsl_pa_decode
 
         n = flydsl_seqs
+        context_bound = _flydsl_context_bound(max_context_length, block_tables, k_cache)
         # Handed in by the caller off the ForwardContext it already holds, not
         # re-read from the thread-local one: under TBO a worker thread that
         # never installed its own context would read whatever the other ubatch
@@ -688,6 +747,13 @@ def run_pa_decode(
             max_context_length=context_bound,
         )
 
+    if not pa_decode_gluon_supports_arch():
+        raise RuntimeError(
+            "This GPU has no gluon paged decode: set ATOM_PA_FLYDSL=1"
+            if not envs.ATOM_PA_FLYDSL
+            else "FlyDSL paged decode does not take this call and this GPU "
+            f"has no gluon fallback (q {tuple(q.shape)}, {compute_type})"
+        )
     return torch.ops.aiter.pa_decode_gluon(
         output,
         q,

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.ops.attention import pa_fwd_asm_supports_arch
 
 from atom.model_ops.minimax_m3.index_topk import build_n_valid_column_per_row
 
@@ -246,7 +247,8 @@ def _sparse_pa_ran_on_asm(
 
     rows, group, head_dim = q.shape
     if (
-        k_scale is None
+        not pa_fwd_asm_supports_arch()
+        or k_scale is None
         or not _is_fp8_kv_cache_tensor(k_cache)
         or group > PA_ASM_MAX_QUERY_GROUP_SIZE
         or rows <= _gluon_one_pass_max_rows(q.device.index)