[yunyang1999]

Motivation

DSA prefill on CUDA has no fused Triton path, and the kernels it does dispatch to are a poor fit for the head count the model actually has after TP.

flash_mla_sparse_fwd requires num_heads % 64 on SM90 and % 128 on SM100+, so _forward_flashmla_sparse zero-pads the head dim before calling it. GLM-5.x at TP8 has 8 real heads, so 7/8 of the tensor-core work on SM90 — 15/16 on SM100+ — is spent on padding, and the split form then writes O(tokens x splits) partials to HBM for a second kernel to merge. What main selects on its own for an FP8 KV cache on SM90 today is flashmla_kv, which is slower still from 4096 tokens up; both are columns in the tables below.

Neither cost is necessary at this size. The attention tile is small enough that one Triton program per query token, holding the online softmax in registers, beats a wide cooperative block, and it lets the head dim be padded into a 16-row MMA tile (8 -> 16) instead of into the grid (8 -> 128) — no partials, no merge.

Modifications

Adds --dsa-prefill-backend triton_sparse_mla, default off, plus one opt-in fast path. The fast path is an algebraic rewrite rather than an approximation, and it is gated against an fp32 reference in the tests.

--dsa-triton-union {2,4}

In DSA prefill the indexer's scores drift slowly in t, so neighbouring query tokens select nearly the same KV rows. Measured on captured GLM-5.1 traces:

capture	valid rows/token	union of 2 tokens	union of 4 tokens
bs1 x 2048	1024	1.00x one token's set	1.01x
bs1 x 4096	1536	1.05x	1.09x
bs1 x 8192	1792	1.12x	1.23x
bs2 x 4096	1536	1.05x	1.09x
bs4 x 2048	1024	1.00x	1.01x

One gather can therefore serve G tokens for ~1.1x the rows. G adjacent tokens share a single gathered index set, and a G * num_heads-row MMA tile then carries real work instead of padding. A per-row ownership bitmask restores each token's own softmax support, so the result is identical to the per-token path.

The union set is built per group, on the device, with nothing that depends on the KV span: each token's topk rows are sorted on their own (a K-wide segment, which torch sorts in shared memory), then one Triton program per group runs a vectorised lower_bound of every row against the other tokens' sorted lists. A token emits a row if no earlier token of the group selected it, and ORs in the bit of every token that did. Cost is O(G * K * log K) gathers per group from lists that stay in L1, so there is no bytemap, no workspace that outlives the call, and no device-to-host read. The union tile must be a power of two in [16, 32] rows (tl.arange needs a power of two, tl.dot needs M >= 16), which the validator enforces at startup: with 8 heads both group sizes are legal, at 16 heads only G=2; on SM90 real indices G=2 is also the better setting (below).

Recommended setting: --dsa-triton-union 2. On SM90 it is the faster of the two at every real-index shape and end to end; on SM120 the two are within 3% of each other at kernel level, and 4 is the one measured end to end there.

One measurement note, because it is easy to get wrong: uniformly random indices have none of that structure and understate this path by about 2x. The benchmark added here builds indices as a calibrated random walk instead — token t+1 is token t with ~2% of rows swapped, plus the -1 padding real prefill has for t < topk — so the synthetic union sizes match the captured ones.

Robustness

