c6a65fe6ad
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.
|
||
|---|---|---|
| .. | ||
| vendor | ||
| ab_vs_tilelang.py | ||
| probe.py | ||
| README.md | ||
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.pystill imports TileLang; a Triton backend behind the samepatch_importshook 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'slseis in log2 units where FlagGems' is natural log. A wrapper must dropmax_logitsand convert, or every downstream consumer reads a 1.44x-wrong scale.