[AMD] Enable Fast Triton Sparse MLA backend by clintg6 · Pull Request #30575 · sgl-project/sglang

clintg6 · 2026-07-09T01:05:05Z

Fast Triton Sparse MLA Kernels for DSA (Prefill + Decode)

Summary

This adds triton as an explicit DSA prefill/decode backend:

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

The new backend provides pure Triton sparse MLA kernels for the fp8 DSA path on ROCm, validated on MI355X (gfx950) and MI300X (gfx942). It replaces the previous Triton prefill env-var gate, SGLANG_DSA_TRITON_PREFILL, with a faster kernel exposed through the normal backend option and adds a fast Triton kernel for the decode path.

What Changed

Adds triton to DSA_CHOICES for --dsa-prefill-backend and --dsa-decode-backend.
Adds startup validation in both server-args paths so Triton DSA is only used with --kv-cache-dtype fp8_e4m3.
Adds standalone Triton prefill and decode dispatch branches in dsa_backend.py.
Ensures HIP Triton DSA uses the raw MLA KV layout (512 nope + 64 rope) instead of the scaled fp8 layout used by other non ROCm backends.
Extends the gfx950 fused-rope skip path to include Triton, matching the TileLang/AITER raw-layout flow.

Kernel Notes

Prefill uses a split-dim 4x128 implementation for the 512-wide nope dimension, plus the 64-wide rope tail.
Prefill can split topk work across CTAs and reduce partial online-softmax state in Triton.
Decode uses an adaptive split-K path with a fused single-pass fast path when kv_splits == 1.
Kernels use exp2 with sm_scale * LOG2E prescaling.

The kernels consume raw fp8 MLA KV cache layout on HIP. They are not intended for NVIDIA's scaled fp8 KV layout.

Validation

GSM8K 400-question validated: TileLang and Triton both produced 95.5% accuracy with 0% invalid answers. Accuracy also verified on gfx942
Isolated prefill microbench on MI355X/GLM-5.2 shape (H=16, D_V=512, D_TAIL=64, topk=2048) showed about 10% speedup for the split-dim prefill kernel versus the earlier 512-wide-dot Triton env gated kernel, with cosine similarity 1.00000.
End-to-end GLM-5.1-MXFP4 serving on TP4 MI355X (ISL=8192, OSL=1024) showed Triton improving throughput by roughly 3-5% versus TileLang in the measured concurrency range.
Performance improvement is seen across all concurrencies, TPs, and for short and long context.

Performance

End-to-end GLM-5.1-MXFP4 serving on TP4 MI355X, concurrency 2:

Backend	Total tok/s	Output tok/s	Mean TTFT (ms)	Mean TPOT (ms)
TileLang	1,005	112	834	17.1
Triton	1,044	116	701	16.6
Improvement	+3.9%	+3.6%	-16.0%	-2.9%

Changed Files

python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py: updates the Triton prefill implementation with split-dim and split-K paths.
python/sglang/srt/layers/attention/dsa/triton_sparse_mla_decode.py: adds the Triton decode implementation.
python/sglang/srt/layers/attention/dsa_backend.py: adds Triton prefill/decode dispatch and removes the old env-var-gated prefill path.
python/sglang/srt/server_args.py: exposes triton as a DSA backend choice.
python/sglang/srt/arg_groups/overrides.py: validates Triton DSA override combinations against fp8 KV cache requirements.
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py: selects raw HIP MLA KV layout for Triton DSA.
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py: includes Triton in the gfx950 fused-rope skip check.

CI States

Latest PR Test (Base): ❌ Run #34293813943
Latest PR Test (Extra): ❌ Run #34293813789
Latest PR Test (AMD ROCm 10): ❌ Run #34293813924

gemini-code-assist · 2026-07-09T01:05:08Z

Warning

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

HaiShaw · 2026-07-09T09:24:31Z

/tag-and-rerun-ci

