feat(h20): FlagGems' Triton sparse MLA matches the TileLang kernel at V4.1's shape

The remaining blocker for a TileLang-free path was sparse MLA: `h20_ops/sparse_mla_sm90.py` is a
231-line TileLang program with no Triton counterpart in this repo. FlagGems already has one, and it
turns out to be drop-in.

`flagos-ai/FlagGems` -> `src/flag_gems/fused/flashmla_sparse.py` -> `flash_mla_sparse_fwd`. Pure
Triton. At V4.1's real dims (h_q=64, h_kv=1, d_qk=512, topk=512) it compiles and runs on H20, and
on 7 shapes it agrees with the TileLang kernel within 2 bf16 ULP -- including `topk_length` and
`attn_sink`, the two optional arguments.

The pairing is closer than expected, which is why it needed no adaptation:

    prefill_attn_varlen_sparse(query, kv, index_topk, softmax_scale, attn_sink, index_length)
    flash_mla_sparse_fwd(q,         kv, indices,    sm_scale,     attn_sink, topk_length)

Causality matches too, and for the same reason. `attn_sm90.py` passes `is_causal=False` because the
indexer only scores positions before each query, so the indices already carry causality -- bounding
on the row index instead collapses the layer to NaN, which already happened once and is documented
at that call site. The FlagGems kernel has no causal parameter either.

Recorded as *not* the answer: `DSA/sparse_mla.py` in the same repo. It hardcodes causal and needs
the V3.2 576-wide split layout; V4.1 folds RoPE inside a single 512-wide latent, so TD = 0 and it
fails to compile on `tl.arange(0, 0)`. Worth writing down because the name looks like the better
match.

What remains, all much smaller than writing a kernel: wire it in behind the existing
`patch_imports` hook; test the decode path separately (it gathers to bf16 with identity indices, a
different call shape than the 7 cases); and convert the return, since the kernel gives
(out, max_logits, lse) with lse in natural log where DeepInfer expects (out, lse) with lse in log2
-- a 1.4427x scale error if dropped.

TODO.md

@@ -630,13 +630,45 @@ quant/MoE/MHC/engram，25 处模块级 import / 23 个 API）。
 4. prompt B 第 1 个 token 就分叉，暗示 prefill 阶段已存在实质差异，需要对比首个 token 的 logits
    才能定位，而不是继续猜。
 
-### 6.6 还差什么才能真正脱离 tilelang
+### 6.6 稀疏 MLA 的 Triton 方案已找到，且已数值验证（2026-10-25）
+
+原判断「SM90 稀疏 MLA 没有 Triton 版本、工作量远大于 compat 层」**需要修正**：
+FlagGems 里已经有，而且是 drop-in 的。
+
+**`flagos-ai/FlagGems` → `src/flag_gems/fused/flashmla_sparse.py` → `flash_mla_sparse_fwd`**
+纯 Triton，实测在 H20 上编译执行通过，**7 组形状与现有 TileLang kernel 全部一致（≤2 bf16 ULP）**。
+
+参数逐项对应，连两个可选参数都有：
+
+| DeepInfer `prefill_attn_varlen_sparse` | FlagGems `flash_mla_sparse_fwd` |
+|---|---|
+| `query [total_q, H, D]` | `q [s_q, h_q, d]` |
+| `kv_a_with_mqa [total_k, 1, D]` | `kv [s_kv, h_kv, d_qk]` |
+| `index_topk [total_q, topk]` | `indices [s_q, h_kv, topk]` |
+| `softmax_scale` | `sm_scale` |
+| `attn_sink [H]` | `attn_sink [h_q]` |
+| `index_length [total_q]` | `topk_length [s_q]` |
+
+**语义也对得上**：`attn_sm90.py:110` 明确让 prefill 走 `is_causal=False`（索引由 indexer 保证
+causal，用行号当因果边界会导致整层 NaN，这个 bug 已经发生过一次）；FlagGems 这个 kernel 同样
+没有 causal 参数，理由一致。V4.1 的 `h_q=64 / h_kv=1 / d_qk=512` 恰好满足它所有 assert。
+
+**不能用的那个**：同仓库 `DSA/sparse_mla.py`（`triton_sparse_mla_fwd_interface`）硬编码
+`is_causal=True`，且要求 V3.2 的 576 拆分布局（latent 512 + rope tail 64）。V4.1 把 RoPE 并进
+512 维内，`TD = 512-512 = 0`，`tl.arange(0,0)` 直接编译失败。别选错。
+
+验证脚本与 vendor 副本：`scripts/sparse_mla_port/`（`README.md` 有完整对照表与未做项）。
+
+**还差的三步**（都比重写 kernel 小得多）：
+1. **接线**：`sparse_mla_sm90.py` 仍 `import tilelang`。用现成的 `patch_imports` 挂一个 Triton
+   后端即可，不动上游调用点。
+2. **decode 路径单独测**：今天 decode 是 gather→bf16 + identity indices，与上面 7 组形状不同，
+   需要对着 fp8 cache 路径另做一次 A/B。
+3. **返回值适配**：它返回 `(out, max_logits, lse)`，DeepInfer 要 `(out, lse)`；且 TileLang 的
+   `lse` 是 log2 单位、FlagGems 是自然对数，**必须转换**，否则下游读到的 scale 差 1.4427 倍。
 
