[ROCm] Add Triton sparse-MLA prefill and decode for gfx950

[fanxingran]

Wire --dsa-prefill-backend triton and --dsa-decode-backend triton on gfx950 with fp8 KV. Prefill uses a fused Triton kernel with XCD swizzle and 64-bit KV offsets; decode adds a separate Triton path for MTP-shaped forwards. Shape-gate misses fall back to TileLang.

Motivation

Sparse MLA is a major GPU cost on gfx950 DSA serving (prefill chunks and every MTP draft/verify decode step). This PR adds triton as an explicit DSA backend for both prefill and decode on gfx950 + fp8 KV:

--dsa-prefill-backend triton --dsa-decode-backend triton

Default remains TileLang until those flags are set. Prefill and decode each have a shape gate; misses log triton_shape_gate and fall back to TileLang so short warmup / out-of-spec batches stay safe.

Replaces the earlier single-file triton_sparse_mla.py stub with split prefill/decode modules registered in the kernel inventory.

Modifications

• Prefill kernel python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py — fused per-query sparse-MLA prefill. XCD-swizzled pid → token bijection, softmax in log2 space, reciprocal-once denominator, unmasked KV loads after page clamp, next-block index prefetch, 64-bit KV offsets past the int32 wrap. Fixed serving config: BLOCK_N=64, num_warps = clamp(H_PAD // 16, 1, 4), num_stages=1 (autotune hits a Triton 3.7.0 LLVM abort on part of the grid).
• Decode kernel python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py — split-K sparse-MLA decode for gfx950 fp8 KV. Tiling (pick_n_groups / pick_block_n_warps / pick_num_stages) targets MTP token counts; q is two strided halves of the fused [T, H, 576] tensor (no extra .contiguous()).
• Serving hook dsa_backend.py — --dsa-prefill-backend triton / --dsa-decode-backend triton in DSA_CHOICES. Construction requires gfx950 + fp8 KV. Prefill shape gate: 8 or 16 heads, d_v=512, rope tail 64, topk 2048, 512–32768 tokens. HIP KV layout selector lists triton next to tilelang / aiter.
• CLI / hisparse — exec_.py, hisparse_hook.py, kv_cache_configurator.py, forward_mla_rocm.py wire the new backend name.
• Tests — padding / XCD bijection / 64-bit KV offset (prefill); decode vs fp32 and vs TileLang; backend shape gates; CLI registry for triton.

Accuracy Tests

Kernel SNR floors (registered tests):

Check	Floor
Prefill vs fp32 reference (padding / fully-padded / interleaved / single-valid-key)	> 30 dB
Prefill vs TileLang, production shapes (seq 1696 / 4096 / 16384 / 32768 × H 8 / 16)	95.5–100.2 dB
Decode vs fp32 oracle	32.2–32.4 dB (floor 30 dB)
Decode vs TileLang	> 30 dB

Bit-exact match is not expected: fp8 accumulation order differs.

Serving (upstream tree, MI355X, TP4/EP4, GLM-5.2-MXFP4, EAGLE MTP, fp8 KV, --dsa-topk-backend sgl-kernel):

Arm	GSM8K (512 threads)
TileLang / TileLang	0.932 (GLM fork baseline, same recipe)
Triton / Triton	0.925

Decode used Triton in production-shaped batches; prefill fell back to TileLang on small token counts (reason=triton_shape_gate), which matches the prefill seq gate design.

Speed Tests and Profiling

Kernel — prefill

MI355X, seq=32768, H=16 (TP4), fp8 KV, topk=2048, production index distribution. Alternating in-process A/B after a 1.5 s clock ramp.

Context (KV working set)	TileLang	Triton	Speedup
32768 tok (18 MB)	9040 µs	4460 µs	2.03x
65536 tok (36 MB)	9067 µs	5087 µs	1.78x
98304 tok (54 MB)	9080 µs	5313 µs	1.71x

Four 32k chunks of a 100k prompt: 36.2 ms → 20.2 ms, 1.79x.

Kernel — decode

Idle MI355X, ctx=100k, auto tiling, µs/call from 16-call timed_graph. 24-point sum: TileLang 510, Triton 293 (1.74x).

    TP8 (H=8)                            TP4 (H=16)
    C  role     T  floor    TL        Triton   floor    TL        Triton
    1  draft    1    2.9  14.3         6.9     2.9  14.0         7.3
    1  verify   6    3.1  14.9         8.0     3.0  15.1         8.8
    4  verify  24    4.8  21.9        11.3     4.9  22.8        13.1
    10 verify  60   14.0  36.7        19.8    13.9  37.3        21.8
    14 verify  84   18.8  39.8        26.8    18.7  40.4        27.7

Serving — AgentX 900 s (GLM-5.2-MXFP4, same node)

TP4/EP4, MTP (EAGLE n=5), HiCache, simulate_acc=3.61. Baseline = tilelang/tilelang; Triton = triton/triton. Error rate 0 on every cell.

Conc	P90 intvty TL	P90 intvty Triton	Δ	out tok/s TL	out tok/s Triton	Δ
8	85.97	103.45	+20.3%	366.02	380.19	+3.9%
10	63.03	73.40	+16.5%	536.55	552.69	+3.0%
14	41.63	52.72	+26.6%	512.27	554.50	+8.2%

Conc	TTFT p50 TL → Triton	TTFT p90 TL → Triton	ITL p90 TL → Triton
8	0.426 → 0.393 s	1.445 → 0.993 s (−31%)	11.63 → 9.67 ms (−17%)
10	0.482 → 0.452 s	2.223 → 1.413 s (−36%)	15.86 → 13.62 ms (−14%)
14	0.560 → 0.534 s	2.474 → 2.176 s (−12%)	24.02 → 18.97 ms (−21%)

Official 3600 s TP4/EP4 sweep (tilelang vs both_triton, conc 8/10) shows similar trends (+0.9% to +1.6% out tok/s, TTFT −3.8% to −7.5%); full 3600 s table to be attached after the node sweep completes.

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.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #34303074580
Latest PR Test (Extra): ❌ Run #34303074357
Latest PR Test (AMD ROCm 7.2): ❌ Run #34303074600