[ROCm] Admit any head count up to the MFMA tile to the gfx950 Triton sparse-MLA prefill path

[chrisaberger]

Motivation

The gfx950 Triton sparse-MLA prefill kernel is admitted only when a rank holds
exactly 16 query heads:

and layer.tp_q_head_num == 16

Any tensor-parallel width that puts fewer than 16 heads on a rank therefore
falls back to the slower TileLang partial+combine path. The fallback is silent
and correct, so nothing fails — it just costs prefill latency. GLM-5.2 has 64
attention heads, so TP8 gives 8 heads per rank and never reaches the Triton
kernel. Any heads / TP < 16 configuration is affected.

The kernel could not simply be dispatched at fewer heads: tl.dot on gfx950
requires a 16-row MFMA tile.

Note that dsa_backend.py already does exactly this padding for the aiter
mla_decode_fwd path, citing this same model and TP width:

# Aiter mla_decode_fwd supports num_heads multiples of 16 in range [16, 128].
# For models with fewer heads per GPU (e.g. GLM-5 64 heads / TP8 = 8), need to pad the heads to 16.
self.need_pad_heads = self.num_q_heads < 16

This PR applies the same idea to the Triton prefill path, and adds the launch
configuration that measured best at those shapes.

Modifications

1. Pad the query-head tile (triton_sparse_mla.py). Pad to
BLOCK_H = max(16, next_power_of_2(H)), mask the padded rows on load and on
store, and exclude them from the online softmax. m_i/l_i/acc are sized to
BLOCK_H. Real values and indices are unchanged.

2. Widen the dispatch gate (dsa_backend.py) from == 16 to <= 16, i.e.
any head count that fits the MFMA tile. It stays narrowly guarded on gfx950,
FP8 KV, v_head_dim == 512, RoPE tail 64 and topk 2048; every other shape keeps
the existing TileLang fallback.

3. Drop one autotune candidate on HIP (triton_sparse_mla.py). The torch
2.11 HIP Triton/LLVM stack aborts the compiler process on BLOCK_N=128, num_warps=1, num_stages=2 for the padded-head path. Autotune cannot recover
from a process abort, so that single candidate is omitted from the HIP grid;
every other backend keeps the full grid.

4. Test (test_triton_sparse_mla.py). New test over H=4, 8, 12 and 16 with
trailing, strided and head-plus-tail sparse index patterns, registered to
stage-b-test-1-gpu-small-amd-mi35x (the gfx950 runner pool; the
jit-kernel-unit-test-amd suite dispatches to mi300/mi325, both gfx942, where
the arch guard would skip every case) and skipped off gfx950.

The padded rows add no memory traffic — the K/V operand is unchanged — and this
kernel is selected-KV bandwidth bound, so the extra MFMA rows occupy issue slots
that were already stalled on HBM. Consistent with that, larger generic KV tiles
and higher warp counts measured neutral-to-slower and are not included.

Accuracy Tests

Against the TileLang kernel it replaces, at the deployed shape (8148 tokens,
H=8, Dqk=576, Dv=512, topk 2048, raw E4M3 Q/KV):

• mean absolute error 8.11e-5, maximum absolute error 0.00586
• zero elements exceeding SGLang's existing atol=0.2, rtol=0.2 gate for this op

Against an independently accumulated FP8 reference, Triton and TileLang show
the identical 0.000772 mismatch fraction. Since the padded path does not
move that fraction at all relative to the kernel it replaces, the residual sits
in reference/production probability rounding rather than in the head padding —
i.e. the padding is numerically inert.

Executed on an MI355X (gfx950:sramecc+:xnack-, ROCm 7.2, torch 2.9.1) against
this branch: 12/12 cases pass (H=4, 8, 12, 16 across all three index
patterns), no skips. A separate case covers the production top-k of 2048, which
is the shape at which the excluded autotune candidate above was found to abort
the compiler on torch 2.11 — note that the timings below were taken on torch
2.9.1, so they do not cover that newer toolchain.

