Repository navigation
Conversation
On gfx1250 the MLA prefill took a different path from gfx950 for three small ops: - DeepseekScalingRotaryEmbedding ran forward_native (about nine torch kernels). It now calls aiter's cached-positions rope kernel in place on contiguous cos/sin copies of the module's own cos_sin_cache, rebuilt when the cache grows. The module stays sglang's because ROCm MLA decode paths read cos_sin_cache, which aiter's rope modules do not have. - RMSNorm used the triton kernel. It now uses aiter rmsnorm2d_fwd and rmsnorm2d_fwd_with_add, as on other ROCm targets. - MLA KV rows went through index_put. They now use store_mla/read_mla, as on gfx950. Also drops the stale "aiter rope does not build on gfx1250" comments and adds a gfx1250-only test comparing the aiter rope path with forward_native. Kimi-K2 TP8 on gfx1250, 8k/1k (measured on an earlier base): prefill kernel time per 8192-token step drops from 361.9 to 312.2 ms. At concurrency 64 median TTFT goes from 1417 to 1316 ms (ABBA, two runs each) with TPOT unchanged; at concurrency 4 TTFT does not move because prefill there is bound by CPU kernel launches. GSM8K stays within noise (0.926-0.938). Co-authored-by: Cursor <cursoragent@cursor.com>
The in-place aiter rope kernel would rotate a key passed as its own query twice (IndexerKPool's K-only path), so such calls keep forward_native, as calls with offsets already do. RL on-policy targets force native rope in the base class; the gfx1250 override no longer bypasses that. Rename _aiter_rope_unsupported_arch to _keeps_sglang_rope_modules, since aiter's rope kernel is now used on gfx1250, and add tests for the dtype rebuild and the forward_native fallbacks.
The cached-positions kernel loads both inputs before storing, so a key passed as its own query (IndexerKPool's K-only path) is rotated once. Drop the forward_native fallback for that case and test the aliased call through the kernel instead.
jiaryang
changed the title
jiaryang
marked this pull request as ready for review
November 3, 2026 12:05
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
jiaryang
requested review from
BBuf,
Edwardf0t1,
HaiShaw,
Ying1123,
ch-wan,
fzyzcjy,
ispobock and
merrymercy
as code owners
November 3, 2026 12:06
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
On gfx1250, MLA prefill runs three small ops on different paths from gfx950:
DeepseekScalingRotaryEmbeddingrunsforward_native(about nine torch kernels), RMSNorm uses a gfx1250-only Triton branch, and the MLA KV rows go throughindex_put. aiter's cached-positions RoPE kernel andrmsnorm2d_fwdbuild and run on gfx1250 with the current gfx1250 aiter pin; the RMSNorm row-tail fix for N % 8 != 0 is ROCm/aiter#6348, which the pin includes. This PR extends the gfx950 paths for these three ops to gfx1250.Modifications
DeepseekScalingRotaryEmbedding: on gfx1250 withSGLANG_USE_AITER=1, callaiter.rope_cached_positions_2c_fwd_inplacein place. The module stays sglang's because ROCm MLA decode paths readcos_sin_cache, which aiter's rope modules do not have. The kernel reads cos/sin without strides, so it gets contiguous query-dtype copies of the module's owncos_sin_cache, built on first use and rebuilt when the cache or the query dtype changes.forward_nativeis kept for calls withoffsetsand whenrl_on_policy_targetforces native rope. A key passed as its own query (theIndexerKPoolK-only path) is rotated once, since the kernel loads both inputs before storing.rmsnorm2d_fwd/rmsnorm2d_fwd_with_add, as on other ROCm targets. This applies to every model on gfx1250 withSGLANG_USE_AITER=1, in prefill and decode.store_mla/read_mlaon gfx1250, as on gfx950.test/registered/rotary/test_rope_rocm.pythat compares the aiter RoPE path withforward_native(neox and GPT-J styles, strided q_pe/k_pe views shaped like the MLA callers, in-place output, cache growth, dtype change, a key passed as its own query, and theforward_nativefallbacks).gfx950 is unchanged. On gfx1250, RoPE now runs in place with query-dtype cos/sin, the same as the aiter rope module on gfx950. MLA decode RoPE on gfx1250 also moves to this kernel, except where #44459 fuses it into the decode kernel; the measurements below were taken with #44458 and #44459 applied, so this PR should land after #44459.
Validated on gfx1250 with Kimi-K2-Thinking-MXFP4 at TP8.
Accuracy Tests
GSM8K, 8-shot, 1319 questions, 3 runs:
Speed Tests and Profiling
bench_serving, 8k input / 1k output, FP8 KV cache, aiter decode backend, page size 64.In the prefill profile (8192-token chunks, GPU kernel time excluding all-reduce, mean over 8 ranks), a chunk goes from 175.0 ms to 153.6 ms at concurrency 4. Elementwise kernels go from 25.2 to 10.0 ms and RMSNorm from 19.4 to 15.6 ms; the new RoPE kernel takes 1.0 ms. Decode kernel time is unchanged.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #41756161586
Latest PR Test (Extra): ❌ Run #41756161550
Latest PR Test (AMD ROCm 10): ✅ Run #41756161600
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: