[AMD] [GLM5] Add opt-in Triton fp8 sparse-MLA prefill kernel for gfx950

[Raiden-Makoto]

Motivation

This affects HIP serving of DSA (DeepSeek Sparse Attention) models; it was found
and validated on GLM-5.1-MXFP4 (which uses the DSA indexer) on MI350X.

On the HIP fp8 path, tilelang_sparse_fwd runs the same TileLang
partial+combine kernel for prefill as for decode. The sparse-MLA attention tile
is tiny — M = 16 heads (one 16x16 MFMA) — so the 256-thread (4-warp)
TileLang block over-parallelizes it and spends most of its time on intra-block
coordination rather than the matmuls (profiling at the prefill config on gfx950
shows the kernel VALU-bound at ~42% with MFMA at ~12% and ~37% occupancy).

A per-query Triton flash kernel with a small 2-warp / BLOCK_N=32 tile fits
the problem shape and saturates the GPU on block count instead. It also reads
q_nope/q_rope directly, skipping the per-layer concat (the kernel splits q
into main/tail internally anyway, so combining them first is wasted work — and
the HIP path uses a plain torch.cat, not the fused CUDA concat). Decode is
left on TileLang.

Modifications

• New python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py: a per-query
  Triton flash kernel over the indexer-selected topk KV. Autotuned over
  BLOCK_N/num_warps/num_stages; guards an all-masked query row against
  NaN (finite softmax shift when the row has no valid key). Reads q_nope
  (width d_v) and q_rope (width dim-d_v) as two separate tensors.
• python/sglang/srt/layers/attention/dsa_backend.py: in forward_extend, route
  the fp8 prefill path to the Triton kernel and pass q_nope/q_rope directly,
  skipping concat_mla_absorb_q_general.
  • Opt-in, default off: enable with SGLANG_DSA_TRITON_PREFILL=1.
  • Gated to gfx950 + the validated shape (num_heads==16, d_v==512,
    tail==64, topk==2048); other archs/shapes use TileLang.
  • Prefill only; decode stays on TileLang.

Accuracy Tests

GSM8K 5-shot, full 1319 questions, GLM-5.1-MXFP4, MI350X (gfx950), tp4,
fp8_e4m3 KV:

	accuracy
TileLang prefill (default)	0.938
Triton prefill (SGLANG_DSA_TRITON_PREFILL=1)	0.941

No regression.

Speed Benchmarks

E2E sglang.bench_serving (random, input 8192 / output 1024, median latencies),
GLM-5.1-MXFP4 / MI350X / tp4 / fp8_e4m3 KV.

Baseline (TileLang prefill):

concurrency	total tok/s	TTFT (ms)	ITL (ms)	E2EL (ms)
2	944	920	17.74	19491
4	1717	1752	18.56	21436
8	2847	2887	20.66	25895
16	4343	5116	24.23	33946
32	6233	9588	28.68	47438
64	8158	18558	35.54	72308

This PR (SGLANG_DSA_TRITON_PREFILL=1):

concurrency	total tok/s	TTFT (ms)	TTFT delta	ITL (ms)	E2EL (ms)
2	957	798	-13.3%	17.71	19228
4	1738	1511	-13.8%	18.58	21126
8	2922	2486	-13.9%	20.73	25217
16	4515	4410	-13.8%	24.22	32655
32	6574	8276	-13.7%	28.56	44814
64	8750	16051	-13.5%	35.36	67431

Median TTFT improves 13.3-13.9% across the sweep; median ITL is unchanged
(within +/-0.5%) and median E2EL is flat-to-better, confirming the change is
prefill-side with no decode regression.

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 #27994768455
Latest PR Test (Extra): ⏳ Run #27994768350

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Raiden-Makoto]

/tag-and-rerun-ci

[HaiShaw]

LGTM

[HaiShaw]

@Raiden-Makoto please add a follow-up PR to update GLM cookbook.
cc @1am9trash

[amd-oshkarav]

Guys, @Raiden-Makoto,
I have tried the feature on GLM 5.2 FP8 which theoretically should work, but the server cannot be started when kv is fp8_e4m3. Could you check up if I do something wrong, or indeed we have an issue?
Here are the details.
Image lmsysorg/sglang-rocm:v0.5.13.post1-rocm720-mi35x-20260624 as it is,
Server start script:

#!/usr/bin/env bash
export SGLANG_USE_AITER="1"
export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
export SAFETENSORS_FAST_GPU="1"
export SGLANG_DSA_TRITON_PREFILL="1"
python -m sglang.launch_server \
  --model zai-org/GLM-5.2-FP8 \
  --host 0.0.0.0 \
  --port 5060 \
  --tp-size 8 \
  --served-model-name "glm-5.2-fp8" \
  --trust-remote-code \
  --watchdog-timeout 12000 \
  --attention-backend dsa \
  --dsa-prefill-backend aiter \
  --dsa-decode-backend tilelang \
  --reasoning-parser glm45 \
  --tool-call-parser glm47 \
  --disable-radix-cache \
  --kv-cache-dtype fp8_e4m3 \
    2>&1 | tee serve.log

Error at the end of server start:

[AITER] /sgl-workspace/aiter/aiter_meta/csrc/py_itfs_cu/asm_mla.cu:227 mla_decode_stage1_asm_fwd: fp8 Q requires q_scale and kv_scale
Fatal Python error: Aborted

Also, when I remove --kv-cache-dtype fp8_e4m3, the server can be started and does not garbage on a simple request.

[Raiden-Makoto]

@amd-oshkarav two things:

1. Wrong prefill backend. You passed --dsa-prefill-backend aiter, but the Triton prefill kernel from this PR only runs under the tilelang backend (--dsa-prefill-backend tilelang, which is the default). With aiter, the Triton kernel never engages — so SGLANG_DSA_TRITON_PREFILL=1 is a no-op in your run and this PR's code path isn't exercised at all. Use --dsa-prefill-backend tilelang (and keep --kv-cache-dtype fp8_e4m3).

2. There is a real bug, but it's separate from this PR. With --dsa-prefill-backend aiter, prefill/extend goes to _forward_aiter_extend, which calls the aiter asm MLA kernel. That code only ever sets kv_scale and leaves q_scale = None:

q_scale = None
kv_scale = None
if kv_cache.dtype == fp8_dtype:
    kv_scale = torch.ones((), dtype=torch.float32, device=q_kernel.device)

GLM-5.2-FP8's MLA query is fp8, so the aiter kernel hits its fp8 Q requires q_scale and kv_scale guard and aborts. When you drop --kv-cache-dtype fp8_e4m3 the if is skipped → no crash, which matches what you saw. This comes from the aiter DSA path added in #26639, not from this PR.