Speed Tests and Profiling

MI355X, seq 8148, topk 2048, Dqk 576, Dv 512, OCP e4m3, 20 warmup + 50
timed iterations. Triton against the TileLang fallback at each head count —
the comparison the dispatch gate actually makes:

H	TileLang	Triton	speedup	vs fallback
4	2.2386 ms	1.3149 ms	1.702x	-41.3%
8	2.2522 ms	1.3170 ms	1.710x	-41.5%
12	2.3062 ms	1.3210 ms	1.746x	-42.7%
16	2.3035 ms	1.3449 ms	1.713x	-41.6%

Triton wins at every head count by a near-constant 1.70–1.75x, so H=4 and H=12
behave exactly like the head counts the original == 16 gate admitted. Both
kernels are near-flat across head count, the expected signature of a
bandwidth-bound operator, and the reason the padded rows are free.

Regression control. The dispatch gate is narrow, but the kernel change also
executes on the shipping H=16 path, so that path was benched with the kernel
installed in the image and with this branch's kernel, back to back:
1.3480 ms → 1.3480 ms, +0.00%.

End-to-end serving, GLM-5.2 MXFP4 on 8x MI355X, TP8/EP1, closed-loop
concurrency 4, 8K input / 1K output, FP8 KV cache, MTP K3 with pinned
acceptance, chunked prefill 8192. Repeated A/B, 4-run control median against
3-run candidate median, for the head-padding change alone:

Metric	TileLang	Triton	Change
Wall time	82.784 s	80.388 s	-2.9%
Output throughput	444.591 tok/s	457.841 tok/s	+3.0%
Median TTFT	429.868 ms	366.780 ms	-14.7%
Median TPOT	8.356 ms	8.176 ms	-2.2%

Candidate runs spanned 80.076–80.406 s against 82.308–82.941 s for the controls.

Future work

The head padding exists because tl.dot needs a 16-row tile on this path. A
Gluon kernel can run BLOCK_M below 16 natively instead — ROCm/aiter#4919 adds
a gfx950 sparse-MLA covering these same GLM-5 prefill shapes and does exactly
that. If SGLang ever dispatches DSA prefill to that kernel, this padding becomes
redundant rather than wrong, and the module docstring records the pointer so a
future reader does not have to rediscover it.

It is not urgent: measured at H=4, 8 and 12 the padded rows add no KV traffic
and cost nothing. Note also that #4919 targets vLLM and consumes those models'
default cache format (fp8 with a per-tensor scale, caller-quantized q), whereas
this ROCm SGLang path stores raw FP8 nope + raw FP8 rope with no scale record —
so it is not a drop-in replacement today. Maintainers who know the intended
direction here, please say so on this PR.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #33835245358
Latest PR Test (Extra): ❌ Run #33835245209
Latest PR Test (AMD ROCm 7.2): ❌ Run #33835245291

[chrisaberger]

Verified on hardware against this branch head (14d8b87)

The numbers in the PR description came from an earlier revision of this patch on
an older base. Re-ran on a real MI355X (gfx950:sramecc+:xnack-, ROCm 7.2,
torch 2.9.1) against the code as it stands on this branch.

Correctness — 6/6 passed, no skips:

test_triton_sparse_mla_raw_fp8[trailing-8]   PASSED
test_triton_sparse_mla_raw_fp8[trailing-16]  PASSED
test_triton_sparse_mla_raw_fp8[strided-8]    PASSED
test_triton_sparse_mla_raw_fp8[strided-16]   PASSED
test_triton_sparse_mla_raw_fp8[head_tail-8]  PASSED
test_triton_sparse_mla_raw_fp8[head_tail-16] PASSED

