# 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 could not start. FlagGems ships a pure-Triton implementation of the same
operator; it is now vendored, verified against the TileLang kernel, and wired in behind a dispatcher
(`deepinfer/h20_ops/sparse_mla.py`). With TileLang unimportable the engine runs end to end on it.

## Source

`flagos-ai/FlagGems`, file `src/flag_gems/fused/flashmla_sparse.py`, function
`flash_mla_sparse_fwd`. Vendored as `deepinfer/h20_ops/vendor/flashmla_sparse.py` from the **PyPI
5.0.2 wheel**. Two reasons, plus a correction to what this file said previously:

* Import cost: `flag_gems/__init__.py` star-imports every op module and pulls in Triton TLE, which
  loads `libflagcx` through ctypes -- an .so this deployment does not otherwise need. One attention
  kernel should not drag that in.
* The released file has no `flag_gems` internal imports at all (only torch and triton), so the
  package bought nothing for this particular operator.

Correction: importing `flag_gems` does **not** register ATen implementations, contrary to the earlier
wording here. It creates the `aten_lib` handle; registration happens only in `enable()` or
`only_enable(include=[...])`. Measured by hashing `torch._C._dispatch_dump_table` before and after
`import flag_gems`: unchanged for `aten::rms_norm`, `aten::silu`, `aten::softmax`, `aten::topk`. The
blast radius lives in `enable()`, not in the import. That distinction is what makes a broader
FlagGems integration tractable -- see `H20-FLAGGEMS-PLAN.md`.

Master's copy (~1200 lines, TLE fast path, fp16 accumulate) also works and was verified first; 5.0.2
is used instead because it is a release, is 329 lines, accumulates in fp32, and already guards the
zero-width rope tail that V4.1 needs.

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_vendored.py`, 7 shapes, TileLang as reference: **all match** within 2 bf16 ULP, including
`topk_length` and `attn_sink`. This runs the in-tree vendored file, so it is what the engine
actually executes.

| 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

* Decode path end to end is covered by `ab_decode.py` and `test_wiring.py` at the entry-point
  level, but the engine's own decode was only exercised through the end-to-end runs below, not
  against a long-context trace.
* Not token-identical to the TileLang backend. See below -- it runs and answers correctly, but the
  two backends do not produce the same tokens.
* `fused_sample`'s rejection-sampling variant still needs TileLang, so the pure-Triton end-to-end run
  has to disable speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`). With MTP on, that one kernel
  is the remaining TileLang dependency.

---

## Wiring status

`h20_ops/sparse_mla.py` selects the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
means TileLang if it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
two entry points were the only call sites to touch. The chosen backend is logged once per rank.

### Verification, in the order it was done

Each step had to pass before the next was attempted, so a failure would point at one thing.

| step | what | result |
|---|---|---|
| 1 | `ab_vendored.py` — vendored kernel vs TileLang, 7 shapes | `ALL_MATCH`, ≤2 bf16 ULP |
| 2 | `ab_decode.py` — decode call shape, fp8 cache, 5 cases incl. two-cache | `ALL_MATCH` |
| 3 | `test_sparse_mla_triton.py` — adapter contract | `ALL_PASSED` (5 tests) |
| 4 | `test_wiring.py` — selection + both entry points | `ALL_PASSED` (4 tests) |
| 5 | end-to-end, TileLang blocked | `exit=0`, 0 fallbacks, correct answers |
| 6 | end-to-end control, TileLang present | `exit=0`, 0 fallbacks, correct answers |

Two things the vendoring had to fix, both caught by step 1 rather than by luck:

* FlagGems 5.0.2 guards each top-k block by its **position** in the index array
  (`if ck * BK <= max_col`). Non-causal, `max_col` is `SKV-1`, so with `topk > SKV` every block past
  the first `SKV` entries is skipped. At SQ=16/SKV=128/topk=512 that is a 2.8e-2 error. Any prompt
  shorter than the 512-entry top-k hits it. Fixed locally, marked in the file; FlagGems master uses
  the data-dependent form.
* The kernel's lse is a natural log; the TileLang one is log2. Passed through unconverted it would
  scale `_apply_sink` by 1.4427x -- plausible output, wrong numbers. `test_sparse_mla_triton.py`
  asserts the ratio is 1.0 rather than 0.693 so this cannot regress silently.

### End-to-end result

Same command, only the TileLang availability differs:

```
DEEPINFER_SIMULATE_NO_TILELANG=1 DEEPINFER_KERNELS=ref DEEPINFER_OPS_STRICT=1 \
DEEPINFER_USE_NEXTN_DECODE=0 bash scripts/h20_e2e_report.sh
```

| | TileLang blocked | TileLang present |
|---|---|---|
| exit code | 0 | 0 |
| sparse MLA backend (per rank) | `triton` x8 | `tilelang` x8 |
| `ModuleNotFoundError` failures | 0 | 0 |
| kernels served by the default backend | *(none)* | *(none)* |
| DP ranks agreeing | 8/8 | 8/8 |
| answer | 中国的首都→北京，日本的首都→东京 | 同上，另有 1949-09-27 定都的展开 |

**Not token-identical.** Comparing the two: prompt A shares 15 tokens then diverges, prompt B shares
1 then diverges. A control run of the TileLang backend twice is bit-identical, so the harness is
deterministic and the difference is real -- it comes from accumulation order, ≤2 ULP per layer
compounding across 40+ layers, consistent with the 2.7% logits drift already recorded in `TODO.md`
for the SM90 quantisation floor. It is not a wiring defect, and not something this kernel can fix on
its own.