DeepInfer/scripts/sparse_mla_port

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.