H=16 regression control. The dispatch gate is narrow, but the kernel change
is not: BLOCK_H, hmask and the masked loads/stores also execute on the
already-shipping H=16 path. Benched the installed kernel and this branch's
kernel back to back, same process shape (seq 8148, Dqk 576, Dv 512, topk
2048, OCP e4m3, 20 warmup + 50 timed iterations):

H=16	median	p10	p90	min
before (installed)	1.3480 ms	1.3444	1.3577	1.3420
after (this branch)	1.3480 ms	1.3436	1.3548	1.3391

+0.00% — at H=16 BLOCK_H is 16 and hmask is all-true, and that is
measurably a no-op rather than only an argued one.

H=8, the newly admitted shape: 1.3133 ms (p10 1.3073, p90 1.3206, min
1.3048), against the 1.280 ms in the description above (+2.6%). Different node
and a uniformly random index distribution rather than the three specific
patterns, so the gap is measurement conditions rather than a discrepancy.

Why the padding is cheap, measured rather than argued. H=8 padded (1.3133 ms)
comes out faster than H=16 (1.3480 ms), despite issuing the identical 16-row
MFMA work. The only difference is the query loads and output stores; the KV
operand is byte-for-byte the same. If the padded rows were expensive, H=8 would
cost what H=16 costs. It does not, which is the bandwidth-bound behaviour this
change relies on.

Not re-measured here: the 1.705x against TileLang (this harness never invokes
the TileLang path) and the end-to-end serving A/B. Those remain as described
above.

One note for reviewers on CI: the test is registered to
stage-b-test-1-gpu-small-amd-mi35x, not jit-kernel-unit-test-amd. The latter
dispatches to mi300/mi325 runners, both gfx942, where _require_gfx950()
would skip every case silently.

[chrisaberger]

Widened the gate to <= 16, on measurements

The gate originally landed as in (8, 16) because those were the only head
counts measured. The kernel itself was already general — it pads to
BLOCK_H = max(16, next_power_of_2(H)) — so the enumerated pair was
under-selling it. Measured the rest of the family on an MI355X
(gfx950:sramecc+:xnack-) and widened the gate accordingly.

Correctness: 12/12 pass — H=4, 8, 12 and 16 across trailing, strided and
head-plus-tail index patterns, against the same reference the test already used.

Median device time by head count (seq 8148, topk 2048, Dqk 576, Dv 512,
OCP e4m3, 20 warmup + 50 timed iterations, all selecting BLOCK_N=32, num_warps=2, num_stages=2):

H	median	p10	p90	vs H=16
4	1.3112 ms	1.3018	1.3184	-2.8%
8	1.3143 ms	1.3065	1.3291	-2.5%
12	1.3193 ms	1.3133	1.3324	-2.2%
16	1.3485 ms	1.3437	1.3521	—

Every padded head count comes out at or below the unpadded H=16 cost —
including H=4, where three quarters of the MFMA tile is padding. The padded
rows add no KV traffic (the K/V operand is identical), and this kernel is
selected-KV bandwidth bound, so the extra rows occupy issue slots that were
already stalled on HBM. That is the justification for gating on the tile size
rather than on an enumerated pair.

Why this is the better gate:

• It fixes the bug class, not one deployment. Any heads / TP < 16 hits the
  same silent TileLang fallback; a 64-head model at TP16 gives H=4 and was
  still excluded under in (8, 16).
• It matches what the kernel already supports, so the gate and the kernel stop
  disagreeing.
• It removes a maintenance trap: the next model with a new head count would
  otherwise need another entry in the tuple, and would silently run slow until
  someone noticed.

The test now parametrizes [4, 8, 12, 16] so the widened gate is covered rather
than asserted.

