fix(attention): keep MTP-5 decode on FlyDSL pa_decode past the gluon query-length limit

atom/model_ops/attention_mha.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
+import logging
 from functools import cache
 
 import aiter
@@ -29,6 +30,12 @@
 
 from .attention_mla import MLAModules
 
+logger = logging.getLogger("atom")
+
+# Once per process: the hint names the opt-in, not the shape, so one line says
+# all a default user needs, and the decode loop pays one bool test after it.
+_flydsl_off_hint_logged = False
+
 
 @cache
 def use_pa_decode_bf16_asm() -> bool:
@@ -181,16 +188,23 @@ def _can_use_prefill_sink_asm(
             return False
         return q.shape[1] % k.shape[1] == 0
 
-    def _reject_per_token_scales_on_unified(self, k_scale, v_scale, phase: str):
+    @staticmethod
+    def _per_token_kv(k_scale, v_scale) -> bool:
+        return (k_scale is not None and k_scale.numel() > 1) or (
+            v_scale is not None and v_scale.numel() > 1
+        )
+
+    def _reject_per_token_scales_on_unified(
+        self, k_scale, v_scale, phase: str, flydsl_would_take: bool = False
+    ):
         """unified_attention carries one descale for the whole tensor.
 
         Feeding it a per-token cache is not an error downstream, just wrong
         numbers, so it has to be refused here. Both phases route into unified on
-        the same conditions, so both check.
+        the same conditions, so both check. ``flydsl_would_take``: the decode
+        is one ATOM_PA_FLYDSL=1 would offer to FlyDSL, so the error says so.
         """
-        if (k_scale is not None and k_scale.numel() > 1) or (
-            v_scale is not None and v_scale.numel() > 1
-        ):
+        if self._per_token_kv(k_scale, v_scale):
             if phase == "decode" and not pa_decode_gluon_supports_arch():
                 reason = (
                     "set ATOM_PA_FLYDSL=1"
@@ -202,9 +216,15 @@ def _reject_per_token_scales_on_unified(self, k_scale, v_scale, phase: str):
                     f"cache on this GPU: {reason}, unified_attention cannot "
                     "carry per-token KV scales, and there is no gluon fallback"
                 )
+            hint = (
+                "; ATOM_PA_FLYDSL=1 lets aiter's FlyDSL pa_decode take it "
+                "(needs aiter with ROCm/aiter#7206)"
+                if flydsl_would_take
+                else ""
+            )
             raise NotImplementedError(
                 f"layer {self.layer_num} takes unified_attention for {phase}, "
-                "which cannot carry its per-token KV scales"
+                f"which cannot carry its per-token KV scales{hint}"
             )
 
     def forward_impl(
@@ -608,19 +628,61 @@ def _shuffle_v_for_flydsl(self, v_cache, k_cache):
             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.
-        """
+    def _pa_decode_kv_scales(self, k_scale, v_scale):
+        """The K/V scales ``paged_attention_triton`` hands ``run_pa_decode``."""
+        if self.kv_cache_dtype == "bf16":
+            return None, None
+        if k_scale is not None and k_scale.numel() > 1:
+            return k_scale.unsqueeze(-1), v_scale.unsqueeze(-1)
+        return k_scale, v_scale
+
+    def _flydsl_route_open(self, attn_metadata) -> bool:
+        """Whether this unified-bound decode is one ATOM_PA_FLYDSL offers to
+        FlyDSL: not forced to unified, paged layout, no window, no sinks, not
+        BF16 KV on a gluon GPU (ATOM's own routing in
+        ``_flydsl_pa_decode_num_seqs``), and on a gluon GPU past the gluon
+        envelope (inside it the gluon route is kept). aiter's predicate is not
+        asked here; the opt-in hint relies on this, so it never names a route
+        ATOM's routing would refuse."""
+        gluon = pa_decode_gluon_supports_arch()
         if (
             envs.ATOM_USE_UNIFIED_ATTN
-            or pa_decode_gluon_supports_arch()
+            or self.use_flash_layout
             or self.sliding_window > 0
+            or self.sinks is not None
+            or (self.kv_cache_dtype == "bf16" and gluon)
         ):
             return False
-        num_seqs = attn_metadata.context_lens.shape[0]
-        return bool(
+        return not gluon or gluon_decode_over_limit(
+            attn_metadata.max_seqlen_q, self.num_heads, self.num_kv_heads
+        )
+
+    def _flydsl_takes_decode(
+        self, q, k_cache, v_cache, k_scale, v_scale, attn_metadata
+    ):
+        """Sequence count FlyDSL runs a decode routed to unified with, or a
+        falsy value to keep it on unified.
+
+        Two cases reach unified with a call FlyDSL may take: a GPU without
+        gluon decode (gfx1250), and a gluon GPU whose query length x group is
+        past the gluon envelope (Qwen3.5 MTP-5: qlen 6, group 8). Inside that
+        envelope a gluon GPU keeps its route. The decision is made per call,
+        on the same tensors and metadata ``run_pa_decode`` then receives, and
+        the count is passed down as ``flydsl_seqs``, so aiter's predicate is
+        asked once per call and the two cannot disagree.
+
+        With ATOM_PA_FLYDSL unset such a call stays on unified;
+        ``paged_attention_unified`` says what the flag would change.
+
+        ATOM_USE_UNIFIED_ATTN keeps the call on unified_attention.
+
+        gfx1250 with per-token FP8 KV past qlen 4 is decided by the same check;
+        it is not exercised by this path's tests or runs (out of scope here).
+        """
+        if not envs.ATOM_PA_FLYDSL or not self._flydsl_route_open(attn_metadata):
+            return False
+        k_scale, v_scale = self._pa_decode_kv_scales(k_scale, v_scale)
+        return (
             flydsl_decode_seqs(
                 output=q,
                 q=q,
@@ -629,21 +691,21 @@ def _flydsl_takes_decode(self, q, k_cache, v_cache, attn_metadata) -> bool:
                 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,
+                k_scale=k_scale,
+                v_scale=v_scale,
                 alibi_slopes=None,
                 sinks=self.sinks,
                 sliding_window=self.sliding_window,
                 ps=True,
             )
+            or False
         )
 
     @mark_trace(prefix="paged_attention_unified", torch_compile=False)
@@ -657,11 +719,40 @@ 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):
+        flydsl_seqs = self._flydsl_takes_decode(
+            q, k_cache, v_cache, k_scale, v_scale, attn_metadata
+        )
+        if flydsl_seqs:
             return self.paged_attention_triton(
-                q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx
+                q,
+                k,
+                v,
+                k_cache,
+                v_cache,
+                k_scale,
+                v_scale,
+                fwd_ctx,
+                flydsl_seqs=flydsl_seqs,
+            )
+        # A default user gets the opt-in named once: in the per-token FP8 KV
+        # refusal (unified cannot carry those scales), else in one info line.
+        flydsl_would_take = not envs.ATOM_PA_FLYDSL and self._flydsl_route_open(
+            attn_metadata
+        )
+        self._reject_per_token_scales_on_unified(
+            k_scale, v_scale, "decode", flydsl_would_take
+        )
+        global _flydsl_off_hint_logged
+        if flydsl_would_take and not _flydsl_off_hint_logged:
+            _flydsl_off_hint_logged = True
+            logger.info(
+                "decode with max_seqlen_q=%d x %d q heads per kv head runs on "
+                "unified_attention because ATOM_PA_FLYDSL is unset; "
+                "ATOM_PA_FLYDSL=1 lets aiter's FlyDSL pa_decode take it "
+                "(needs aiter with ROCm/aiter#7206)",
+                attn_metadata.max_seqlen_q,
+                self.num_heads // self.num_kv_heads,
             )
-        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"):
             o = torch.empty(*q.shape, dtype=torch.bfloat16, device=q.device)
@@ -696,13 +787,23 @@ def paged_attention_unified(
 
     @mark_trace(prefix="paged_attention_triton", torch_compile=False)
     def paged_attention_triton(
-        self, q, k, v, k_cache, v_cache, k_scale, v_scale, fwd_ctx: ForwardContext
+        self,
+        q,
+        k,
+        v,
+        k_cache,
+        v_cache,
+        k_scale,
+        v_scale,
+        fwd_ctx: ForwardContext,
+        flydsl_seqs: int | None = None,
     ):
         """Decode through ``run_pa_decode``: gluon, or opted-in FlyDSL.
 
         Callers must have cleared ``gluon_decode_over_limit`` in
-        ``_dispatch_decode`` or, on a GPU without gluon, established through
-        ``_flydsl_takes_decode`` that FlyDSL accepts the call.
+        ``_dispatch_decode`` or established through ``_flydsl_takes_decode``
+        that FlyDSL accepts the call; the latter pass the sequence count it
+        returned as ``flydsl_seqs``, so ``run_pa_decode`` does not ask again.
         """
 
         attn_metadata = fwd_ctx.attn_metadata
@@ -723,9 +824,7 @@ def paged_attention_triton(
             max_context_partition_num = 1
             context_partition_size = 128
 
-        if k_scale is not None and k_scale.numel() > 1:
-            k_scale = k_scale.unsqueeze(-1)
-            v_scale = v_scale.unsqueeze(-1)
+        k_scale, v_scale = self._pa_decode_kv_scales(k_scale, v_scale)
 
         compute_type = (
             torch.bfloat16 if self.kv_cache_dtype == "bf16" else aiter.dtypes.fp8
@@ -743,8 +842,8 @@ def paged_attention_triton(
             context_partition_size=context_partition_size,
             compute_type=compute_type,
             q_scale=None,
-            k_scale=None if self.kv_cache_dtype == "bf16" else k_scale,
-            v_scale=None if self.kv_cache_dtype == "bf16" else v_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
             alibi_slopes=None,
             sinks=self.sinks,
             sliding_window=self.sliding_window,
@@ -755,6 +854,7 @@ def paged_attention_triton(
             # max_context_length: the block-table bound is the same in graph
             # and eager, so both pick the same kernel specialization.
             work_plan=getattr(attn_metadata, "flydsl_work_plan", None),
+            flydsl_seqs=flydsl_seqs,
         )
 
         return o
@@ -1086,9 +1186,10 @@ def prefill_attention_triton(
     def _dispatch_decode(self, max_qlen: int = 1):
         """Select the coarse decode runner from model shape and configuration.
 
-        Cache-dependent dispatch happens later: on a GPU without gluon,
-        ``paged_attention_unified`` hands FlyDSL calls inside its supported
-        domain and keeps the rest on unified attention. The vLLM bridge shares
+        Cache-dependent dispatch happens later: on a GPU without gluon, or
+        past the gluon envelope on one with it, ``paged_attention_unified``
+        hands FlyDSL calls inside its supported domain and keeps the rest on
+        unified attention. The vLLM bridge shares
         the constants and predicate, not this order -- its arms differ (no
         unified fallback, no sliding-window arm). Both
         paged kernels stop short of unified, at different points, and drafting