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

atom/model_ops/attention_mha.py

@@ -19,6 +19,7 @@
     cp_mha_gather_cache,
     dense_decode_splits,
     flydsl_decode_seqs,
+    flydsl_decode_static_ok,
     gluon_decode_over_limit,
     run_pa_decode,
     run_pa_fwd_asm,
@@ -701,7 +702,9 @@ def paged_attention_triton(
         """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
+        ``_dispatch_decode``, or established there through
+        ``flydsl_decode_static_ok`` that FlyDSL serves this layer past that
+        limit, or, on a GPU without gluon, established through
         ``_flydsl_takes_decode`` that FlyDSL accepts the call.
         """
 
@@ -755,6 +758,10 @@ 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),
+            # Past the gluon envelope only FlyDSL may run it; see _dispatch_decode.
+            require_flydsl=gluon_decode_over_limit(
+                attn_metadata.max_seqlen_q, num_q_heads_total, num_kv_heads
+            ),
         )
 
         return o
@@ -1107,6 +1114,23 @@ def _dispatch_decode(self, max_qlen: int = 1):
         over_gluon = not pa_decode_gluon_supports_arch() or gluon_decode_over_limit(
             max_qlen, self.num_heads, self.num_kv_heads
         )
+        # FlyDSL takes any query length, so a layer it will serve stays on the
+        # paged path past the gluon envelope (Qwen3.5 MTP-5: qlen 6, group 8).
+        # Only evaluated when the limit is actually hit, so every within-limit
+        # route -- including MTP-3 -- is decided exactly as before. A GPU
+        # without gluon already reaches FlyDSL through paged_attention_unified.
+        flydsl_past_gluon = False
+        if over_gluon and pa_decode_gluon_supports_arch() and not self.use_flash_layout:
+            flydsl_past_gluon = flydsl_decode_static_ok(
+                head_dim=self.head_dim,
+                num_heads=self.num_heads,
+                num_kv_heads=self.num_kv_heads,
+                kv_cache_dtype=self.kv_cache_dtype,
+                block_size=get_current_atom_config().kv_cache_block_size,
+                sliding_window=self.sliding_window,
+                has_sinks=self.sinks is not None,
+            )
+            over_gluon = not flydsl_past_gluon
         over_asm = (
             max_qlen * (self.num_heads // self.num_kv_heads)
             > PA_ASM_MAX_QUERY_GROUP_SIZE
@@ -1138,7 +1162,10 @@ def _dispatch_decode(self, max_qlen: int = 1):
 
         if wants_unified:
             return self.paged_attention_unified
-        if self.use_triton_attn:
+        # Straight to the FlyDSL-capable runner, not through the ASM arm: the
+        # ASM envelope is a group bound only and this query length is past
+        # anything ASM was ever routed here with.
+        if self.use_triton_attn or flydsl_past_gluon:
             return self.paged_attention_triton
 
         # use_pa_decode_bf16_asm() requires ATOM_USE_UNIFIED_ATTN, which the

atom/model_ops/base_attention.py

@@ -360,6 +360,49 @@ def _flydsl_pa_decode_num_seqs(
     return num_seqs
 
 
+def flydsl_decode_static_ok(
+    *,
+    head_dim: int,
+    num_heads: int,
+    num_kv_heads: int,
+    kv_cache_dtype: str,
+    block_size: int,
+    sliding_window: int,
+    has_sinks: bool,
+) -> bool:
+    """Whether FlyDSL will take a layer's dense paged decode, from what is known
+    before the call.
+
+    The layer-static half of ``_flydsl_pa_decode_num_seqs`` as
+    ``paged_attention_triton`` calls it (fp8 KV, page-16 cache with dim -2 the
+    block size, no alibi, no q_scale, ps=True, a 256 tile and at most
+    PA_DENSE_SPLIT_MAX partitions). It is what lets ``_dispatch_decode`` keep a
+    query length past the gluon envelope on the paged path: FlyDSL takes any
+    positive query length. The per-call clauses (q dtype, contiguity, the
+    row/sequence rectangle, batch cap) cannot be known here, so
+    ``run_pa_decode(require_flydsl=True)`` refuses rather than handing such a
+    call to gluon out of its envelope.
+    """
+    if not envs.ATOM_PA_FLYDSL:
+        return False
+    if has_sinks or sliding_window > 0:
+        return False
+    if kv_cache_dtype != "fp8":
+        return False
+    if block_size not in _FLYDSL_PA_BLOCK_SIZES:
+        return False
+    if not (head_dim == 64 or (head_dim % 128 == 0 and head_dim <= 1024)):
+        return False
+    if num_kv_heads < 1 or num_heads % num_kv_heads != 0:
+        return False
+    # paged_attention_triton's dense tile and split ceiling; both are fixed.
+    if _FLYDSL_PA_TILE != 256 or PA_DENSE_SPLIT_MAX > _FLYDSL_PA_MAX_PARTITIONS:
+        return False
+    if not _flydsl_arch_supported():
+        return False
+    return _flydsl_explicit_plan_api()
+
+
 _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
@@ -644,6 +687,7 @@ def run_pa_decode(
     plan_step_owner=None,
     plan_cache_owner=None,
     plan_partition_cap: int | None = None,
+    require_flydsl: bool = False,
 ):
     """Run the AITER paged-attention decode kernel.
 
@@ -659,6 +703,11 @@ def run_pa_decode(
     ``exp_sums`` / ``max_logits`` / ``temporary_output`` are gluon scratch.
     Omitted, ``pa_decode_gluon`` sizes and allocates its own, so a FlyDSL
     call, which uses its plan's cached scratch, allocates none.
+
+    ``require_flydsl``: the caller dispatched here only because FlyDSL takes a
+    query length gluon does not (``flydsl_decode_static_ok``); on a GPU with
+    gluon, a call FlyDSL then rejects raises instead of falling back to gluon
+    outside its envelope.
     """
     flydsl_seqs = flydsl_decode_seqs(
         output=output,
@@ -802,6 +851,12 @@ def run_pa_decode(
             else "FlyDSL paged decode does not take this call and this GPU "
             f"has no gluon fallback (q {tuple(q.shape)}, {compute_type})"
         )
+    if require_flydsl:
+        raise NotImplementedError(
+            f"pa_decode: max_seqlen_q={max_seqlen_q} is past the gluon envelope "
+            f"and FlyDSL rejected this call (rows={q.shape[0]}, "
+            f"padded_seqs={context_lens.shape[0]}, q dtype {q.dtype})"
+        )
     return torch.ops.aiter.pa_decode_gluon(
         output,
         q,

tests/test_paged_attention_dispatch.py

@@ -21,6 +21,7 @@
 import importlib.util
 import sys
 from types import SimpleNamespace
+from typing import ClassVar
 
 import pytest
 
@@ -295,10 +296,23 @@ class _Layer:
     through once already.
     """
 
-    def __init__(self, sliding_window=-1, num_heads=16, num_kv_heads=1, **flags):
+    def __init__(
+        self,
+        sliding_window=-1,
+        num_heads=16,
+        num_kv_heads=1,
+        head_dim=128,
+        kv_cache_dtype="fp8",
+        sinks=None,
+        **flags,
+    ):
         self.sliding_window = sliding_window
         self.num_heads = num_heads
         self.num_kv_heads = num_kv_heads
+        # Read only past the gluon envelope, by the FlyDSL static gate.
+        self.head_dim = head_dim
+        self.kv_cache_dtype = kv_cache_dtype
+        self.sinks = sinks
         self.use_triton_attn = flags.get("use_triton_attn", False)
         self.use_flash_layout = flags.get("use_flash_layout", False)
         for name in (
@@ -310,11 +324,26 @@ def __init__(self, sliding_window=-1, num_heads=16, num_kv_heads=1, **flags):
             setattr(self, name, name)
 
 
-def _route(monkeypatch, max_qlen, block_size=128, unified=False, force=False, **kw):
+def _route(
+    monkeypatch,
+    max_qlen,
+    block_size=128,
+    unified=False,
+    force=False,
+    flydsl=False,
+    flydsl_arch=True,
+    **kw,
+):
     from atom.model_ops import attention_mha as mha
+    from atom.model_ops import base_attention as ba
 
     monkeypatch.setattr(mha.envs, "ATOM_USE_UNIFIED_ATTN", unified)
     monkeypatch.setattr(mha.envs, "ATOM_FORCE_ATTN_TRITON", force)
+    # Pinned, not inherited: the container sets ATOM_PA_FLYDSL=1, and the two
+    # capability probes below need a GPU / the aiter API, not this test.
+    monkeypatch.setattr(ba.envs, "ATOM_PA_FLYDSL", flydsl)
+    monkeypatch.setattr(ba, "_flydsl_arch_supported", lambda: flydsl_arch)
+    monkeypatch.setattr(ba, "_flydsl_explicit_plan_api", lambda: True)
     monkeypatch.setattr(
         mha,
         "get_current_atom_config",
@@ -401,6 +430,120 @@ def test_persistent_asm_is_not_bounded_by_the_run_pa_fwd_envelope(
         )
 
 
+class TestFlyDSLPastGluonRouting:
+    """Past the gluon envelope a FlyDSL-served layer stays on the paged path.
+
+    Qwen3.5 TP4 MTP-5: 8 q heads / 1 kv head per rank, head_dim 256, fp8 KV,
+    qlen 6 -> 48 rows. FlyDSL takes any query length; unified cannot carry the
+    per-token KV scales and raises at capture.
+    """
+
+    QWEN35: ClassVar[dict] = {
+        "num_heads": 8,
+        "num_kv_heads": 1,
+        "head_dim": 256,
+        "use_triton_attn": True,
+    }
+
+    @pytest.mark.parametrize("max_qlen", [5, 6, 8])
+    def test_qwen35_mtp5_takes_paged_with_flydsl(self, monkeypatch, max_qlen):
+        assert (
+            _route(monkeypatch, max_qlen, block_size=16, flydsl=True, **self.QWEN35)
+            == "paged_attention_triton"
+        )
+
+    @pytest.mark.parametrize("max_qlen", [1, 2, 4])
+    def test_within_gluon_unchanged_by_flydsl(self, monkeypatch, max_qlen):
+        """MTP-3 (qlen 4) and the draft steps (qlen 1) route identically."""
+        for kw in (self.QWEN35, {}):
+            assert _route(
+                monkeypatch, max_qlen, block_size=16, flydsl=True, **kw
+            ) == _route(monkeypatch, max_qlen, block_size=16, flydsl=False, **kw)
+
+    def test_flydsl_off_still_unified(self, monkeypatch):
+        assert (
+            _route(monkeypatch, 6, block_size=16, flydsl=False, **self.QWEN35)
+            == "paged_attention_unified"
+        )
+
+    @pytest.mark.parametrize(
+        "override",
+        [
+            {"block_size": 256},
+            {"block_size": 1024},
+            {"kv_cache_dtype": "bf16"},
+            {"sinks": object()},
+            {"sliding_window": 128},
+            {"head_dim": 192},
+            {"use_flash_layout": True},
+            {"flydsl_arch": False},
+        ],
+    )
+    def test_non_flydsl_configs_still_unified(self, monkeypatch, override):
+        kw = {**self.QWEN35, "block_size": 16, "flydsl": True, **override}
+        assert _route(monkeypatch, 6, **kw) == "paged_attention_unified"
+
+    def test_unified_env_still_wins(self, monkeypatch):
+        assert (
+            _route(
+                monkeypatch, 6, block_size=16, flydsl=True, unified=True, **self.QWEN35
+            )
+            == "paged_attention_unified"
+        )
+
+    def test_no_gluon_gpu_keeps_unified_first(self, monkeypatch):
+        """Without gluon (gfx1250) paged_attention_unified already hands FlyDSL
+        its calls per call; the static shortcut must not bypass that."""
+        from atom.model_ops import attention_mha as mha
+
+        monkeypatch.setattr(mha, "pa_decode_gluon_supports_arch", lambda: False)
+        assert (
+            _route(monkeypatch, 6, block_size=16, flydsl=True, **self.QWEN35)
+            == "paged_attention_unified"
+        )
+
+    def test_past_gluon_with_flydsl_skips_asm(self, monkeypatch):
+        """qlen 5 x group 3 = 15 is inside ASM's group bound; it must still
+        go to the FlyDSL-capable runner, never run_pa_fwd_asm at qlen 5."""
+        assert (
+            _route(
+                monkeypatch, 5, block_size=16, flydsl=True, num_heads=3, num_kv_heads=1
+            )
+            == "paged_attention_triton"
+        )
+
+    def test_run_pa_decode_refuses_gluon_past_envelope(self, monkeypatch):
+        """require_flydsl: a call FlyDSL rejects raises, never reaches gluon."""
+        import torch
+
+        from atom.model_ops import base_attention as ba
+
+        monkeypatch.setattr(ba.envs, "ATOM_PA_FLYDSL", True)
+        monkeypatch.setattr(ba, "_flydsl_pa_decode_num_seqs", lambda **k: None)
+        q = torch.empty(12, 8, 256)
+        with pytest.raises(NotImplementedError, match="gluon envelope"):
+            ba.run_pa_decode(
+                output=q,
+                q=q,
+                k_cache=None,
+                v_cache=None,
+                context_lens=torch.empty(2, dtype=torch.int32),
+                block_tables=None,
+                softmax_scale=1.0,
+                max_seqlen_q=6,
+                max_context_partition_num=1,
+                context_partition_size=256,
+                compute_type=None,
+                q_scale=None,
+                k_scale=None,
+                v_scale=None,
+                exp_sums=None,
+                max_logits=None,
+                temporary_output=None,
+                require_flydsl=True,
+            )
+
+
 class _FakePlan:
     """Just the fields the op and the scratch helper read off a real plan."""