Related work, for reviewers' awareness: ROCm/aiter#4919 adds a Triton/Gluon
gfx950 sparse-MLA kernel that covers GLM-5-style sparse prefill and handles
H < 16 natively at BLOCK_M = next_pow2(H) rather than padding to 16. If
SGLang eventually dispatches DSA prefill to that kernel, the in-tree Triton path
this PR fixes becomes legacy. Two things make that a later question rather than
a blocker here: that PR targets vLLM and consumes the models' default cache
format (fp8 with a per-tensor scale, plus caller-quantized q), whereas this
ROCm SGLang path stores raw FP8 nope + raw FP8 rope with no scale record, so it
is not a drop-in; and this change needs no new dependency to unblock TP8 today.

[chrisaberger]

Correction: the head-count table above used the wrong baseline

My previous comment compared Triton at H=4/8/12 against Triton at H=16. That
is not the choice the gate makes. For a given head count the gate picks between
this Triton kernel and the TileLang fallback, and H=16 is not an alternative a
workload can select — the model and TP width fix the head count. The H-to-H
table only supports the narrower claim that padding costs nothing; it is not
evidence for widening the gate.

Here is the comparison the gate actually makes. Same inputs handed to both
routes, tilelang_sparse_fwd against triton_sparse_mla_fwd, MI355X
(gfx950:sramecc+:xnack-), seq 8148, topk 2048, Dqk 576, Dv 512, OCP e4m3,
20 warmup + 50 timed iterations:

H	Triton	TileLang	speedup	Triton vs its fallback
4	1.3149 ms	2.2386 ms	1.702x	-41.3%
8	1.3170 ms	2.2522 ms	1.710x	-41.5%
12	1.3210 ms	2.3062 ms	1.746x	-42.7%
16	1.3449 ms	2.3035 ms	1.713x	-41.6%

Triton wins at every head count, by a near-constant 1.70–1.75x. H=4 and
H=12 — the two the original in (8, 16) gate excluded — are no different from
the two it admitted. That is the justification for gating on the MFMA tile size
rather than on an enumerated pair.

Two things worth noting from this run:

The 1.705x in the description reproduces independently. That figure came
from our internal harness at H=8 (TileLang 2.183 ms → Triton 1.280 ms). This
run measures 2.2522 → 1.3170 = 1.710x. Absolute times are ~3% higher on both
sides here (different node, and a uniformly random index distribution rather
than the trailing/strided/head-plus-tail patterns), but the ratio lands within
0.3%. The speed claim in the description is no longer inherited — it has been
re-measured against this branch.

Both kernels are near-flat across head count — TileLang 2.24–2.31 ms,
Triton 1.31–1.34 ms — which is what you would expect if the operator is
selected-KV bandwidth bound in both implementations. It is also why the padding
is free: the padded query rows do not touch the K/V operand, and the K/V operand
is what sets the time.

[chrisaberger]

Closing in favor of #30575, which landed this capability as part of the new triton DSA backend.

This PR widened the gfx950 Triton sparse-MLA prefill dispatch from tp_q_head_num == 16 to any head count that fits the 16-row MFMA tile, padding the query-head tile to BLOCK_H = 16 and masking the padded rows on load and store — so GLM-5.2 at TP8 (64 heads / 8 ranks = 8 heads per rank) reaches the Triton kernel instead of silently falling back to TileLang.

#30575 covers that ground. On current main, python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py computes n_head_blocks = (H + BLOCK_H - 1) // BLOCK_H with BLOCK_H = 16 and masks partial head blocks with h_mask = h_offs < H, and the old SGLANG_DSA_TRITON_PREFILL env gate is replaced by --dsa-prefill-backend triton. I checked that no head-count gate remains on the dispatch path, so there is nothing left here worth maintaining separately.

Thanks @clintg6. For whatever it is worth as corroboration: we measured the padded-head path at 1.70-1.75x over the TileLang fallback at H=4, 8, 12 and 16 (seq 8148, topk 2048) on MI355X, and roughly +10% end-to-end output throughput on GLM-5.2 MXFP4 at TP8, 8K/1K, concurrency 16. The sub-16-head coverage is a real win for anyone serving at those TP widths.