[gfx1250] feat(mla): unfused fallback for gather_kv_b_proj

atom/model_ops/attention_mla.py

@@ -498,6 +498,11 @@ def _maybe_view_mxfp4_weight_for_gather(
     from atom.model_ops.utils import quark_post_load_weights
 
 
+# Unfused torch fallback for `_kv_b_proj_gather`, gated by
+# ATOM_UNFUSED_GATHER_KV_B_PROJ. Pure torch, so unlike the two kernel backends
+# it is always importable and needs no availability flag.
+from atom.model_ops.mla_unfused_gather import unfused_gather_kv_b_proj
+
 # Optional flydsl backend for `_kv_b_proj_gather`, gated by
 # ATOM_USE_FLYDSL_GATHER_KV_B_PROJ.
 try:
@@ -706,6 +711,21 @@ def __init__(
         # kernels with an unpadded 576-wide q_out. The triton path never uses seg.
         self.use_seg_mla = (not self.use_triton_mla) and envs.ATOM_MLA_PAGE_SIZE > 1
         self.use_flydsl_gather_kv_b_proj = bool(envs.ATOM_USE_FLYDSL_GATHER_KV_B_PROJ)
+        # Unfused torch gather: an escape hatch for targets where Triton cannot
+        # compile the fused kernel for chunked-prefill shapes. Refuse the one
+        # combination it cannot serve here, at construction, rather than from
+        # inside the first chunked prefill.
+        self.use_unfused_gather_kv_b_proj = bool(envs.ATOM_UNFUSED_GATHER_KV_B_PROJ)
+        if (
+            self.use_unfused_gather_kv_b_proj
+            and self.use_triton_mla
+            and envs.ATOM_USE_TRITON_MLA_SHUFFLE_KV
+        ):
+            raise RuntimeError(
+                "ATOM_UNFUSED_GATHER_KV_B_PROJ=1 cannot read the shuffled-KV "
+                "cache layout. Set ATOM_USE_TRITON_MLA_SHUFFLE_KV=0, or unset "
+                "ATOM_UNFUSED_GATHER_KV_B_PROJ."
+            )
         # Resolved on the first gather, when the weights and cache exist; see
         # `_kv_b_proj_gather`. None = not asked yet.
         self._flydsl_gather_ok: bool | None = None
@@ -1695,7 +1715,14 @@ def _kv_b_proj_gather(
                     return k_gather, v_gather, *kv_out_scales
                 return None
 
-        gather_kv_b_proj(
+        # Checked after FlyDSL: where the fused kernel compiles at all, it wins.
+        # This path trades throughput for existing.
+        gather = (
+            unfused_gather_kv_b_proj
+            if self.use_unfused_gather_kv_b_proj
+            else gather_kv_b_proj
+        )
+        gather(
             kv_buffer,
             self._k_scale,
             kv_indptr,

atom/model_ops/mla_unfused_gather.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: MIT
+"""Unfused torch fallback for aiter's Triton ``gather_kv_b_proj``.
+
+The fused kernel does the whole MLA cached-prefix expansion in one launch: row
+gather, KV-cache dequant, ``kv_b_proj``, the k_nope/v split and the k_pe concat.
+On some architectures Triton has no working backend for the shapes a *chunked*
+prefill produces and codegen asserts inside LLVM before any kernel runs::
+
+    llvm/IR/Instructions.h: PHINode::setIncomingValue:
+    Assertion `getType() == V->getType()' failed.
+
+That is a compile-time failure with no fallback: the process aborts, so a long
+prompt at concurrency takes the engine down. The FlyDSL gather is not an
+alternative on such a target -- it is gfx950-only and aborts in its own
+compiler if forced.
+
+Every step of the chain is an ordinary torch op, so an unfused version needs no
+codegen and cannot miscompile. It is slower and allocates per chunk; it exists
+to keep a bring-up target serving, not to compete with the fused kernel, and is
+off by default (``ATOM_UNFUSED_GATHER_KV_B_PROJ``).
+
+What it deliberately does NOT implement, raising instead of guessing: the
+shuffled-KV cache layout and a preshuffled weight (both private layouts of the
+fused kernel) and MXFP4 ``kv_b_proj`` weights. Producing subtly wrong K/V would
+be far worse than refusing -- nothing downstream would notice.
+"""
+
+from __future__ import annotations
+
+import torch
+
+__all__ = ["unfused_gather_kv_b_proj"]
+
+# Rows per chunk. Bounds the temporary [rows, num_heads * (nope + v_dim)] GEMM
+# output; the whole prefix at once is a large transient on long contexts.
+_DEFAULT_ROWS = 4096
+
+
+def _mm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
+    """``x @ w.T`` for a ``[N, K]`` weight, routed away from hipBLASLt.
+
+    ``torch.matmul`` lands on hipBLASLt, i.e. a Tensile ``Cijk_*_SAV_UserArgs_``
+    kernel that reads its pointers from a device-side argument buffer. Combined
+    with hipGraph replay and prefix caching that page-faults. aiter's Triton
+    a16w16 GEMM is correct here and much faster on this shape. Called directly
+    rather than through ``tgemm.mm`` so a tuned-CSV row cannot route it back.
+    """
+    try:
+        from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16
+    except Exception:  # noqa: BLE001 -- optional kernel; torch is the fallback
+        return torch.matmul(x, w.t())
+    # gemm_a16w16 reads N from w.shape[0] and returns [M, N].
+    return gemm_a16w16(x, w, bias=None, dtype=x.dtype)
+
+
+def _dequant_weight(
+    weight: torch.Tensor, scale: torch.Tensor | None, dtype: torch.dtype
+) -> torch.Tensor:
+    """Dequantize ``kv_b_proj`` once per call, matching the fused kernel.
+
+    Three shapes of scale, all of which appear in shipped checkpoints: none at
+    all (a bf16 ``kv_b_proj``, e.g. Kimi-K3), one per output row, and the
+    128x128 block scale DeepSeek ships.
+    """
+    if scale is None:
+        return weight if weight.dtype == dtype else weight.to(dtype)
+    scale = _scale_as_float(scale)
+    if scale.dim() == 1 or (scale.dim() == 2 and scale.shape[1] == 1):
+        return (weight.to(torch.float32) * scale.reshape(-1, 1)).to(dtype)
+    scale_n, scale_k = scale.shape
+    n, k = weight.shape
+    blocked = weight.to(torch.float32).view(
+        scale_n, n // scale_n, scale_k, k // scale_k
+    )
+    return (blocked * scale[:, None, :, None]).reshape(n, k).to(dtype)
+
+
+def _scale_as_float(scale: torch.Tensor) -> torch.Tensor:
+    """Weight scale as fp32 multipliers.
+
+    An E8M0 block scale is ``torch.float8_e8m0fnu`` where torch has it, and
+    ``.to(float32)`` decodes that. Without it aiter's ``dtypes.fp8_e8m0`` falls
+    back to ``uint8`` holding the biased exponent, which a plain cast would read
+    as 127 instead of 1.0 -- so decode those bytes explicitly.
+    """
+    if scale.dtype == torch.uint8:
+        return torch.exp2(scale.to(torch.float32) - 127.0)
+    return scale.to(torch.float32)
+
+
+def _row_addresses(
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    total_kv: int,
+    block_size: int,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Map each output row to its ``(block, slot)`` in the paged cache.
+
+    Vectorized on purpose: a per-sequence Python loop would force a host sync
+    per request in the prefill hot path. ``page_size`` 1 needs no arithmetic at
+    all; the paged form walks ``kv_indptr`` / ``cu_seqlens_k`` the way the fused
+    kernel's non-flat grid does.
+    """
+    if block_size == 1:
+        return kv_indices[:total_kv].long(), None
+    token = torch.arange(total_kv, device=device)
+    # cu_seqlens_k = [0, n0, n0 + n1, ...]: a token's sequence is the last entry
+    # that is <= it.
+    seq = torch.searchsorted(cu_seqlens_k[1:].contiguous(), token, right=True)
+    within = token - cu_seqlens_k[seq]
+    block = kv_indices[(kv_indptr[seq] + within // block_size).long()].long()
+    return block, (within % block_size).long()
+
+
+def unfused_gather_kv_b_proj(
+    k_buffer: torch.Tensor,
+    k_scale: torch.Tensor | None,
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    kv_prefix_sum_context_lens: torch.Tensor,
+    kv_proj_weight: torch.Tensor,
+    kv_proj_scale: torch.Tensor | None,
+    k_prefix: torch.Tensor,
+    v_prefix: torch.Tensor,
+    weight_preshuffle: bool = False,
+    shuffled_kv_cache: bool = False,
+    rows_per_chunk: int = _DEFAULT_ROWS,
+) -> None:
+    """Drop-in replacement for ``aiter.ops.triton.gather_kv_b_proj``.
+
+    Writes into ``k_prefix`` / ``v_prefix`` in place and returns None, exactly
+    as the fused op does.
+    """
+    if shuffled_kv_cache:
+        raise NotImplementedError(
+            "ATOM_UNFUSED_GATHER_KV_B_PROJ does not implement the shuffled-KV "
+            "cache layout (the intra-block shuffle is the fused kernel's "
+            "private layout). Set ATOM_USE_TRITON_MLA_SHUFFLE_KV=0."
+        )
+    if weight_preshuffle:
+        raise NotImplementedError(
+            "ATOM_UNFUSED_GATHER_KV_B_PROJ cannot read a preshuffled "
+            "kv_b_proj weight (is_shuffled=True)."
+        )
+    # Packed MXFP4 arrives as float4_e2m1fn_x2 where torch has that dtype, and
+    # as raw uint8 where it does not (`_maybe_view_mxfp4_weight_for_gather`
+    # leaves it alone). No real kv_b_proj weight is uint8, so refuse both.
+    fp4 = getattr(torch, "float4_e2m1fn_x2", None)
+    if kv_proj_weight.dtype == torch.uint8 or (
+        fp4 is not None and kv_proj_weight.dtype == fp4
+    ):
+        raise NotImplementedError(
+            "ATOM_UNFUSED_GATHER_KV_B_PROJ does not implement MXFP4 kv_b_proj "
+            "weights; unset the env for this checkpoint."
+        )
+
+    _, block_size, hidden = k_buffer.shape
+    total_kv, num_heads, k_width = k_prefix.shape
+    v_dim = v_prefix.shape[2]
+    kv_c_dim = kv_proj_weight.shape[1]
+    pe_dim = hidden - kv_c_dim
+    nope = k_width - pe_dim
+    dtype = k_prefix.dtype
+
+    weight = _dequant_weight(kv_proj_weight, kv_proj_scale, dtype)
+    block, slot = _row_addresses(
+        kv_indptr,
+        kv_indices,
+        kv_prefix_sum_context_lens,
+        total_kv,
+        block_size,
+        k_prefix.device,
+    )
+
+    # A bf16 cache must NOT be scaled. The fused kernel hardcodes exactly this
+    # (its k_scalar_scale is 1.0 for a bf16 cache and a load of k_scale
+    # otherwise); applying k_scale anyway scales all of K and V by a constant,
+    # which nothing downstream would flag. The GEMM is linear in kv_c, so
+    # scaling after the matmul is exact.
+    scale_cache = k_scale is not None and k_buffer.dtype is not torch.bfloat16
+    if scale_cache:
+        # MLAAttention keeps _k_scale as a host tensor. A 0-dim one would still
+        # broadcast against CUDA operands, but a [1] one would not (and would
+        # promote the bf16 products to fp32); make it a 0-dim device scalar once
+        # rather than depend on its shape in every chunk.
+        k_scale = k_scale.reshape(()).to(k_prefix.device, non_blocking=True)
+
+    for start in range(0, total_kv, rows_per_chunk):
+        stop = min(start + rows_per_chunk, total_kv)
+        rows = block[start:stop]
+        latent = (
+            k_buffer[rows, 0] if slot is None else k_buffer[rows, slot[start:stop]]
+        ).to(dtype)
+
+        projected = _mm(latent[:, :kv_c_dim], weight)
+        k_pe = latent[:, kv_c_dim:]
+        if scale_cache:
+            projected = projected * k_scale
+            k_pe = k_pe * k_scale
+        projected = projected.view(stop - start, num_heads, nope + v_dim)
+
+        k_prefix[start:stop, :, :nope] = projected[:, :, :nope]
+        k_prefix[start:stop, :, nope:] = k_pe.unsqueeze(1).expand(-1, num_heads, -1)
+        v_prefix[start:stop] = projected[:, :, nope:]

atom/utils/envs.py

@@ -201,6 +201,13 @@ def _positive_float_env(name: str, default: str) -> float:
     "ATOM_USE_FLYDSL_GATHER_KV_B_PROJ": lambda: (
         os.getenv("ATOM_USE_FLYDSL_GATHER_KV_B_PROJ", "1") == "1"
     ),
+    # gather_kv_b_proj: replace the Triton op with an unfused torch chain (row
+    # gather -> dequant -> GEMM -> split/concat). Escape hatch for targets where
+    # Triton cannot compile the fused kernel for chunked-prefill shapes and
+    # aborts the process; slower, and off unless asked for.
+    "ATOM_UNFUSED_GATHER_KV_B_PROJ": lambda: (
+        os.getenv("ATOM_UNFUSED_GATHER_KV_B_PROJ", "0") == "1"
+    ),
     # FlyDSL FP8 prefill with fused QKV quantization and direct FP8 gather output
     # where supported. Unsupported attention inputs raise. Added 2026-09-10.
     "ATOM_USE_FLYDSL_FP8_PREFILL_ATTN": lambda: (

tests/test_envs.py

@@ -37,6 +37,7 @@
     "ATOM_USE_CUSTOM_ALL_GATHER",
     "ATOM_ENABLE_RELAXED_MTP",
     "ATOM_USE_FLYDSL_GATHER_KV_B_PROJ",
+    "ATOM_UNFUSED_GATHER_KV_B_PROJ",
     "ATOM_USE_FLYDSL_FP8_PREFILL_ATTN",
 ]
 
@@ -128,6 +129,10 @@ def test_atom_enable_gdn_decode_lossy_fast_default(self):
     def test_use_flydsl_gather_kv_b_proj_default(self):
         assert _get_envs().ATOM_USE_FLYDSL_GATHER_KV_B_PROJ is True
 
+    def test_unfused_gather_kv_b_proj_default(self):
+        """Off: the fused kernel is the right answer wherever it compiles."""
+        assert _get_envs().ATOM_UNFUSED_GATHER_KV_B_PROJ is False
+
     def test_use_flydsl_fp8_prefill_attn_default(self):
         assert _get_envs().ATOM_USE_FLYDSL_FP8_PREFILL_ATTN is False
 
@@ -242,6 +247,14 @@ def test_use_flydsl_gather_kv_b_proj_only_one_enables(self, monkeypatch):
         monkeypatch.setenv("ATOM_USE_FLYDSL_GATHER_KV_B_PROJ", "true")
         assert _get_envs().ATOM_USE_FLYDSL_GATHER_KV_B_PROJ is False
 
+    def test_unfused_gather_kv_b_proj_enabled(self, monkeypatch):
+        monkeypatch.setenv("ATOM_UNFUSED_GATHER_KV_B_PROJ", "1")
+        assert _get_envs().ATOM_UNFUSED_GATHER_KV_B_PROJ is True
+
+    def test_unfused_gather_kv_b_proj_only_one_enables(self, monkeypatch):
+        monkeypatch.setenv("ATOM_UNFUSED_GATHER_KV_B_PROJ", "true")
+        assert _get_envs().ATOM_UNFUSED_GATHER_KV_B_PROJ is False
+
 
 class TestIsSet:
     """Test the is_set() helper function."""

tests/test_mla_flydsl_gather_dispatch.py

@@ -31,6 +31,7 @@ def _impl(**over):
         "kv_b_proj": SimpleNamespace(weight=SimpleNamespace(is_shuffled=True)),
         "_k_scale": None,
         "use_flydsl_gather_kv_b_proj": True,
+        "use_unfused_gather_kv_b_proj": False,
         "_flydsl_gather_ok": None,
         "_flydsl_gather_fp8_ok": None,
     }
@@ -93,6 +94,34 @@ def test_env_off_never_asks(spies):
     assert spies.asked == 0
 
 
+@pytest.mark.parametrize("flydsl_ok, flydsl, unfused", [(True, 1, 0), (False, 0, 1)])
+def test_unfused_replaces_only_the_triton_op(
+    spies, monkeypatch, flydsl_ok, flydsl, unfused
+):
+    """ATOM_UNFUSED_GATHER_KV_B_PROJ is the last resort, not the first.
+
+    It is a torch chain that exists so a target whose Triton cannot compile the
+    fused kernel can still serve. Where a kernel backend works, it wins -- so the
+    env must displace the Triton op and nothing above it.
+    """
+    calls = []
+    monkeypatch.setattr(
+        mla, "unfused_gather_kv_b_proj", lambda *a, **k: calls.append(a)
+    )
+    spies.answer = flydsl_ok
+    _gather(_impl(use_unfused_gather_kv_b_proj=True))
+    assert (spies.flydsl, len(calls), spies.triton) == (flydsl, unfused, 0)
+
+
+def test_unfused_off_keeps_the_triton_op(spies, monkeypatch):
+    monkeypatch.setattr(
+        mla, "unfused_gather_kv_b_proj", Mock(side_effect=AssertionError("unexpected"))
+    )
+    spies.answer = False
+    _gather(_impl())
+    assert spies.triton == 1
+
+
 @pytest.mark.parametrize("supported", [False, True])
 def test_fp8_gather_checks_output_capability_once(monkeypatch, supported):
     impl = _impl()