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:
@@ -608,19 +615,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.
+    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_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 and one info
+        line per process says that the flag would route it to FlyDSL.
 
         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 (
             envs.ATOM_USE_UNIFIED_ATTN
-            or pa_decode_gluon_supports_arch()
+            or self.use_flash_layout
             or self.sliding_window > 0
         ):
             return False
-        num_seqs = attn_metadata.context_lens.shape[0]
-        return bool(
+        if pa_decode_gluon_supports_arch() and not gluon_decode_over_limit(
+            attn_metadata.max_seqlen_q, self.num_heads, self.num_kv_heads
+        ):
+            return False
+        if not envs.ATOM_PA_FLYDSL:
+            global _flydsl_off_hint_logged
+            if 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,
+                )
+            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 +678,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,9 +706,20 @@ 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,
             )
         self._reject_per_token_scales_on_unified(k_scale, v_scale, "decode")
 
@@ -696,13 +756,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 +793,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 +811,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 +823,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 +1155,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