-`h20_ops/sparse_mla_sm90.py` 是 `@tilelang.jit` 实现，且 `h20-sm90-triton` 分支上也是如此
-（该分支只是没屏蔽 tilelang，不是真的无依赖）。**SM90 稀疏 MLA 没有 Triton 版本**，这是
-「国产卡完全不用 tilelang」的最后一个硬缺口，工作量远大于这次的 compat 层。
 `tilelang_ops/cuda/` 那 29 个模块已有 `tilelang_ops/triton/` 对应版本，靠
-`DEEPINFER_OPS_BACKEND=triton` 切换，不是缺口。
+`DEEPINFER_OPS_BACKEND=triton` 切换，本来就不是缺口。
 
 ### 6.4 端到端曾经为什么跑不起来（环境 + 排查记录）
 

scripts/sparse_mla_port/README.md

@@ -0,0 +1,64 @@
+# FlagGems sparse-MLA as the SM90 attention backend
+
+The SM90 sparse MLA kernel (`deepinfer/h20_ops/sparse_mla_sm90.py`) is TileLang-only, so with
+TileLang blocked the model cannot start. FlagGems ships a pure-Triton implementation of the same
+operator, and at DeepSeek-V4.1's exact shape it agrees with the TileLang kernel.
+
+## Source
+
+`flagos-ai/FlagGems`, file `src/flag_gems/fused/flashmla_sparse.py`, function
+`flash_mla_sparse_fwd`. Copied verbatim into `vendor/` for the tests; it depends only on
+`flag_gems.utils.triton_version_utils` (a Triton version probe), stubbed in `vendor/flag_gems/`.
+
+Do not confuse it with `src/flag_gems/fused/DSA/sparse_mla.py`
+(`triton_sparse_mla_fwd_interface`). That one hardcodes `is_causal=True` and requires the V3.2
+576-wide split layout (latent 512 + rope tail 64). V4.1 folds RoPE inside a single 512-wide
+latent, so `TD = 512 - 512 = 0` and the kernel fails to compile on `tl.arange(0, 0)`. It is not a
+candidate.
+
+## Why the pairing works
+
+V4.1's dims satisfy every assert in the FlagGems kernel: `h_q = 64`, `h_kv = 1`, `d_qk = 512`.
+
+The argument lists correspond item for item, including the two optional ones:
+
+| DeepInfer `prefill_attn_varlen_sparse` | FlagGems `flash_mla_sparse_fwd` |
+|---|---|
+| `query [total_q, H, D]` | `q [s_q, h_q, d]` |
+| `kv_a_with_mqa [total_k, 1, D]` | `kv [s_kv, h_kv, d_qk]` |
+| `index_topk [total_q, topk]` | `indices [s_q, h_kv, topk]` |
+| `softmax_scale` | `sm_scale` |
+| `attn_sink [H]` | `attn_sink [h_q]` |
+| `index_length [total_q]` | `topk_length [s_q]` |
+
+Causality also lines up: `attn_sm90.prefill_attn_varlen_sparse` passes `is_causal=False` and
+documents why -- the indexer only scores positions before each query, so the indices carry
+causality, and bounding on the row index instead collapses the layer to NaN (that already happened
+once, see the comment at `attn_sm90.py:110`). The FlagGems kernel has no causal parameter for the
+same reason.
+
+## Result
+
+`ab_vs_tilelang.py`, 7 shapes, TileLang as reference: **all match** within 2 bf16 ULP, including
+`topk_length` and `attn_sink`.
+
+| case | max abs diff | tol |
+|---|---|---|
+| prefill, distinct top-k (SQ=64, SKV=2048) | 1.22e-04 | 1.45e-04 |
+| prefill, duplicate top-k | 1.22e-04 | 1.97e-04 |
+| decode (SQ=1, SKV=4096) | 6.10e-05 | 1.20e-04 |
+| odd SQ (7) | 6.10e-05 | 1.43e-04 |
+| short KV < topk | 2.44e-04 | 2.90e-04 |
+| with topk_length | 2.44e-04 | 3.20e-04 |
+| with attn_sink | 1.10e-04 | 1.41e-04 |
+
+## What is not done yet
+
+* Not wired in. `sparse_mla_sm90.py` still imports TileLang; a Triton backend behind the same
+  `patch_imports` hook is the next step.
+* Decode path untested end to end. Today decode gathers to bf16 and calls the kernel with identity
+  indices and no causal; that is a different call shape than the ones above and needs its own test
+  against the fp8-cache path.
+* The kernel returns `(out, max_logits, lse)` while DeepInfer's returns `(out, lse)`, and TileLang's
+  `lse` is in log2 units where FlagGems' is natural log. A wrapper must drop `max_logits` and
+  convert, or every downstream consumer reads a 1.44x-wrong scale.

