f25e5d8b6c
The last thing that kept the SM90 path on TileLang was the sparse MLA attention kernel. FlagGems
has a pure-Triton implementation of the same operator, and it turns out to be drop-in once two
differences are handled. With TileLang unimportable, the engine now runs end to end.
`h20_ops/sparse_mla.py` picks the backend: `DEEPINFER_SPARSE_MLA_BACKEND=tilelang|triton`, unset
means TileLang when it imports and Triton otherwise. `attn_sm90` imports from the dispatcher, so its
two entry points were the only call sites touched. Auto resolution catches `ModuleNotFoundError`
rather than testing for the package by name, because a missing TileLang fails inside
`sparse_mla_sm90`, not at the module lookup -- and it only swallows a *missing* dependency, so a
present-but-broken module still surfaces instead of being swapped out.
Vendored rather than depended upon: `flag_gems/__init__.py` registers ATen implementations
process-wide (`torch.library.Library("aten", "IMPL")` then `replace_customized_ops`). Importing the
package to reach one kernel would swap operators under the whole engine. The PyPI 5.0.2 file has no
`flag_gems` internal imports at all, so there was nothing to gain; it is 329 lines, accumulates in
fp32, and already guards the zero-width rope tail that V4.1 needs (RoPE folded inside the 512-wide
latent). The file header carries the source md5 and both deviations.
Two things it needed, both found by comparison rather than by reading:
* 5.0.2 guards each top-k block by its **position** in the index array (`if ck * BK <= max_col`).
Non-causal that is `SKV-1`, so when `topk > SKV` every block past the first `SKV` entries is
skipped. At SQ=16/SKV=128/topk=512 the output is off by 2.8e-2 -- and any prompt shorter than the
512-entry top-k hits it. Restored the data-dependent guard FlagGems master uses, marked in place.
* Its lse is a natural log; the TileLang one is log2. Unconverted, `_apply_sink` would be off by
1.4427x: plausible output, wrong numbers. The adapter converts, and the test asserts the ratio is
1.0 rather than 0.693 so it cannot regress quietly.
The kernel's operands have no batch axis, while decode hands it a different KV window per query.
Rather than loop, the adapter concatenates the windows and shifts query b's indices by b*seq_kv --
one launch, no extra memory, since the concatenated buffer is the tensor decode already built.
Verification ran in order, each step gating the next: 7 shapes against TileLang (`ALL_MATCH`, <=2
bf16 ULP, including `topk_length` and `attn_sink`), the decode call shape incl. the SWA+full
two-cache case (`ALL_MATCH`), the adapter contract (5 tests), the dispatcher and both `attn_sm90`
entry points (4 tests), then end to end with TileLang blocked: exit=0, 8/8 DP ranks agreeing, zero
kernel fallbacks, correct answers. A control run with TileLang present is unchanged.
Not token-identical, and the doc says so: prompt A shares 15 tokens then diverges, prompt B shares
1. A control run of the TileLang backend twice is bit-identical, so this is accumulation order, not
noise -- consistent with the 2.7% logits drift already recorded for the SM90 quantisation floor.
One TileLang dependency remains, and it is the deliberate one: `fused_sample`'s rejection-sampling
variant, which needs TileLang's Philox stream to reproduce the same tokens. So the pure-Triton run
disables speculative decoding (`DEEPINFER_USE_NEXTN_DECODE=0`).
6.6 KiB
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, which has no flag_gems internal imports at all -- only torch and triton. That
matters: importing flag_gems itself registers ATen implementations process-wide
(torch.library.Library("aten", "IMPL") then runtime.replace_customized_ops), so depending on the
package to reach one kernel would swap operators under the whole engine. The file header records the
source md5 and the two local deviations.
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.pyandtest_wiring.pyat 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_colisSKV-1, so withtopk > SKVevery block past the firstSKVentries 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_sinkby 1.4427x -- plausible output, wrong numbers.test_sparse_mla_triton.pyasserts 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.