…nce CSV; rules: known-boxes m12-17 + no-SIGKILL-on-GPU-procs Adds the MI355X (MXFP4) prefill kernel breakdown (tuned MoE, sgl-project#30575/sgl-project#30519/sgl-project#30715 annotations) and the 1k/1k conc64 MI355X-vs-B200 decode reference CSV. Documents the m12-17 GLM-5.2 box and the rule against pkill -9 on ROCm GPU server processes (triggers 100-200GB gpucore dumps that fill the shared disk).

Jacob0226 · 2026-07-16T10:31:38Z

Hi @clintg6
I ran some decode-stage unit tests comparing Tilelang vs Triton across concurrency 4–256. Your Triton kernel performs well at low concurrency, but the advantage fades as concurrency increases.

In the perspective of E2E in SGLang, TTFT actually speedup >10% and it also helps TPOT be faster as waiting for prefill is reduced. TPOT on high conc shows regression on i1k but i8k still benefit from TTFT and Triton attention kernel in i8k take small runtime so it didn't add the show the regression in the end.

From an end-to-end perspective in SGLang, TTFT speeds up by >10%, which in turn helps TPOT since requests spend less time waiting on prefill.
On i1k, high concurrency shows a slight TPOT regression — this tracks with the kernel-level regression above, because at short context the decode attention kernel is a meaningful share of each decode step, so the Triton slowdown shows through.
On i8k, there's no TPOT regression. At long context the decode attention kernel plays a small fraction of the per-step runtime, so the Triton kernel's regression is amortized away. 😊
Test were done on rocm/sgl-dev:v0.5.14-rocm720-mi35x-20260708.

My only open question: is the regression at high concurrency expected behavior, or is there room to improve it?

1am9trash

LGTM

clintg6 · 2026-09-09T21:48:00Z

/rerun-failed-ci

fanxingran · 2026-09-10T05:10:43Z

Hi @clintg6 @1am9trash — sharing independent MI355X validation on GLM-5.2 AgentX (TP4/EP4, conc=8, 3600s, 064fc7c72), which is important for our upcoming GLM-5.2 AgentX benchmark (9/11).

vs TileLang (c8):

P90 Interactivity: 106.6 → 114.4 (+7.3%)
Output tok/s: 347.1 → 356.7 (+2.8%)
TTFT p90: 1.244s → 1.116s
GSM8K: 0.932 → 0.933

Head-to-head vs our parallel PR #38601 (same recipe):

P90 intvty: 114.4 vs 112.9 ([AMD] Enable Fast Triton Sparse MLA backend #30575 slightly better on decode tail)
TTFT p90: 1.116s vs 1.078s ([ROCm] Add Triton sparse-MLA prefill and decode for gfx950 #38601 slightly better on prefill)
Output tok/s: 356.7 vs 357.6 (+0.2%)
GSM8K: 0.933 vs 0.933

All gaps are ~1–3% (noise). AMD internal has reviewed this; we plan to close #38601 once #30575 lands.
Thanks!

HaiShaw · 2026-09-10T05:27:17Z

Error: PR Test Base / base-c-test-8-gpu-b300 - https://github.com/sgl-project/sglang/actions/runs/34293813943/job/102658495822?pr=30575 is addressed in #38588

HaiShaw

LGTM

HaiShaw · 2026-09-10T06:52:34Z

@ispobock @xiezhq-hermann @BBuf @hzh0425 please have a review.

fanxingran · 2026-09-10T07:16:57Z

Error: PR Test Base / base-c-test-8-gpu-b300 - https://github.com/sgl-project/sglang/actions/runs/34293813943/job/102658495822?pr=30575 is addressed in #38588

Hi, we noticed #38588 has been merged. Happy to leave next steps！Thanks!

sgl-project/sglang#30575 exposes the Fast Triton Sparse MLA path as an explicit backend; the 20260910 image already contains it. Co-authored-by: Cursor <cursoragent@cursor.com>

* update config for glm5.2-mxfp4 on mi355x gpu * update pr link * Update changelog link for replacement PR Co-authored-by: Cursor <cursoragent@cursor.com> * Update changelog link for upstream PR Co-authored-by: Cursor <cursoragent@cursor.com> * Raise SGLang /health timeout so eval-only GSM8K survives AITER JIT. Cold mha_varlen compile after ready is ~40s; the default 20s generate health check marked the server UnHealthy and lm_eval hit a dead port. Co-authored-by: Cursor <cursoragent@cursor.com> * Bump the GLM-5.2 MI355X SGLang daily image to 20260910. Co-authored-by: Cursor <cursoragent@cursor.com> * Lower the TP HiCache host-pool ratio from 1.5 to 1.0. Eval and sweep keep the same ratio-based sizing so a green evals-only run at ratio 1.0 remains comparable for this sweep PR. TP4 pinned host DRAM drops from ~679 GB to ~453 GB. Co-authored-by: Cursor <cursoragent@cursor.com> * Keep HiCache write_through_selective instead of restoring write_through. Matches main and the ratio-1.0 eval PR so a green evals-only run is the same host-tier write path as this sweep. Co-authored-by: Cursor <cursoragent@cursor.com> * Rename --cuda-graph-max-bs for the v0.5.19 image. The 20260910 tag splits that flag into prefill and decode names, so the unsuffixed form is rejected as ambiguous and the server never starts. Co-authored-by: Cursor <cursoragent@cursor.com> * Restore HiCache write-through policy Co-authored-by: Cursor <cursoragent@cursor.com> * Use canonical decode CUDA graph limit Co-authored-by: Cursor <cursoragent@cursor.com> * Switch GLM-5.2 DSA prefill/decode to Triton. sgl-project/sglang#30575 exposes the Fast Triton Sparse MLA path as an explicit backend; the 20260910 image already contains it. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix changelog --------- Co-authored-by: LI MOU <lxglbk@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Chun Fang <chun.fang@amd.com>

clintg6 requested review from 1am9trash, Fridge003, HaiShaw, Qiaolin-Yu, YAMY1234, Ying1123, ch-wan, fzyzcjy, hebiao064, hnyls2002, hubertlu-tw, ispobock, kkHuang-amd, merrymercy and rainj-me as code owners July 9, 2026 01:05

clintg6 changed the title ~~Enable Fast Triton Sparse MLA backend~~ [AMD] Enable Fast Triton Sparse MLA backend Jul 9, 2026

sogalin reviewed Jul 9, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/dsa_backend.py Outdated

Comment thread python/sglang/srt/layers/attention/dsa_backend.py

github-actions Bot added the run-ci label Jul 9, 2026

Raiden-Makoto added a commit to Raiden-Makoto/squidward that referenced this pull request Jul 10, 2026

run_glm52.sh: use triton DSA prefill/decode backend (sgl-project#30575)

61a883c

amd-bot mentioned this pull request Jul 11, 2026

[CI Monitor] Daily Report - 2026-07-11 bingxche/sglang-ci-bot#132

Open

HaiShaw and others added 7 commits September 1, 2026 15:04

Merge branch 'main' into feat/triton-sparse-mla

38c26ca

Merge branch 'main' into feat/triton-sparse-mla

2aefc32

Fix linting issue

91f70a4

Fix linting issue

9ae98e7

Merge branch 'main' into feat/triton-sparse-mla

77bf906

Merge branch 'main' into feat/triton-sparse-mla

ca7797c

Merge branch 'main' into feat/triton-sparse-mla

feea2f8

1am9trash mentioned this pull request Sep 8, 2026

[ROCm] Fuse the MLA q absorb into the RoPE + KV-write kernel on gfx950 #38340

Open

clintg6 and others added 2 commits September 8, 2026 14:09

Merge upstream main and resolve conflicts

7182bce

Merge branch 'main' into feat/triton-sparse-mla

064fc7c

HaiShaw assigned alexnails, alisonshao, RolaoDenthu and 1am9trash Sep 9, 2026

1am9trash approved these changes Sep 9, 2026

View reviewed changes

JohnQinAMD mentioned this pull request Sep 10, 2026

[FlyDSL] [Kernel] gfx950 FP8 sparse MLA prefill and decode (FlyDSL) ROCm/aiter#5301

Open

HaiShaw approved these changes Sep 10, 2026

View reviewed changes

HaiShaw merged commit 8a6ab89 into sgl-project:main Sep 10, 2026
562 of 643 checks passed

salexspb mentioned this pull request Sep 10, 2026

[1/N] [AMD] Fix HiSparse slot translation in the fused MLA KV writer #38776

Open

chrisaberger mentioned this pull request Sep 11, 2026

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

Closed

5 tasks

Navigation Menu

[AMD] Enable Fast Triton Sparse MLA backend - #30575

[AMD] Enable Fast Triton Sparse MLA backend#30575
HaiShaw merged 41 commits into
sgl-project:mainfrom
clintg6:feat/triton-sparse-mla

clintg6 commented Jul 9, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 9, 2026

Uh oh!

Uh oh!

Uh oh!

HaiShaw commented Jul 9, 2026

Uh oh!

Jacob0226 commented Jul 16, 2026

Uh oh!

1am9trash left a comment

Uh oh!

clintg6 commented Sep 9, 2026

Uh oh!

fanxingran commented Sep 10, 2026 •

edited

Loading

Uh oh!

HaiShaw commented Sep 10, 2026 •

edited

Loading

Uh oh!

HaiShaw left a comment

Uh oh!

HaiShaw commented Sep 10, 2026

Uh oh!

fanxingran commented Sep 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

Conversation

clintg6 commented Jul 9, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Fast Triton Sparse MLA Kernels for DSA (Prefill + Decode)

Summary

What Changed

Kernel Notes

Validation

Performance

Changed Files

CI States

Uh oh!

gemini-code-assist Bot commented Jul 9, 2026

Uh oh!

Uh oh!

Uh oh!

HaiShaw commented Jul 9, 2026

Uh oh!

Jacob0226 commented Jul 16, 2026

Uh oh!

1am9trash left a comment

Choose a reason for hiding this comment

Uh oh!

clintg6 commented Sep 9, 2026

Uh oh!

fanxingran commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HaiShaw commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HaiShaw left a comment

Choose a reason for hiding this comment

Uh oh!

HaiShaw commented Sep 10, 2026

Uh oh!

fanxingran commented Sep 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

clintg6 commented Jul 9, 2026 •

edited by github-actions Bot

Loading

fanxingran commented Sep 10, 2026 •

edited

Loading

HaiShaw commented Sep 10, 2026 •

edited

Loading