scripts/sparse_mla_port/ab_vs_tilelang.py

@@ -0,0 +1,153 @@
+"""FlagGems' `flash_mla_sparse_fwd` vs DeepInfer's TileLang sparse MLA, on the engine's own contract.
+
+Why this pairing and not the DSA kernel: `attn_sm90.prefill_attn_varlen_sparse` calls the TileLang
+kernel with `is_causal=False`, documented in the source -- the indexer only scores positions before
+each query, so the indices carry causality, and bounding on the row index instead collapses the
+layer to NaN (that bug already happened once). `flash_mla_sparse_fwd` has no causal parameter for
+the same reason. The DSA kernel (`triton_sparse_mla_fwd_interface`) hardcodes causal and needs the
+V3.2 576-wide split layout, which V4.1 does not have, so it is not a candidate here.
+
+Argument correspondence, which is closer than expected:
+
+    prefill_attn_varlen_sparse(query, kv, index_topk, softmax_scale,
+                               attn_sink, index_length)
+    flash_mla_sparse_fwd(q,     kv, indices,    sm_scale,
+                         attn_sink=..., topk_length=...)
+
+Both return (out, ...). `index_length` is a per-query valid count applied as -1 padding upstream;
+`topk_length` is the same thing applied inside the kernel.
+"""
+
+import sys
+from pathlib import Path
+
+import torch
+
+# The vendored copy imports `flag_gems.utils...`; its stub sits next to it, so put that on the
+# path here rather than making the caller set PYTHONPATH for a layout detail.
+sys.path.insert(0, str(Path(__file__).resolve().parent / "vendor"))
+
+from vendor.fg_flashmla_sparse import flash_mla_sparse_fwd
+from deepinfer.h20_ops.sparse_mla_sm90 import sparse_mla_fwd_attn
+
+DEV = "cuda"
+HQ, HKV, D, TOPK = 64, 1, 512, 512
+FAILED = []
+
+
+def tl_call(q3, kv3, idx4, scale):
+    """DeepInfer's shape convention: (S,H,D)/(S,1,D)/(S,1,topk) -> pad a batch axis."""
+    s_q = q3.shape[0]
+    s_kv = kv3.shape[0]
+    out, lse = sparse_mla_fwd_attn(
+        q3.view(1, s_q, HQ, D),
+        kv3.view(1, s_kv, 1, D),
+        idx4.view(1, s_q, 1, -1),
+        sm_scale=scale,
+        is_causal=False,
+    )
+    return out.view(s_q, HQ, D), lse.view(s_q, HQ)
+
+
+def compare(tag, a, b, extra=""):
+    af, bf = a.float(), b.float()
+    if af.shape != bf.shape:
+        print(f"  FAIL {tag}: shape {tuple(af.shape)} vs {tuple(bf.shape)}")
+        FAILED.append(tag)
+        return
+    diff = (af - bf).abs()
+    scale = af.abs().max().item()
+    tol = 2.0**-8 * max(scale, 1e-9) * 2
+    finite = bool(torch.isfinite(bf).all())
+    ok = diff.max().item() <= tol and finite
+    print(
+        f"  {'ok  ' if ok else 'FAIL'} {tag}: max|ref| {scale:.4f}  max diff {diff.max().item():.3e}"
+        f"  mean {diff.mean().item():.3e}  tol {tol:.3e}  finite {finite}{extra}"
+    )
+    if not ok:
+        FAILED.append(tag)
+
+
+def make(SQ, SKV, seed, distinct):
+    torch.manual_seed(seed)
+    q = (torch.randn((SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
+    kv = (torch.randn((SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1).contiguous()
+    if distinct and SKV > TOPK:
+        rows = [torch.randperm(SKV, device=DEV)[:TOPK] for _ in range(SQ)]
+        idx = torch.stack(rows).view(SQ, 1, TOPK).to(torch.int32)
+    else:
+        idx = torch.randint(0, SKV, (SQ, 1, TOPK), device=DEV, dtype=torch.int32)
+    return q, kv, idx.contiguous()
+
+
+def case(tag, SQ, SKV, seed, distinct=True, use_length=False, use_sink=False):
+    q, kv, idx = make(SQ, SKV, seed, distinct)
+    scale = D**-0.5
+    print(f"\n[{tag}] SQ={SQ} SKV={SKV} topk={TOPK} distinct={distinct}"
+          f"{' +topk_length' if use_length else ''}{' +attn_sink' if use_sink else ''}")
+
+    # DeepInfer applies `index_length` by masking the tail of the index row to -1 upstream
+    # (`_mask_beyond_length`), so the reference call is the same kernel with a padded index.
+    idx_ref = idx
+    tl_len = None
+    if use_length:
+        torch.manual_seed(seed + 100)
+        tl_len = torch.randint(64, TOPK, (SQ,), device=DEV, dtype=torch.int32)
+        idx_ref = idx.clone()
+        ar = torch.arange(TOPK, device=DEV).view(1, 1, TOPK)
+        idx_ref[(ar >= tl_len.view(SQ, 1, 1)).expand(SQ, 1, TOPK)] = -1
+        idx_ref = idx_ref.contiguous()
+
+    tl_out, _ = tl_call(q, kv, idx_ref, scale)
+    torch.cuda.synchronize()
+
+    sink = None
+    if use_sink:
+        torch.manual_seed(seed + 200)
+        sink = (torch.randn((HQ,), device=DEV, dtype=torch.float32) * 0.5)
+        # DeepInfer's sink is folded into the output after the kernel; feed the kernel no sink and
+        # compare the raw attention, then check the sink formula separately below.
+    fg_out, fg_max, fg_lse = flash_mla_sparse_fwd(
+        q, kv, idx_ref, sm_scale=scale, d_v=D,
+        attn_sink=sink,
+        topk_length=tl_len if use_length else None,
+    )
+    torch.cuda.synchronize()
+
+    if use_sink:
+        # Compare against the reference with the sink applied the same way the kernel documents:
+        # out * exp(lse) / (exp(lse) + exp(sink)).
+        tl_lse = sparse_mla_fwd_attn(
+            q.view(1, SQ, HQ, D), kv.view(1, SKV, 1, D), idx_ref.view(1, SQ, 1, -1),
+            sm_scale=scale, is_causal=False,
+        )[1].view(SQ, HQ)
+        import math
+        w = torch.exp(tl_lse.float() * math.log(2))  # TileLang lse is log2 units
+        # tl_out is (SQ, HQ, D), so the weight must gain a trailing axis to broadcast over D.
+        s = sink.view(1, HQ, 1)
+        gate = w.unsqueeze(-1) / (w.unsqueeze(-1) + torch.exp(s))
+        applied = tl_out.float() * gate
+        compare("sink applied", applied, fg_out)
+    else:
+        compare("out", tl_out, fg_out)
+
+
+CASES = [
+    ("prefill, distinct top-k", dict(SQ=64, SKV=2048, seed=0)),
+    ("prefill, duplicate top-k", dict(SQ=32, SKV=512, seed=1, distinct=False)),
+    ("decode (SQ=1)", dict(SQ=1, SKV=4096, seed=2)),
+    ("odd SQ", dict(SQ=7, SKV=1024, seed=3)),
+    ("short KV (< topk)", dict(SQ=16, SKV=128, seed=4, distinct=False)),
+    ("with topk_length", dict(SQ=32, SKV=2048, seed=5, use_length=True)),
+    ("with attn_sink", dict(SQ=16, SKV=1024, seed=6, use_sink=True)),
+]
+
+if __name__ == "__main__":
+    # Run from the repo root: PYTHONPATH=. python scripts/sparse_mla_port/ab_vs_tilelang.py
+    for tag, kw in CASES:
+        case(tag, **kw)
+    print()
+    if FAILED:
+        print(f"MISMATCH {len(FAILED)}: {FAILED}")
+        sys.exit(1)
+    print(f"ALL_MATCH ({len(CASES)} cases)")

scripts/sparse_mla_port/probe.py

@@ -0,0 +1,55 @@
+"""Can FlagGems' Triton sparse-MLA kernel run on H20 (SM90) at DeepSeek-V4.1's shape?
+
+Deliberately narrow: first prove the kernel compiles and produces finite output, then (in
+sparse_mla_ab.py) compare it against the TileLang kernel DeepInfer uses today. Mixing the two into
+one script makes a compile failure look like a numerical mismatch.
+
+Shape comes from the real config: 64 query heads, 1 kv head, head_dim 512, index_topk 512.
+"""
+
+import sys
+from pathlib import Path
+import traceback
+
+import torch
+
+# The vendored copy imports `flag_gems.utils...`; its stub sits next to it, so put that on the
+# path here rather than making the caller set PYTHONPATH for a layout detail.
+sys.path.insert(0, str(Path(__file__).resolve().parent / "vendor"))
+
+from vendor.fg_flashmla_sparse import flash_mla_sparse_fwd
+
+torch.manual_seed(0)
+DEV = "cuda"
+
+B, SQ, HQ, HKV, D, TOPK, SKV = 1, 8, 64, 1, 512, 512, 64
+
+q = torch.randn((B, SQ, HQ, D), device=DEV, dtype=torch.bfloat16) * 0.1
+kv = torch.randn((B, SKV, HKV, D), device=DEV, dtype=torch.bfloat16) * 0.1
+
+# Valid indices into SKV, no duplicates within a row is not required by the kernel.
+idx = torch.randint(0, SKV, (B, SQ, HKV, TOPK), device=DEV, dtype=torch.int32)
+
+# The kernel takes varlen shapes: fold the batch dim away (B is 1 here).
+q2 = q.view(B * SQ, HQ, D).contiguous()
+kv2 = kv.view(B * SKV, HKV, D).contiguous()
+idx2 = idx.view(B * SQ, HKV, TOPK).contiguous()
+
+print(f"q {tuple(q2.shape)} kv {tuple(kv2.shape)} idx {tuple(idx2.shape)}", flush=True)
+print("h_q=64, h_kv=1, d=512 -- all inside the kernel's asserts", flush=True)
+
+try:
+    out, max_logits, lse = flash_mla_sparse_fwd(
+        q2, kv2, idx2, sm_scale=D**-0.5, d_v=D
+    )
+    torch.cuda.synchronize()
+    print(f"OK   out {tuple(out.shape)} {out.dtype}")
+    print(f"     max_logits {tuple(max_logits.shape)} {max_logits.dtype}")
+    print(f"     lse {tuple(lse.shape)} {lse.dtype}")
+    print(f"     finite: out={torch.isfinite(out).all().item()} "
+          f"lse={torch.isfinite(lse).all().item()}")
+    print(f"     mean |out| {out.float().abs().mean().item():.6f}")
+except Exception:
+    print("FAIL")
+    traceback.print_exc()
+    sys.exit(1)

scripts/sparse_mla_port/vendor/flag_gems/__init__.py

@@ -0,0 +1 @@
+"""Minimal stand-in so the FlagGems fused kernels can be imported standalone."""

scripts/sparse_mla_port/vendor/flag_gems/utils/triton_version_utils.py

@@ -0,0 +1,27 @@
+"""Stand-in for FlagGems' version probe.
+
+The real module gates TLE (Triton Language Extensions) on the triton version. Here it reports
+whatever the installed triton actually supports, because the point is running the kernel, not
+reproducing FlagGems' packaging.
+"""
+
+import triton
+
+
+def _version_tuple() -> tuple[int, ...]:
+    parts = []
+    for piece in triton.__version__.split(".")[:3]:
+        digits = "".join(c for c in piece if c.isdigit())
+        parts.append(int(digits) if digits else 0)
+    return tuple(parts)
+
+
+def has_triton_tle(*minimum: int) -> bool:
+    try:
+        import triton.experimental.tle.language  # noqa: F401
+    except Exception:
+        return False
+    return _version_tuple() >= tuple(minimum)
+
+
+HAS_TLE = has_triton_tle(3, 6, 0)