[Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for MLA decode on gfx1250 by jiaryang · Pull Request #44459 · sgl-project/sglang

jiaryang · 2026-10-15T09:50:54Z

Motivation

On gfx950, aiter MLA decode runs RoPE, the q/k concat and the KV-cache write in one fused_qk_rope_cat_and_cache_mla kernel. On gfx1250 these are still separate: native RoPE, two torch.cat and set_kv_buffer, more than a dozen elementwise kernels per layer. This PR extends the gfx950 fused branch to gfx1250.

Modifications

With SGLANG_USE_AITER=1 on gfx1250, MLA decode with the aiter or triton backend skips the standalone RoPE and takes the fused branch, as aiter MLA does on gfx950. gfx950 is unchanged.

As on gfx950, the gate is evaluated per call: decode, the aiter or triton backend, no DSA, one DCP rank, a rope module of exactly RotaryEmbedding or DeepseekScalingRotaryEmbedding, and a bf16 or FP8 KV pool. The pool check reads the pool's dtype, so an fp16 pool, which the kernel rejects, stays on the unfused path.
sglang's rope keeps cos and sin in one cos_sin_cache; the kernel gets the two halves as strided views.
Q keeps q_nope_out's dtype, which is bf16 on gfx1250: MLA decode there takes an unquantized Q even over an FP8 KV cache.
Prepare's rope-skip condition moves into _applies_rope_in_prepare, so the unit test can check that prepare skips RoPE exactly when core fuses it.
CPU unit tests are added to test/registered/unit/models/test_rocm_mla_absorb_fusion_gate.py.

Requires ROCm/aiter#5501, which both ROCm Docker pins in docker/rocm.Dockerfile already contain. On main, gfx1250 MLA decode runs the triton backend, which the gate covers, and the first set of numbers below is main with the triton decode backend. With #44458's aiter decode backend the same fused branch applies; the second table measures that.

Validated on gfx1250 with Kimi-K2-Thinking-MXFP4 at TP8, with the triton decode backend on main and the aiter decode backend from #44458.

Accuracy Tests

The fused kernel matches the unfused path (q and the written KV rows) for bf16 and FP8 KV, batch 1 to 128.

GSM8K, 8-shot, 1319 questions, triton decode backend on main:

	KV cache	Runs	Accuracy
Before	FP8	4	0.942 (0.950, 0.939, 0.939, 0.939)
After	FP8	4	0.940 (0.942, 0.938, 0.934, 0.945)
Before	bf16	1	0.947
After	bf16	1	0.938

Unchanged main moves by 0.011 between server launches (0.950 and 0.939), so the single bf16 run is within that spread.

Speed Tests and Profiling

RoPE, concat and KV write per layer, Kimi-K2 TP8 shapes, CUDA graph replay:

batch	unfused (us)	fused (us)
1	37.6	1.9
4	40.8	1.6
32	55.6	1.7
128	81.6	1.9

bench_serving, 8k input / 1k output, FP8 KV, triton decode backend on main:

Concurrency		Throughput per GPU (tok/s/GPU)	Interactivity (tok/s/user)	Median TTFT (ms)	Median TPOT (ms)
4	Before	171.4	39.7	416.8	25.20
4	After	217.7	51.2	406.4	19.54
4	Change	+27.0%	+29.0%	-2.5%	-22.5%
16	Before	398.8	22.9	440.6	43.63
16	After	483.4	28.0	427.1	35.68
16	Change	+21.2%	+22.3%	-3.1%	-18.2%
64	Before	647.9	9.4	677.0	106.87
64	After	745.6	10.7	518.1	93.66
64	Change	+15.1%	+14.1%	-23.5%	-12.4%

In the decode profile (concurrency 4, 8k context), each layer's RoPE, concat and KV write go from about 14 PyTorch kernels taking about 82 us to the fused kernel taking about 3 us, and the decode step goes from 30.0 ms to 23.6 ms. The triton decode attention kernels are unchanged.

Concurrency 4 with #44458's aiter decode backend on both sides:

	Throughput per GPU (tok/s/GPU)	Interactivity (tok/s/user)	Median TTFT (ms)	Median TPOT (ms)
Before	304.7	71.4	271.3	14.00
After	443.1	104.9	262.6	9.53
Change	+45.4%	+46.9%	-3.2%	-31.9%

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. (No user-facing option changes.)
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

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
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
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 #41622351290
Latest PR Test (Extra): ❌ Run #41622351313
Latest PR Test (AMD ROCm 10): ✅ Run #41622351159

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

jiaryang · 2026-10-22T13:48:09Z

/rerun-failed-ci

jiaryang · 2026-10-23T00:19:44Z

/rerun-failed-ci

jiaryang · 2026-10-28T01:37:53Z

/rerun-failed-ci

jiaryang · 2026-10-28T10:01:55Z

/rerun-failed-ci

jiaryang · 2026-10-29T01:00:20Z

/rerun-failed-ci

jiaryang · 2026-10-29T02:02:59Z

@amd-bot ci-status

bingxche · 2026-10-29T02:03:24Z

@jiaryang

CI Status for PR #44459

Warning

The Claude Code agent did not work, so the CI status check could not be produced.
Error: git worktree add failed: fatal: invalid reference: HEAD
No API fallback is used. Re-run the command, or check the amd-bot runners.

Generated by amd-bot

jiaryang · 2026-10-30T08:18:23Z

/rerun-failed-ci

With SGLANG_USE_AITER=1 on gfx1250, aiter or triton MLA decode skips the standalone RoPE and uses aiter fused_qk_rope_cat_and_cache_mla, the same fused branch aiter MLA takes on gfx95, with a bf16 Q. sglang's RotaryEmbedding keeps cos and sin in one cos_sin_cache, so the kernel gets its two halves as strided views. Limited to plain rope modules whose table covers the whole rope head, without DSA or DCP, and to bf16 or fp8 e4m3 KV caches. Requires aiter with ROCm/aiter#5501. The layer-static terms of the gate are computed once in __init__. A CPU unit test pins down each term of the gate and the cos/sin order of the tables. Co-authored-by: Cursor <cursoragent@cursor.com>

Address review feedback. - The gate is evaluated per call as on gfx95 instead of caching the layer-static terms in __init__, so init_mla_rocm_forward is removed. - The rope check is the exact type RotaryEmbedding or DeepseekScalingRotaryEmbedding; subclasses such as Bailing MoE v3's video mrope compute their rope differently. - The KV check reads the pool's logical dtype (bf16 or fp8) instead of the kv_cache_dtype server arg, so fp8-quantized checkpoints and draft runners whose server arg is "auto" are covered, and an fp16 pool stays unfused. - Q keeps q_nope_out's dtype on gfx1250, which is bf16 there. - Prepare's rope-skip condition moves into _applies_rope_in_prepare so the tests can check that prepare skips rope exactly when core fuses it. - The CPU cases join test_rocm_mla_absorb_fusion_gate.py, and test_rocm_mla_gfx1250_fused_rope_gate.py is removed. Co-authored-by: Cursor <cursoragent@cursor.com>

main moved ForwardMode into sglang.srt.model_executor.forward_mode and its check-indirect-imports hook rejects importing it through forward_batch_info. Co-authored-by: Cursor <cursoragent@cursor.com>

On gfx1250 with SGLANG_USE_AITER=1, _fused_rope_cat_and_cache takes the gfx1250 branch and reads rotary_emb.cos_sin_cache, which the test's stub lacks. That branch excludes DSA, so the test's DSA pools pin _use_gfx1250_fused_mla_rope to False.

jiaryang · 2026-11-03T12:22:04Z

/rerun-failed-ci

kkHuang-amd

LGTM

jiaryang · 2026-11-04T08:37:22Z

/rerun-failed-ci

github-actions Bot added the amd label Oct 15, 2026

jiaryang force-pushed the amd/gfx1250-mla-fused-rope-kv-cache branch from 68175b3 to f142e83 Compare October 20, 2026 15:55

github-actions Bot added the deepseek label Oct 20, 2026

jiaryang marked this pull request as ready for review October 20, 2026 15:58

jiaryang requested review from Fridge003, ch-wan, fzyzcjy, ispobock and merrymercy as code owners October 20, 2026 15:58

sogalin added the run-ci CI: run the baseline test suite on this PR label Oct 22, 2026

jiaryang force-pushed the amd/gfx1250-mla-fused-rope-kv-cache branch from 217506b to b377194 Compare October 22, 2026 16:26

jiaryang changed the title ~~[Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for triton MLA decode on gfx1250~~ [Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for MLA decode on gfx1250 Oct 22, 2026

mqhc2020 mentioned this pull request Oct 29, 2026

[Perf] [AMD] attention: Switch MLA decode to Aiter on gfx1250 #47193

Draft

5 tasks

jiaryang and others added 2 commits November 3, 2026 06:45

jiaryang force-pushed the amd/gfx1250-mla-fused-rope-kv-cache branch from e99dec6 to 59901a8 Compare November 3, 2026 06:51

This comment has been minimized.

Sign in to view

jiaryang and others added 4 commits November 3, 2026 07:15

Merge branch 'main' into amd/gfx1250-mla-fused-rope-kv-cache

245dc70

[Test] Import ForwardMode from forward_mode

1240927

main moved ForwardMode into sglang.srt.model_executor.forward_mode and its check-indirect-imports hook rejects importing it through forward_batch_info. Co-authored-by: Cursor <cursoragent@cursor.com>

Merge branch 'main' into amd/gfx1250-mla-fused-rope-kv-cache

0e8b07e

jiaryang mentioned this pull request Nov 3, 2026

[Perf] [AMD] mla: Use the gfx950 aiter RoPE, RMSNorm and MLA KV-row paths on gfx1250 #48154

Open

4 of 5 tasks

kkHuang-amd approved these changes Nov 4, 2026

View reviewed changes

Navigation Menu

Repository navigation

[Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for MLA decode on gfx1250 - #44459

[Perf] [AMD] mla: Fuse RoPE, concat and KV-cache write for MLA decode on gfx1250#44459
jiaryang wants to merge 6 commits into
sgl-project:mainfrom
jiaryang:amd/gfx1250-mla-fused-rope-kv-cache

jiaryang commented Oct 15, 2026 •

edited

Loading

Uh oh!

jiaryang commented Oct 22, 2026

Uh oh!

jiaryang commented Oct 23, 2026

Uh oh!

jiaryang commented Oct 28, 2026

Uh oh!

jiaryang commented Oct 28, 2026

Uh oh!

jiaryang commented Oct 29, 2026

Uh oh!

jiaryang commented Oct 29, 2026

Uh oh!

bingxche commented Oct 29, 2026

Uh oh!

jiaryang commented Oct 30, 2026

Uh oh!

This comment has been minimized.

jiaryang commented Nov 3, 2026

Uh oh!

kkHuang-amd left a comment

Uh oh!

jiaryang commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

jiaryang commented Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

jiaryang commented Oct 22, 2026

Uh oh!

jiaryang commented Oct 23, 2026

Uh oh!

jiaryang commented Oct 28, 2026

Uh oh!

jiaryang commented Oct 28, 2026

Uh oh!

jiaryang commented Oct 29, 2026

Uh oh!

jiaryang commented Oct 29, 2026

Uh oh!

bingxche commented Oct 29, 2026

CI Status for PR #44459

Uh oh!

jiaryang commented Oct 30, 2026

Uh oh!

This comment has been minimized.

jiaryang commented Nov 3, 2026

Uh oh!

kkHuang-amd left a comment

Choose a reason for hiding this comment

Uh oh!

jiaryang commented Nov 4, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jiaryang commented Oct 15, 2026 •

edited

Loading