• Token offsets inside the kernel are int64, so a large extend batch at 64 or 128 heads cannot wrap into another token's rows.
• The launcher steps the shared-memory tile down rather than propagating OutOfResources to the request, and remembers the tile that fit per (num_heads, G, device), so the retry runs once per shape rather than on every layer.
• The startup validator rejects what cannot work or exceeds what was measured: ROCm, SM < 90, more than 32 query heads per rank, index_kpool > 1 (the pooled indexer tail is only rerouted for flashmla_sparse), and union tiles outside the power-of-two [16, 32] window. Tiles are tuned on SM90 and SM120; other SM90+ architectures run an untuned default tile and log a warning at startup (the SM100 CI lane exercises this path).
• Union is forced off under --enable-deterministic-inference (a token's output would depend on which tokens share its group) and whenever the stream itself is capturing (the FULL prefill graph); the per-token path gives the same result there. Inside the breakable prefill graph, where attention runs eagerly between the graph segments, union stays on.
• The wrapper refuses a non-bf16 KV cache with a clear error instead of reading packed FP8 bytes at a bf16 stride, and takes topk_length from the metadata the backend already holds rather than recomputing it per layer.
• Known follow-up, not in this PR: get_topk_transform_method keeps MIXED batches on the PAGED transform, so --enable-mixed-chunk with an FP8 KV cache reaches this backend with the packed pool and hits that error on the first mixed batch. flashmla_sparse has the same routing gap on main, so both will be fixed together in a separate change.

Files

New kernel module (no environment variables); the dispatch is folded into the existing flashmla_sparse branch of dsa_backend.py (the input preparation is identical; only the final kernel call differs), with a startup validator and the option read at construction; one server arg (--dsa-triton-union) in arg_groups/fields/exec_.py; the SM120 GLM-FP8 allowlist in flash_mla_sm120.py admits triton_sparse_mla for prefill, backed by the end-to-end launch below; and registration in the DSA variant matrix so test/registered/attention/unittests/dsa/test_dsa.py exercises this backend too; and the triton_sparse_mla row plus the --dsa-triton-union entry in docs/docs/advanced_features/attention_backend.mdx and server_arguments.mdx.

Accuracy Tests

Kernel — test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py, against an fp32 reference: base path at three shapes, -1/ragged padding, union exactness on overlapping-but-unequal selections at G=2 and G=4 including a T % G != 0 tail (T=1023, G=4), the union dedup contract on its own (each row once, ownership bits equal to the OR of its owners, -1 dropped — a duplicated row is weighted twice, which can hide under the cosine gate), the fall-back to the per-token path when the union tile is illegal (h=4, G=2), int32-vs-int64 gather addressing equivalence, determinism, two forced shared-memory step-downs (an oversized tile, so the retry path runs on the CI GPU too; the recorded fit must differ from the requested tile), and the backend method itself driving the real kernel with the [tokens, heads, d_v] contract. CPU — test/registered/unit/attention/test_triton_sparse_mla_prefill.py, 20 cases: the validator's boundaries (including the new power-of-two, head-cap and index_kpool rejections), the deterministic, stream-capture and breakable-graph gates, topk_length plumbing and its row-count guard, the RAGGED top-k routing regression, CLI registration, and that the SM120 GLM FP8 auto-selection still resolves to flashinfer_sparse_mla with this backend registered.

Every timed point below is correctness-gated against the same fp32 reference before it is timed — 5/5 real-index shapes PASS at cosine 1.0000 on each architecture.

Two independent model-level checks, at different tolerances, because they answer different questions.

(a) The kernel's own output tensor — GLM-5.2 config, TP8 on SM120, --load-format dummy (fixed per-parameter seed, so both runs get identical weights) and a seeded prompt at 4096 tokens. The reference is flashinfer_sparse_mla, the backend SM120 selects on its own (flashmla_sparse is not admitted on that platform). Comparing the layer-0 prefill attention output — the tensor this backend actually computes — in fp64, against a control of the reference rerun against itself:

comparison	cosine	rel L2	max abs
control: flashinfer_sparse_mla vs itself (rerun)	0.9999998258	5.90e-04	3.05e-05
triton_sparse_mla	0.9996746051	2.55e-02	4.39e-03
+ --dsa-triton-union 4	0.9996746126	2.55e-02	4.39e-03
--dsa-triton-union 4 vs the per-token path	0.9999994603	1.04e-03	6.10e-05

Two readings. The union path sits on top of the per-token path at the control's own noise level (rel L2 1.0e-03 against 5.9e-04, max abs 6.1e-05 against 3.1e-05): it is the same computation with the rows accumulated in a different order. The 2.5% between this kernel and flashinfer_sparse_mla is the distance between an FP8-native kernel and a bf16 one, not an error in either: against the bf16 flash_mla_sparse_fwd on SM90 (same config, same seed) this kernel is within 1.5x of rerun noise, and against an fp32 reference on the five real captures it scores cosine 1.0000 on every path.

SM90, GLM-5.2 config, reference flashmla_sparse	cosine	rel L2	max abs
control: flashmla_sparse vs itself (rerun)	0.9999999134	4.16e-04	3.05e-05
triton_sparse_mla	0.9999998054	6.24e-04	6.10e-05
+ --dsa-triton-union 4	0.9999996655	8.18e-04	6.10e-05
+ --dsa-triton-union 2	0.9999996657	8.18e-04	6.10e-05
--dsa-triton-union 4 vs the per-token path	0.9999997687	6.80e-04	3.05e-05

Against the bf16 reference the per-token path sits within 1.5x of the rerun noise and the union paths within 2x. The pre-review revision measured the per-token row at 0.9999997999 / 6.33e-04 on the same seed, which is what an unchanged kernel should look like; the union rows are the new per-group dedup.

This is the tight measurement: it reads the kernel's output directly, before 78 layers of MoE can amplify anything. GLM-5.1 shares every attention dimension and reproduces these figures.

(b) What the model actually emits, on real weights — GLM-5.2 (78 layers, FP8, TP8 on SM90), seeded prompt, next-token logits over the full 154880-wide vocabulary in fp64, against a rerun control and against main's own default for this platform:

tokens	comparison	cosine	rel L2	argmax
2048	every pair	1.0000000	0	same
4096	control: flashmla_sparse vs itself	0.9985622	7.49e-02	same
4096	auto (flashmla_kv, main's pick)	0.9987185	6.46e-02	same
4096	triton_sparse_mla	0.9987416	5.02e-02	same
4096	+ --dsa-triton-union 4	0.9984925	5.87e-02	same
4096	+ --dsa-triton-union 2	0.9990947	4.25e-02	same
8192	control: flashmla_sparse vs itself	0.9995941	2.85e-02	same
8192	auto (flashmla_kv, main's pick)	0.9993871	3.53e-02	differs
8192	triton_sparse_mla	0.9994118	3.53e-02	same
8192	+ --dsa-triton-union 4	0.9990762	5.36e-02	same
8192	+ --dsa-triton-union 2	0.9991140	4.72e-02	differs

At 2048 tokens every backend emits bit-identical logits: SM90 routes that length to the dense one-shot path and the stack is deterministic there. From 4096 up the tolerances are three orders looser than in (a) — that is 78 layers of 256-expert FP8 MoE, not the kernel, and the control shows it: the reference rerun against itself moves by 3-7% in rel L2. The per-token path sits inside the control's own band at both lengths. The union paths are inside it at 4096 and within 2x of it at 8192, the same order as main's own alternative kernel, which also flips the argmax at 8192. A greedy-token match is therefore not a usable metric on this stack; the layer-0 tensor in (a) is the measurement that isolates the kernel.

(c) SM120, real weights, first-token agreement — GLM-5.2-NVFP4, TP8, greedy, one generated token, comparing the argmax token and its logprob across the three prefill backends on otherwise identical servers. Each row of the control is a separate server instance, so it measures run-to-run noise of the stack itself:

prompts	comparison	argmax agreement	mean |Δ logprob| (agreeing)	max
50 natural-language	control: flashinfer_sparse_mla vs its own rerun	50/50	0.0000	0.0000
50 natural-language	triton_sparse_mla vs flashinfer_sparse_mla	47/50	0.090	0.255
50 natural-language	+ --dsa-triton-union 4 vs triton_sparse_mla	50/50	0.0000	0.0000
50 random-word strings	control: flashinfer_sparse_mla vs its own rerun	50/50	0.0000	0.0000
50 random-word strings	triton_sparse_mla vs flashinfer_sparse_mla	28/50	0.358	0.854
50 random-word strings	+ --dsa-triton-union 4 vs triton_sparse_mla	50/50	0.0000	0.0000

Unlike the SM90 FP8 stack in (b), this deployment is bit-reproducible across server instances (the controls are exact), so the differences against flashinfer_sparse_mla are real and come from the arithmetic: that kernel consumes the FP8 pool with its scales directly, while this one takes the dequantized bf16 rows the RAGGED transform produces and runs a bf16 MMA. The three natural-prompt disagreements are all near-ties — "A year on Earth has approximately 365" (top-1 logprob −0.48 vs −0.71), "A hexagon has" (−0.94 vs −0.86), "The opposite of hot is" (−0.76 vs −0.44) — and on agreeing tokens the logprob moves by 0.09 on average. Random-word prompts have flat next-token distributions, so the same-sized perturbation flips the argmax far more often; they are shown for completeness, not as an accuracy metric. The union path is bit-identical to the per-token path on every prompt of both sets, and the greedy 48-token continuation of a fixed prompt is identical between them too.

Speed Tests and Profiling

Kernel-level numbers are on real captured GLM-5.1 index sets at five (batch, seqlen) points, not synthetic indices; times in ms, median of 15, every variant gated at cosine 1.0000 against the fp32 reference. End-to-end numbers are on real model weights except the two 8-layer dummy-weight configs, which are labelled as such; every end-to-end table below is the current revision.

SM90 — kernel level, vs flash_mla_sparse_fwd (padded 8 -> 64 heads, as the backend does)

capture	tokens	flashmla	ours	+union2	+union4
bs1 x 2048	2048	4.516	1.209	0.850	0.940
bs1 x 4096	4096	8.857	3.154	2.237	2.421
bs1 x 8192	8192	17.632	7.142	5.173	5.710
bs2 x 4096	8192	17.629	6.217	4.348	4.642
bs4 x 2048	8192	17.630	4.373	3.090	3.294
geomean			3.12x	4.41x	4.06x

The spread is wide but monotone in the right variable: shorter requests favour this kernel most, because more of their rows are -1 padding that the per-token early exit skips and the padded dense kernel does not. On real indices G=2 is the better union setting: the captured selections overlap so tightly that a group of 4 gathers barely fewer rows than a group of 2, while its membership test costs G*(G-1) searches instead of 2. On the calibrated synthetic benchmark (SM90, 8 heads) the per-token path takes 1.16 / 3.15 / 7.14 ms at 2048 / 4096 / 8192 tokens and union=4 0.96 / 2.26 / 4.88 ms, i.e. up to 1.46x; at 16 heads only G=2 is legal and it does not pay off there (7.36 -> 8.92 ms), so union is a TP8 setting.

SM120 — kernel level, this kernel's paths on the same five captures

capture	tokens	ours	+union2	+union4	best vs ours
bs1 x 2048	2048	1.087	0.656	0.645	1.69x
bs1 x 4096	4096	3.288	1.845	1.735	1.90x
bs1 x 8192	8192	7.675	4.361	4.379	1.76x
bs2 x 4096	8192	6.653	3.708	3.592	1.85x
bs4 x 2048	8192	4.583	2.705	2.586	1.77x

Union pays more on SM120 than on SM90 (1.7-1.9x over the per-token path against 1.3-1.4x) because the gather is a larger share of the kernel on this part. This table deliberately carries no flashinfer_sparse_mla column: that kernel takes the packed FP8 pool with its scales, which the real-index harness does not synthesise, so the like-for-like comparison against it is the end-to-end table below rather than a kernel-level number.

SM120 — end to end, GLM-5.2-NVFP4, real weights, TP8, prefill latency

Launched from this branch with sglang.launch_server --attention-backend dsa --kv-cache-dtype fp8_e4m3 --quantization modelopt_fp4 --tp 8, batch 1, one generated token, median of 5 requests with fresh random prompts so the radix cache never serves the prefix. The baseline is what mainline selects on SM120 by itself (prefill=flashinfer_sparse_mla, decode=flashinfer_sparse_mla).

tokens	baseline	triton_sparse_mla	+ --dsa-triton-union 4	ours	+union4
2048	556.9 ms	550.6 ms	550.4 ms	1.01x	1.01x
4096	1137.9 ms	1131.4 ms	1008.5 ms	1.01x	1.13x
8192	2315.6 ms	2295.5 ms	2034.8 ms	1.01x	1.14x
16384	4848.7 ms	4806.9 ms	4266.5 ms	1.01x	1.14x

Two things this table says plainly. The per-token path alone barely moves the end-to-end number on this deployment: at TP8 with NVFP4 experts the attention kernel is a small share of the prefill, so the per-token path's kernel-level gain shows up as ~1% of the request. Union is where the end-to-end gain is (1.13-1.14x from 4096 tokens up), because it removes most of the gather rather than shaving the MMA. Every run counts invocations of this kernel (union=0 during CUDA-graph capture, union=4 on the served requests), so a silent no-op cannot be mistaken for a win, and the greedy continuation of a fixed prompt is identical between the per-token and union paths.

SM90 — end to end, 8-layer configs, TP8, prefill latency at 4096 tokens

Small configs with --load-format dummy, kept because they cover both model generations side by side and are cheap enough to rerun on every revision; these rows are the current code. Two reference runs are shown to make the run-to-run spread visible and the ratios are against their mean. auto is what main selects on its own for an FP8 KV cache on SM90 today (prefill=flashmla_kv), i.e. the number a user who changes nothing gets.

model	auto (flashmla_kv)	flashmla_sparse	triton_sparse_mla	+ union 4	+ union 2
GLM-5.1	131.4 ms	112.3 / 110.6 ms	80.7 ms (1.38x; 1.63x vs auto)	75.3 ms (1.48x)	75.3 ms (1.48x)
GLM-5.2	126.3 ms	107.4 / 107.5 ms	77.0 ms (1.40x; 1.64x vs auto)	71.5 ms (1.50x)	70.3 ms (1.53x)

For the record, the same launch in July measured flashmla_sparse at 128 / 125 ms and this kernel at 82 / 78 ms (1.57x / 1.61x): the reference path has become ~14% faster on main since then while this kernel's time is unchanged, so the ratio on this small config is 1.4x today, not 1.6x.

The same kernel serves both generations unchanged — GLM-5.2 has the same DSA shape (kv_lora_rank 512 + qk_rope_head_dim 64 = 576, d_v 512, topk 2048, 8 heads at TP8) and the same GlmMoeDsaForCausalLM architecture. Its layers are not uniform (indexer_types mixes full and shared), so the run counts kernel invocations to confirm coverage rather than infer it: 16 per rank for an 8-layer model over two prefills, i.e. every DSA layer.

SM90 — end to end, GLM-5.2 (78 layers, real FP8 weights), TP8, prefill latency

Each length is visited twice in-process and the second visit is reported: bench_one_batch warms up only at input_len[0], so a plain sweep would charge the first visit to each new shape with Triton JIT and allocator growth, which lands on this backend harder than on a precompiled CUDA kernel. Two independent reference runs agree to within 0.8% at every length. auto is main's own pick for an FP8 KV cache on SM90 today (prefill=flashmla_kv). Ratios are against the flashmla_sparse mean, then against auto.

tokens	auto (flashmla_kv)	flashmla_sparse (x2)	triton_sparse_mla	+ union 4	+ union 2
2048	0.211 s	0.213 / 0.212 s	0.212 s	0.213 s (rerun)	0.214 s
4096	1.072 s	0.888 / 0.888 s	0.582 s (1.53x; 1.84x)	0.529 s (1.68x; 2.03x)	0.516 s (1.72x; 2.08x)
8192	2.111 s	1.891 / 1.892 s	1.197 s (1.58x; 1.76x)	1.087 s (1.74x; 1.94x)	1.048 s (1.80x; 2.01x)

The kernel is not invoked at 2048 tokens on SM90 (next paragraph), so that row is one code path for every column; the union-4 run's first pass measured 0.285 s there and its rerun 0.213 s, with 0.529 / 1.086 s at the other two lengths (within 0.1% of the first pass), so the 0.285 was a one-off and the rerun is shown. Against the bytemap union this PR replaced (0.521 / 1.065 s at 4096 / 8192 in July), the per-group dedup costs 1.5-2% end to end, the price of removing the three defects. G=2 is the better setting on real weights as it was on the captured indices: the real selections overlap so tightly that a group of 4 gathers barely fewer rows than a group of 2, while its membership test does six searches per row instead of two.

Short contexts are unchanged by design, not by accident. At 2048 tokens the configurations are identical because the kernel is never called: on SM90 and SM100, dsa_backend.py routes prefill to a dense MHA one-shot path when max_kv_len <= SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD (default 2048). Counting invocations confirms it — 312 calls per rank across the sweep, i.e. 78 layers x the 4 prefills that exceed the threshold, and zero for the rest.

A benchmark is included at test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py. Absolute numbers here are from SM90 and SM120 parts that are not CI runners, so the benchmark is written to make the relative comparison reproduce on 1-gpu-large.

Checklist

• Format your code according to Format code with pre-commit
• Add unit tests according to Run and add unit tests
• Update documentation — attention_backend.mdx (DSA sub-backend table) and server_arguments.mdx (--dsa-prefill-backend choices, --dsa-triton-union)
• Provide accuracy and speed benchmark results
• Follow the SGLang code style guidance

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): 🚫 Run #37255031786
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ❌ Run #37255031958

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[yunyang1999]

Hi maintainers — this is a fork PR so I can't add the run-ci label myself. Could someone add it (or run /tag-and-rerun-ci) to kick off CI? Thanks.

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

@Fridge003 Could you help us to review this GLM-5.1/5.2 SM90/SM120 perf optimization? Thanks!

[yunyang1999]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[yunyang1999]

@ch-wan thanks for the careful read -- all eleven were real and the two High ones were structural. Everything below is in the branch as of the latest push; one summary comment rather than eleven replies, in your order:

1. Union tile not validated. _validate_triton_sparse_mla_backend now rejects any num_q_heads * union that is not a power of two in [16, 32], and _union_path returns False for the same shapes so a direct caller falls back to the per-token path instead of hitting the Triton compile error. CPU and GPU regression tests (h=4/G=2, h=12/G=2, h=6/G=4) both fail on the old code.
2. _UNION_WS unbounded. Gone with the bytemap. The union set is built per group: each token's K rows are sorted on their own (torch.sort on a K-wide segment, in shared memory) and one Triton program per group does a vectorised lower_bound membership test, O(GKlog K) L1-resident gathers whatever the KV span. Per-call tensors only, nothing persists, nothing for mem_fraction_static to miss. test_union_dedup_contract checks the output against a Python set reference (each row once, bits = OR of owners, -1 dropped); it caught an off-by-one in the search depth while I built it.
3. Large head counts. Token offsets in the base kernel are int64, and the validator caps num_q_heads at 32, the largest count measured. Tiling over heads is the right follow-up if a deployment needs 64/128 here; I would rather land it measured than promise it.
4. index_kpool > 1. Rejected at construction, pointing at flashmla_sparse.
5. Host sync per layer. Gone with the bytemap: no .tolist(), no other host read.
6. Capture and deterministic guards. Union is forced off under --enable-deterministic-inference at construction (with a warning, for the batch-invariance reason you gave) and whenever the stream itself is capturing. Following the second review it is gated on torch.cuda.is_current_stream_capturing() alone, so breakable-graph replays, where attention runs eagerly between segments, keep union on.
7. topk_length. Passed from metadata.dsa_cache_seqlens_int32 with the same row-count guard _forward_flashmla_sparse uses.
8. Step-down on every call. _tile_candidates remembers the tile that fit per (path, num_heads, G, device), so the OutOfResources retry runs once per shape. Kernel-level time of the base path is unchanged (7.143 vs 7.144 ms at T=8192).
9. Docstring promises. The >= S claim is gone (-1 marks an invalid slot, no upper bound is checked) and the wrapper raises TypeError on a non-bf16 KV cache.
10. do_not_specialize. epoch and base no longer exist; the remaining integer arguments are per-config constants.
11. Stale references and em-dashes. The GLM_NSA_UNION / decision-ledger / cuda-agent notes are removed, the module docstring is cut to the interface contract, and the new code is ASCII-only.

Description rewritten against the diff: one server arg, no dense-prefix material, the union section describing the per-group dedup, and every table re-measured on the current kernel with a rerun control and main's own auto pick as a column (on SM90 with an FP8 KV cache that is now flashmla_kv, not flashmla_sparse). Recommended setting is --dsa-triton-union 2.

SM120, launched end to end from this branch (GLM-5.2-NVFP4, TP8, lmsysorg/sglang:v0.5.20-cu130): mainline's DSA path does run there today (it auto-selects flashinfer_sparse_mla), so the earlier statement was stale and is gone. Our per-token path is ~1% faster end to end (attention is a small share of that deployment's prefill); with union it is 1.13-1.14x from 4096 tokens up (1008 vs 1138 ms, 2035 vs 2316 ms, 4267 vs 4849 ms). Greedy first token over 50 natural prompts: flashinfer vs its own rerun 50/50, triton_sparse_mla vs flashinfer 47/50 with the three disagreements near-ties (FP8-native vs bf16 arithmetic), union vs per-token 50/50 with identical logprobs. Given that, I kept the flash_mla_sm120.py allowlist entry; if you would rather see it land separately I will split it out.

SM90 was re-launched on the same revision as well (SM90, 8 GPUs, TP8): 8-layer dummy configs for the layer-0 tensor check (per-token path within 1.5x of rerun noise, union within 2x) and GLM-5.2-FP8 real weights for the latency and logits tables: 1.53x / 1.58x over flashmla_sparse at 4096 / 8192 tokens (1.84x / 1.76x over flashmla_kv), 1.72x / 1.80x with --dsa-triton-union 2.

[yunyang1999]

@nvpohanh thanks, all six were right; they are in the latest push, in your order:

1. Stale union numbers. The description now carries only numbers measured on the current kernel (SM120 runs 09-29, SM90 runs 09-30; the kernel has since changed only in its docstring and by dropping a constexpr-dead FP8 branch), and the recommendation is --dsa-triton-union 2. On the 78-layer GLM-5.2-FP8 SM90 launch the union gain at 8192 tokens is now 149 ms for G=2 and 110 ms for G=4 against the bytemap's 186 ms, so your estimate pointed the right way; the net cost is below 1.3 ms x 78 because the bytemap's own mark / compact kernels and host sync went with it (new union 4: 1.087 s vs 1.065 s, +2%).
2. BCG replay. The gate is torch.cuda.is_current_stream_capturing() alone, set under FULL prefill capture and clear inside the breakable graph's eager attention segments; the comment says why. The old test is replaced by test_union_stays_on_when_only_the_breakable_graph_flag_is_set, which fails on the previous code (0 != 4).
3. Docs. triton_sparse_mla row in the DSA table of attention_backend.mdx (prefill only, CUDA SM90+, at most 32 query heads per rank, not with index_kpool, union rule), the backend in the --dsa-prefill-backend options and a --dsa-triton-union row in server_arguments.mdx. Thanks for pointing at the pages.
4. Step-down tests. Both pass an oversized tile ((256, 8, 4) / (256, 4, 4)), clear _FIT_TILE first and assert the recorded fit is not the requested tile; on SM90 they settle on (128, 4) and (128, 2), so the retry loop runs on every GPU instead of being assumed.
5. Vacuous defaults test. Replaced by a run of the SM120 GLM FP8 arm of _dsa_split_backend_resolution with capability / platform / model config / is_deepseek_dsa mocked: flashinfer_sparse_mla for both phases when nothing is set, and a user-set triton_sparse_mla prefill kept as given.
6. Docstrings. The module docstring is the interface contract only; test and benchmark docstrings are cut the same way; the one remaining union-size note (the benchmark's _make_inputs) quotes the same 1.00-1.12x / 1.01-1.23x range as the PR table.

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[ch-wan]

Thanks for the update. Everything from the last round is addressed: the per-group dedup removes the workspace cache, the span budget and the per-layer host sync; the validator now enforces the union tile shape, the head cap and index_kpool; union is off under deterministic inference and under stream capture; and the description now matches the diff. Thanks also for running triton_sparse_mla end to end from this branch on SM120.

This looks ready to me apart from two small items. Details are inline:

• Low: no test covers the union tail (T % G != 0). A T=1023, G=4 case would.
• Nit: the arch gate admits SM100/SM103 with untuned tiles. The description says the validator "rejects what has not been measured", which doesn't cover that. Please reword it or restrict the gate.

Follow-up, not blocking: get_topk_transform_method (python/sglang/srt/layers/attention/dsa_backend.py:3809) keeps MIXED batches on the PAGED transform. So with an FP8 KV cache, --enable-mixed-chunk plus triton_sparse_mla hits the new TypeError on the first mixed batch. That includes SM120, where the allowlist now admits this backend. flashmla_sparse has the same gap on main, so it is better fixed separately.

[nvpohanh]

@yunyang1999 Could you address the final comments? After that, I will monitor the CI and merge this. Thanks!

[nvpohanh]

/tag-and-rerun-ci

[yunyang1999]

@ch-wan thanks for the approval, both items are in the latest push; @nvpohanh over to you for the merge, THX!

• Union tail test. Added your T=1023 / G=4 case to test_union_is_exact_on_overlapping_selections as suggested. On the real code it passes at cosine 0.999998, and the three tail rows are bit-identical to the per-token path's. To confirm it is a real guard I also ran it with the tail hand-off in _union_path deliberately disabled: it then fails at cosine 0.9987 against the 0.999 gate.
• Arch gate vs description. Went with the rewording: the gate stays at SM90+, and the description, the attention_backend.mdx row and the test comment now say that tiles are tuned on SM90 and SM120 while other SM90+ architectures run the untuned default tile and log a warning at startup. I kept SM100 admitted rather than restricting the gate because the 4-GPU SM100 CI lane already runs the DSA variant suite through this backend and passes, so the warning path is exercised rather than assumed; an SM100 sweep that adds a _PINNED entry is the natural follow-up.
• MIXED batches with an FP8 KV cache. Agreed that it belongs in a separate change together with flashmla_sparse; the description now lists it as a known follow-up so nobody trips over it in the meantime.

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

/rerun-failed-ci

[kpham-sgl]

We can merge, the failing test is flaky on main (fixing it)

[yunyang1999]

@ch-wan @nvpohanh the mixed-chunk follow-up from the review is up as #44176: with an FP8 KV cache, --enable-mixed-chunk killed the server on the first mixed batch for flashmla_sparse, flashmla_sparse_q8 and triton_sparse_mla (the attention metadata was built from the MIXED batch and routed PAGED while the EXTEND-replaying prefill graph asked for RAGGED). MIXED now routes exactly like EXTEND, with a kit case, a routing-table pin and an end-to-end check on real weights.