[fengshaoyi-amd]

Motivation

On ROCm, Qwen3 prefill runs QK-norm and RoPE as five launches per layer: apply_qk_norm copies the strided q and k views of the packed qkv to contiguous tensors, runs two aiter RMSNorms, then rotary_emb launches rotary_embedding. The fused path in apply_qk_norm is CUDA-only. In a Qwen3-8B prefill profile on MI355X (batch 1/8, input 2048/8192, main 7fe5b127e4) these kernels take 4.8% to 6.8% of GPU kernel time.

The diffusion JIT kernel fused_inplace_qknorm_rope already does this in one in-place launch (strided q/k, GQA, fp32 cos/sin cache table), but it does not compile on HIP: bare __shfl_sync with a 32-bit mask hits HIP's 64-bit-mask static_assert. So the ROCm path enabled in #35573 silently falls back to the unfused ops; the kernel's tests are registered for CUDA CI only. #40585 fixes the same shuffles for MiniMax-H3; this PR carries an equivalent fix so it does not depend on it. #47001 (open) adds the same in-place fusion for XPU through sgl-kernel-xpu, and #47785 (open) a CUDA prefill fusion; all three change forward_prepare_native, so whichever lands later needs a small rebase.

Modifications

Scope: ROCm only. On HIP the kernel's NeoX variants with round_norm_before_rope=True now build. Two ROCm callers request them: Qwen3 attention, which also serves the models built on Qwen3Model (such as Qwen3-Embedding and InternVL) when the gate below accepts them, and ERNIE-Image, which keeps the fused kernel only if its first call matches its reference bit for bit. NeoX with round_norm_before_rope=False and KV packing stay off on HIP, so their callers, such as Cosmos3 video, keep the separate ops as on main. GPT-J variants, CUDA, XPU, NPU, and the aiter decode fusion are unchanged.

• qknorm_rope.cuh: on HIP, shuffle with __shfl / __shfl_xor at width 32 (the kernel's logical warp on wave64). CUDA builds are unchanged.
• qknorm_rope_jit.py: on HIP, can_use_fused_inplace_qknorm_rope returns False and a direct call raises NotImplementedError for the two variants without AMD numerical coverage, KV packing and NeoX with round_norm_before_rope=False.
• models/utils.py: add can_use_jit_qk_norm_rope, which gates on HIP, --enable-fused-qk-norm-rope not set to off, plain RotaryEmbedding with an fp32 cache, plain RMSNorm weights in the activation dtype, and no deterministic / on-policy / inductor prefill mode, and builds the kernel at init so it is never first built during graph capture; add apply_qk_norm_rope, which runs the kernel in place with round_norm_before_rope=True or falls back to apply_qk_norm + rotary_emb.
• qwen3.py: outside the CUDA decode branch from [Perf] [NVIDIA] qwen3: Reuse fused QK-norm/RoPE kernel for decode #46908, forward_prepare_native calls apply_qk_norm_rope. On ROCm it serves prefill / extend, and decode when the aiter decode fusion from [Perf] [AMD] qwen3: Fuse QK-norm + RoPE + KV cache write for dense Qwen3 decode #44267 does not apply.
• test_qk_norm_rope_hip.py (new, AMD CI): fused vs unfused path on strided views over head/rotary dims, GQA ratios, neox and GPT-J, 0 to 4096 tokens, fp16, int32 positions; the fused result must be no farther from an fp64 reference than the unfused one; the kernel gate rejects the two variants above on HIP; YaRN, mRoPE, eps mismatch, fp32 norm weights, deterministic, on-policy, inductor and --no-enable-fused-qk-norm-rope fall back.
• test_qknorm_rope_out_of_place.py: register for AMD CI; skip the two-rows-per-warp comparison on HIP, which builds only the one-row-per-warp kernel.

Accuracy Tests

MI355X (gfx950), torch 2.11.0+rocm10.0.0, sglang-kernel 0.4.10, aiter 285579b62, SGLANG_USE_AITER=1, decode CUDA graphs on. Qwen/Qwen3-8B bf16, otherwise default server args. Base is main 7fe5b127e4. GSM8K with python -m sglang.test.run_eval --eval-name gsm8k --num-examples 1319 --temperature 0 --max-tokens 4096 (1314 questions after the 5 few-shot examples):

Attention backend	Base	This PR
aiter	0.9444 / 0.9422 / 0.9452 (mean 0.9439)	0.9482 / 0.9475 / 0.9460 (mean 0.9472)
triton	0.9467	0.9437

The fused kernel reads the fp32 cos/sin table, while the unfused HIP RoPE rounds it to bf16 first, so outputs differ by 1 to 2 ulp and are closer to fp64.

Unit tests, AMD CI (MI300X) on this head: jit-kernel-unit-test-amd runs test_qk_norm_rope_hip.py (130 passed) and test_qknorm_rope_out_of_place.py (1 passed, 12 skipped); stage-a-test-1-gpu-small-amd runs test_qwen3_aiter_fused_qk_norm_rope.py (decode path, 6 passed). With main's qknorm_rope.cuh the kernel fails to build for head_dim 128 neox, so can_use_jit_qk_norm_rope returns False and test_qk_norm_rope_hip.py fails at its gate assertion.

MI325X (gfx942), torch 2.11.0+rocm10.0.0, aiter 285579b62, main cb55d77d37 with this patch: GSM8K Qwen3-8B 0.9437 on main vs 0.9475 here.

--enable-torch-compile already fails on main with a memory fault during decode graph capture in the inductor output for aiter's decode op fused_qk_norm_rope_cache_pts_quant_shuffle, which this PR does not touch, so torch.compile is not covered.

Speed Tests and Profiling

Microbench of one Qwen3-8B layer on this PR's code, apply_qk_norm + rotary_emb vs apply_qk_norm_rope(use_jit=True) (32 q / 8 kv heads, head_dim 128, bf16, q/k strided views of qkv), CUDA events, median of 50 after 10 warmup:

Tokens	Current: copies + 2 RMSNorm + RoPE (µs)	Fused (µs)	Speedup
2048	66.6	23.8	2.80x
8192	156.9	48.0	3.27x
16384	301.8	80.1	3.77x
65536	1194.5	436.0	2.74x

Prefill profile, python -m sglang.benchmark.one_batch --model-path Qwen/Qwen3-8B --attention-backend aiter --batch-size <B> --input-len <ISL> --output-len 1 --profile --profile-stage prefill, one process per shape, GPU kernel time summed over 36 layers:

Batch × input	Kernel total, base → PR (ms)	QK-norm + copies + RoPE, base → PR (ms)	Launches
1 × 2048	28.770 → 27.876 (−3.1%)	1.959 → 0.506	180 → 36
1 × 8192	119.697 → 116.452 (−2.7%)	6.357 → 1.985	180 → 36
8 × 2048	203.402 → 198.568 (−2.4%)	12.198 → 3.926	180 → 36
8 × 8192	986.627 → 960.753 (−2.6%)	47.313 → 15.531	180 → 36

End to end, Qwen3-8B on one MI355X, with a fresh server for every run:

SGLANG_USE_AITER=1 python -m sglang.launch_server --model-path Qwen/Qwen3-8B --tp 1 \
  --attention-backend aiter --disable-radix-cache
python -m sglang.bench_serving --backend sglang --dataset-name random --random-range-ratio 1.0 \
  --random-input-len <ISL> --random-output-len <OSL> --max-concurrency <C> --num-prompts <N> \
  --seed 1 --warmup-requests 8

<N> is 64 at concurrency 4 and 256 at concurrency 16; at concurrency 64 it is 1024 for W2 and 320 for W3. Blocks of 4 runs in the order base, PR, PR, base; the first run of each block is dropped and each block reports the median of runs 2-4. Base and PR are the means of their two blocks; noise is the second base block vs the first.

Workload (input / output, concurrency)	Median TTFT, base → PR (ms)	Input throughput	Median TPOT, base → PR (ms)	Noise, TTFT / throughput
W1 4096 / 64, c4	222.7 → 217.2 (−2.5%)	+0.6%	9.08 → 9.06 (−0.3%)	0.0% / +0.6%
W1 4096 / 64, c16	518.5 → 504.3 (−2.7%)	+1.9%	11.44 → 11.30 (−1.3%)	+0.7% / −0.2%
W1 8192 / 128, c4	395.0 → 382.5 (−3.2%)	+1.3%	10.31 → 10.24 (−0.7%)	−0.3% / +0.1%
W1 8192 / 128, c16	1125.5 → 1095.0 (−2.7%)	+1.6%	13.93 → 13.79 (−1.0%)	−0.2% / +0.2%
W2 2048 / 256, c16	285.5 → 274.1 (−4.0%)	+0.7%	5.52 → 5.51 (−0.1%)	−0.2% / +0.2%
W2 2048 / 256, c64	1075.6 → 1034.2 (−3.9%)	+1.9%	11.18 → 11.04 (−1.3%)	−1.1% / +1.0%
W3 1024 / 1024, c64	608.5 → 569.5 (−6.4%)	+0.7%	7.41 → 7.38 (−0.4%)	0.0% / 0.0%

TTFT improves by 2.5% to 6.4% in every workload, above the base-to-base noise of at most 1.1%; throughput gains are 0.6% to 1.9%, and TPOT changes by at most 1.3%.

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. N/A: no new server argument or environment variable.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

Found by Hyperloom.

---

CI States

Latest PR Test (Base): ⏳ Run #41993879604
Latest PR Test (Extra): ❌ Run #41993879563
Latest PR Test (AMD ROCm 10): ✅ Run #41993879438

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

[fengshaoyi-amd]

@HaiShaw @BBuf Could you take a look and add run-ci? Qwen3 prefill on ROCm now runs QK-norm + RoPE as one in-place JIT kernel launch instead of five; CUDA is unchanged. Tested on MI355X and MI325X, numbers in the description. The pr-gate failures are only the missing label. Thanks!

[haishuok0525]

/tag-and-rerun-ci amd-extra

[BBuf]

Please address the [P1] items, then re-ping.

_{Agent-assisted pre-review, findings verified against 039301ba3897.}

[BBuf]

[P1] This also enables Cosmos3's packed-KV path on ROCm: the video path probes pack_kv=True and then consumes the fused prefix/suffix buffers. Neither AMD-registered test here invokes that variant, and its existing tests are CUDA-only. Please add AMD reference coverage for packed K/V, or keep this variant gated off on HIP.

[fengshaoyi-amd]

Kept off on HIP. Your 98cbb77c rejects KV packing there; 86343f05d3 moves that check into _unvalidated_on_hip, which can_use_fused_inplace_qknorm_rope (returns False) and the direct JIT call (raises NotImplementedError) share, and test_kernel_gate_rejects_unvalidated_variants pins it on AMD CI. So on ROCm Cosmos3 video never takes the packed-KV path. After the merge with main, the kDLGPU line this comment is on comes from #47300 and is no longer part of this diff.

[BBuf]

[BBuf Review Agent] Checked this head: packed KV is blocked on HIP in both the capability check and direct JIT calls. This addresses the packed-KV concern.

[BBuf]

[P2] Non-blocking: kUseHalfWarp is hardcoded false under USE_ROCM, so both iterations launch the full-warp kernel. This does not validate the claimed half-warp shuffle behavior; the newly AMD-registered out-of-place variant comparisons have the same issue. Please remove these redundant HIP comparisons or replace them with independent-reference coverage.

[fengshaoyi-amd]

Removed test_half_warp_matches_full_warp, and test_two_rows_per_warp_match_one_row_per_warp in the out-of-place file now skips on HIP. On HIP the shuffles stay covered by test_matches_unfused_path against an fp64 reference, and the out-of-place kernel by test_out_of_place_qknorm_rope_matches_inplace_and_keeps_inputs, bit for bit against the in-place one.

[BBuf]

[BBuf Review Agent] Confirmed the redundant HIP comparisons are removed or skipped, and the enabled path keeps the fp64 reference checks. This resolves the P2.

[BBuf]

@fengshaoyi-amd I pushed 98cbb77c to your branch: Disable unvalidated packed-KV dispatch and direct JIT calls on HIP.
Please pull before your next push. The review runs again on the new head.

_{Automated pre-review for 039301ba3897; it re-checks when the PR is updated.}

[BBuf]

Please address the [P1] items, then re-ping.

_{Agent-assisted pre-review, findings verified against 98cbb77c0d70.}

[BBuf]

[P1] Cosmos3 video with T > 1 still reaches this unrounded NeoX path on ROCm: rejecting pack_kv makes it use the in-place fusion. The new numerical tests use round_norm_before_rope=True, while independent coverage for this variant is CUDA-only. Could we gate this specialization on HIP or add AMD numerical and Cosmos3 validation before enabling it?

[fengshaoyi-amd]

Gated on HIP in 86343f05d3. This shuffle is in the unrounded path, and _unvalidated_on_hip now also rejects NeoX with round_norm_before_rope=False, so can_use_fused_inplace_qknorm_rope returns False and a direct call raises NotImplementedError. On ROCm, Cosmos3 video at T > 1 then goes through _apply_qwen3_qk_norm_rope to the separate QK-norm and RoPE ops, as on main. The only NeoX variant enabled on HIP is the rounded one (the shuffle at line 369), which test_qk_norm_rope_hip.py checks against an fp64 reference.

[BBuf]

[BBuf Review Agent] Confirmed the shared gate rejects unrounded NeoX on HIP, and Cosmos3 video with T > 1 falls back to separate QK-norm and RoPE. This addresses the concern.

[BBuf]

/rerun-test test/registered/kernels/ops/attention/test_qk_norm_rope_hip.py test/registered/kernels/ops/diffusion/test_qknorm_rope_out_of_place.py

Targeted validation for 98cbb77c0d70.

[github-actions]

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/diffusion/test_qknorm_rope_out_of_place.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/diffusion/test_qknorm_rope_out_of_place.py

⛔ test/registered/kernels/ops/attention/test_qk_norm_rope_hip.py: test/registered/kernels/ops/attention/test_qk_norm_rope_hip.py is registered for AMD, not for CUDA or CPU; rerun-test.yml has no AMD job. Rerun it with /rerun-failed-ci, or dispatch the AMD workflow manually.

[fengshaoyi-amd]

@BBuf Pulled 98cbb77c, merged main, and pushed 86343f05d3 for both P1s and the P2; replies are inline. In the qwen3.py conflict with #46908, its CUDA decode branch stays first and the other paths call apply_qk_norm_rope. I also corrected the Scope line in the description: on ROCm the rounded NeoX variant now serves Qwen3 attention and ERNIE-Image, the latter behind its own bit-exact check.

AMD CI on this head: jit-kernel-unit-test-amd runs test_qk_norm_rope_hip.py (128 passed) and test_qknorm_rope_out_of_place.py (1 passed, 12 skipped); stage-a-test-1-gpu-small-amd runs test_qwen3_aiter_fused_qk_norm_rope.py (6 passed). PTAL.

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[fengshaoyi-amd]

Thanks for rerunning. /rerun-failed-ci checks out the original merge commit again (01fe155, this head on top of main 572d3dc from 02:11 UTC), so it cannot pick up the fixes that landed on main later, and the deterministic failures come back:

• test_cfg_parallel_warmup.py and test_fused_lora_compose.py (multimodal-gen unit, CUDA and AMD): fixed on main by [Test] diffusion: align warmup and LoRA fixtures with runtime contracts #47835.
• test_linear_parallel_groups.py (base-b-test-1-gpu-small (4)): broken by [Fix] linear: Keep replicated outputs out of TP symmetric memory #47744, fixed on main by [Test] linear: Expect replicated rows to skip TP symmetric memory #47837.
• test_inkling_checkpoint_indices.py (jit-kernel-unit-test-1-gpu-large (1)): broken by [Refactor] [NPU] graph: Remove the NPU graph runner #47173, fix open in [Fix] spec: Give the mocked draft decode backend a cache_loc_dtype in the inkling test #47860.
• AMD: test_minimax_h3_dit_contract.py needs flash_attn, and Flux3 disagg on 2-GPU is tracked in [CI] [AMD] diffusion/flux3: Make exact-quality VAE convolutions deterministic on ROCm #46947.
• The PR Test Extra and PR Test (MLU) gates fail because this PR is not labelled run-ci-nv-extra / run-ci-mlu; the rest are fail-fast cascades.

The tests for this change pass on both vendors: test_qknorm_rope_out_of_place.py on H100 (13 passed), plus the MI300X jobs linked in the description. I can merge main once #47860 lands so the next run starts from a fixed main.

[BBuf]

/rerun-failed-ci

[siliangchen-amd]

/rerun-failed-ci

[fengshaoyi-amd]

CI triage for bd3d8b0 (this head merged into main dbeecc9). None of these failures come from this PR. Its change only reaches the ROCm callers of the fused QK-norm + RoPE JIT kernel (Qwen3 attention and ERNIE-Image), and none of the failing tests goes through that kernel.

Still failing:

• PR Test Extra and PR Test (MLU) pr-gate, plus pr-test-extra-finish: this PR is not labelled run-ci-nv-extra / run-ci-mlu, so the gates fail by design, as on recently merged PRs.
• AMD multimodal-gen-test-1-gpu-amd (0): test_minimax_h3_dit_contract.py::test_group_attend_skips_only_the_trailing_alignment_rows raises ModuleNotFoundError: No module named 'flash_attn' from the SDPA varlen path (flash_attention_v3.py:239); the AMD image has no flash_attn. The same test fails the same way on another PR's AMD run.
• AMD multimodal-gen-test-2-gpu-amd (0): test_disagg_server.py::TestDisaggFlux3Image::test_t2i_and_edit_match_monolithic, the ROCm VAE determinism issue that [CI] [AMD] diffusion/flux3: Make exact-quality VAE convolutions deterministic on ROCm #46947 fixes.
• AMD stage-c-test-large-8-gpu-amd-mi35x (0): test_kimi_k25_mxfp4.py fails while the server starts. Every TP rank raises miopenStatusInternalError because MIOpen cannot open its SQLite database on the runner (sqlite_db.cpp:224: unable to open database file).
• pr-test-amd-finish will fail on these three.

Being rerun after runner-side failures:

• PR Test Base: github-actions[bot] cancelled the first attempt at 05:22 UTC, together with 58 other PR Test runs, while 16 jobs were still queued or running. pr-test-finish and diffusion-coverage-check failed only because those jobs were missing; no CUDA test failed.
• Arm64 build-test: test/registered/cpu/arm64/test_moe.py fails because sgl_kernel has no fused_experts_cpu. [CI] lint: Reject importing a name through a module that only re-imports it #47992 fails identically on main f6db18e.
• NPU base-b-test-4-npu-a3 / run (0): the runner hit rtsMallocHost ... driver error: out of memory, then its container step failed.

The tests covering this change pass: jit-kernel-unit-test-amd and stage-a-test-1-gpu-small-amd on MI300X on this head, and test_qknorm_rope_out_of_place.py on H100 (13 passed) on the previous head, whose kernel code is identical.

[BBuf]

LGTM. Both earlier findings are resolved.

[BBuf]

/rerun-failed-ci

[fengshaoyi-amd]

@BBuf thanks for the approval. Could you merge this when you have a moment? The required checks (Check PR title, lint) pass, and you are the Codeowner for both CODEOWNERS-protected files in the diff. The reruns keep hitting the same failures, which come from main or from the runners:

• AMD: the 21:11 UTC rerun failed again in multimodal-gen-test-1-gpu-amd (0) (flash_attn is not installed on the AMD image), multimodal-gen-test-2-gpu-amd (0) (Flux3 disagg, fixed by [CI] [AMD] diffusion/flux3: Make exact-quality VAE convolutions deterministic on ROCm #46947) and stage-c-test-large-8-gpu-amd-mi35x (0) (MIOpen cannot open its SQLite database), with the same errors as in the triage above. The Failed to load JIT QK-Norm kernel warnings in those logs come from elementwise/qknorm.cuh, not this PR's kernel, and show up on other PRs' AMD runs too.
• Arm64 and Xeon build-test: the CPU sgl_kernel those jobs load has no fused_experts_cpu / init_cpu_threads_env op. [CI] lint: Reject importing a name through a module that only re-imports it #47992 hits the Arm64 one too, and this PR does not touch CPU kernels.
• PR Test Base attempt 2: multimodal-gen-test-1-gpu (0) missed its latency limits on runner h100-computable-d-gpu-1, including model load time (60.5 s against a 57.1 s limit), and test_request_cancel.py missed its backpressure assertion. This PR leaves CUDA behaviour unchanged. Attempt 3 has mostly been waiting for runners since 21:12 UTC.
• The Extra and MLU gates need labels this PR does not carry.

[BBuf]

/rerun-failed-ci

[BBuf]

[BBuf Review Agent] @fengshaoyi-amd these CI failures look related to this change, could you take a look?

• multimodal-gen-test-2-gpu-amd (rocm10, linux-mi300-2gpu-sglang, 0): E Mismatched elements: 101444 / 177099 (57.3%)" — TestDisaggFlux3Image.test_t2i_and_edit_match_monolithic compares generated images; Flux3's native QKV path calls the changed fused QKNorm/RoPE gate

[fengshaoyi-amd]

@BBuf This failure isn't from this change. It's the ROCm VAE nondeterminism that #46947 fixed, and the reruns here can't pick that fix up.

• It fails the same way without this PR. [CI] [AMD] diffusion/flux3: Make exact-quality VAE convolutions deterministic on ROCm #46947 describes it: text-to-image matches, but the edit differs in ~57% of pixels because MIOpen's fp32 Conv2d in the FLUX.3 VAE is non-deterministic. [Perf] qwen-image21: Pack BF16 gate/up projections on SM120 #47090 doesn't contain this change and fails the same test with 101329 / 177099 (57.2%) (job). Here it's 101444 / 177099 (57.3%).
• The tested commit predates the fix. [CI] [AMD] diffusion/flux3: Make exact-quality VAE convolutions deterministic on ROCm #46947 merged at 06:34 UTC Nov 3, 14 commits after dbeecc9. The AMD run on this head and every /rerun-failed-ci of it test 07d07df (bd3d8b0 merged into dbeecc9), so none of them include it.
• Flux3 calls the gate, but not with any variant this PR changes. On ROCm, Flux3 reaches this kernel through apply_qk_norm_rope (the native path of flux3_qkv_qknorm_rope) and through _fused_qknorm_rope_enabled in flux3.py. Both request is_neox=False and round_norm_before_rope=False with no KV packing, which is the GPT-J variant.
  • The only gate change, _unvalidated_on_hip, rejects KV packing and unrounded NeoX, so it gives Flux3 the same answer as main.
  • The kernel edits are warp shuffles inside if constexpr (kIsNeox) branches and in fused_qknorm_rope_half_warp. HIP never instantiates the half-warp kernel, since kUseHalfWarp is already false under USE_ROCM on main.
  • So the GPT-J kernel Flux3 builds on HIP is the same as main's.

I'll merge main so the next run includes #46947.

[BBuf]

[BBuf Review Agent] @fengshaoyi-amd You're right, Flux3's HIP GPT-J path is unchanged by this diff, so I'm dropping the attribution of this failure to the PR. The current head also includes #46947 after the main merge.

[BBuf]

[BBuf Review Agent] Could you rebase this onto the latest main? It conflicts with main now.

[fengshaoyi-amd]

@BBuf Merged the latest main instead of rebasing, so the reviewed commits aren't rewritten; it's squash-merged either way. The conflict was the sglang.srt.models.utils import in qwen3.py, which #44491 also changed, and the merge keeps both sides' names.

[fengshaoyi-amd]

Lint fails on main itself since #47632: check-indirect-imports flags test/registered/unit/layers/attention/test_kv_translate_ownership.py, which this PR doesn't touch. #48514 and #48516 fix it; I'll merge main once one of them lands.

[BBuf]

[BBuf Review Agent] @fengshaoyi-amd Confirmed: the hook flags the ForwardMode import added by #47632, and that file is unchanged by this PR. This lint failure is unrelated to your changes.

[BBuf]

LGTM.

[peymanr]

PR Pundit review

(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)

#		Finding	Where
1	🟠	Fused ROCm qknorm_rope kernel runs on decode despite prefill-only docstrings, untested for decode/graph-replay	python/sglang/srt/models/qwen3.py:259
2	🟠	On ROCm --no-enable-fused-qk-norm-rope cannot disable fusion: this PR's HIP gate ignores the flag	python/sglang/srt/models/utils.py:671
3	🟠	Model-level ROCm wiring and 2-D-position/dtype fallbacks untested: this PR tests only helper functions	test/registered/kernels/ops/attention/test_qk_norm_rope_hip.py:175

1. 🟠 Should fix: Fused ROCm qknorm_rope kernel runs on decode despite prefill-only docstrings, untested for decode/graph-replay

python/sglang/srt/models/qwen3.py:259

forward_prepare_native has no forward_mode gate, so on ROCm decode also runs the fused JIT qknorm_rope kernel whenever the aiter fused decode path is not taken. That happens without SGLANG_USE_AITER, with an aiter build lacking the kernel, or with a non-NHD/other-dtype KV cache. These calls previously ran the same unfused apply_qk_norm + rotary_emb that the fused kernel is tested against, so this is an equivalent substitution, not a correctness problem. The kernel is built at init, so graph capture does not compile it. The remaining gap: docstrings say "prefill", and decode/graph replay are not covered by a test or benchmark; adjust the wording or add a decode case.

Who is affected: ROCm Qwen3 users whose decode does not take the aiter fused path (SGLANG_USE_AITER unset, aiter missing the kernel, or an unsupported KV-cache layout/dtype).

Checked against the code by an independent pass.

2. 🟠 Should fix: On ROCm --no-enable-fused-qk-norm-rope cannot disable fusion: this PR's HIP gate ignores the flag

python/sglang/srt/models/utils.py:671

can_use_jit_qk_norm_rope turns the fused JIT QK-norm+RoPE on by default on HIP and never reads get_exec().kernel.enable_fused_qk_norm_rope, while the CUDA use_table_qk_norm_rope path honors it. The flag's help says fusion is off by default outside XPU and that --no-enable-fused-qk-norm-rope forces it off; neither is true for this HIP path. Only deterministic inference, rl_on_policy_target or the inductor prefill compiler disable it. The numerical effect is small: the fp32 cos/sin cache is used instead of a bf16-rounded one, a 1-2 ulp difference the PR's test accepts. Consider gating on the flag (treating None as on for HIP) so users can opt out.

Who is affected: Qwen3 (dense) users on ROCm with bf16/fp16 and plain RotaryEmbedding, on prefill and on non-aiter decode.

Checked against the code by an independent pass.

Related past reviews: #20667 on python/sglang/srt/models/utils.py

3. 🟠 Should fix: Model-level ROCm wiring and 2-D-position/dtype fallbacks untested: this PR tests only helper functions

test/registered/kernels/ops/attention/test_qk_norm_rope_hip.py:175

The new HIP test checks can_use_jit_qk_norm_rope and apply_qk_norm_rope in isolation; no ROCm test builds a Qwen3Attention to check that use_jit_qk_norm_rope is set and dispatched, or that graph replay matches eager (test_qwen3_qk_norm_rope.py is CUDA-only and unchanged). Fallbacks for 2-D positions, a q.dtype mismatch with use_jit=True, and rl_on_policy_target are not exercised. The fp32-norm-weight, MRoPE, YaRN, eps-mismatch, deterministic and inductor fallbacks are already covered by test_gate_falls_back and test_gate_rejects_mrope. Optionally assert len(calls)==2 after the int32 re-run.

Who is affected: Maintainers relying on ROCm CI to catch regressions in Qwen3 prefill dispatch to the fused kernel.

Checked against the code by an independent pass.

Related past reviews: #46908 on python/sglang/srt/models/qwen3.py

_{Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at 89cb5645; findings checked and posted by @peymanr; points already raised on this PR are left out.}

[fengshaoyi-amd]

@peymanr Thanks. By number:

1. Right: on ROCm, decode takes the fused kernel whenever the aiter decode fusion doesn't apply. The Modifications section says so, but the test docstring said prefill only; fixed. Graph capture never builds the kernel: the gate at init builds the same module the call uses, and the call has no host sync, so replay only relaunches it. Decode batches are the small token counts the test already covers (1 and 7).

2. Fixed: --no-enable-fused-qk-norm-rope now keeps ROCm Qwen3 on the unfused path. Unset still enables it, and the flag's help says so. test_gate_falls_back covers the opt-out.

3. test_gate_falls_back now also covers on-policy mode and the opt-out, and the int32 re-run asserts len(calls) == 2. The other two fallbacks don't arise in Qwen3: 2-D positions come only with MRoPE, which the gate rejects (test_gate_rejects_mrope), and activations keep the dtype the gate checked at init. Qwen3 calls the gate once at init and passes the result through as use_jit; the profile in the description (Qwen3-8B prefill, 180 → 36 launches, GSM8K on par) exercises that end to end, so I didn't add a model-level unit test.

[BBuf]

[BBuf Review Agent] @fengshaoyi-amd Checked this head: the HIP opt-out is honored and tested, and the docstring now covers decode. The init-time gate and use_jit wiring match your explanation; the existing coverage is enough without another model-level test.

[peymanr]

PR Pundit review

(Testing PR Pundit Code Review agent. Not all comments may be relevant. This is a one-time automated pass: please address what is useful, no need to request a re-review. Please reply to each point, whether it applies or not: we harvest these replies, and both kinds are enormously helpful for tuning the agent. Reach out to @peymanr for suggestions.)

#		Finding	Where
1	🟠	Dumped ROCm server args show enable_fused_qk_norm_rope=None despite Qwen3 fusing: default resolved inline	python/sglang/srt/models/utils.py:672
2	🟠	Hand-built deterministic Qwen3 on ROCm uses fused kernel not bit-identical to unfused: gate checks only env var	python/sglang/srt/models/utils.py:673
3	🟠	Duplicated ROCm shuffle wrappers must be kept in sync with warp.cuh: this PR added per-file #ifdef helpers	python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh:80-98

1. 🟠 Should fix: Dumped ROCm server args show enable_fused_qk_norm_rope=None despite Qwen3 fusing: default resolved inline

python/sglang/srt/models/utils.py:672

On ROCm the --enable-fused-qk-norm-rope default is resolved inline in can_use_jit_qk_norm_rope (enable_fused_qk_norm_rope is not False), unlike XPU, which resolves None->True centrally in handle_xpu_backends (platform_hook.py). The server arg therefore stays None on ROCm, so logged/dumped args do not show the fusion is on. The other readers (qwen3.py:271, qwen3_moe.py:566, mellum.py:232) are gated on _is_cuda/_is_xpu, so no ROCm path reads None as off; the difference is only stylistic and in observability. Resolving the ROCm default in platform_hook.py would keep one convention.

Who is affected: ROCm users running Qwen3 who inspect logged or dumped server args; runtime behaviour is unaffected.

Checked against the code by an independent pass.

Related past reviews: #20667 on python/sglang/srt/models/utils.py, #43345 on python/sglang/srt/arg_groups/fields/exec_.py, #34602 on python/sglang/srt/models/utils.py

2. 🟠 Should fix: Hand-built deterministic Qwen3 on ROCm uses fused kernel not bit-identical to unfused: gate checks only env var

python/sglang/srt/models/utils.py:673

can_use_jit_qk_norm_rope checks only envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE, while the CUDA sibling gate use_table_qk_norm_rope (qwen3.py) also checks get_exec().deterministic.enable_deterministic_inference. Normal server and offline Engine launches are unaffected: resolve_once runs handle_environment_variables, which writes the env var into os.environ before the scheduler processes spawn. The mismatch matters only for code that builds a Qwen3 model directly from a deterministic exec config without resolving ServerArgs. Checking get_exec().deterministic.enable_deterministic_inference as well would make it consistent with the sibling gate.

Who is affected: Only ROCm unit tests or tools that build Qwen3 directly with deterministic inference in the exec config but never run ServerArgs.resolve_once.

Checked against the code by an independent pass.

3. 🟠 Should fix: Duplicated ROCm shuffle wrappers must be kept in sync with warp.cuh: this PR added per-file #ifdef helpers

python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh:80-98

The new shfl_idx32/shfl_xor32 wrappers re-implement shuffles that sgl_kernel/warp.cuh already makes HIP-portable. The half-warp butterfly loops are exactly device::warp::reduce_sum<kRowLanes, 1>(part[i]), and the 32-wide indexed shuffle matches warp::broadcast<device::kWarpThreads>. On ROCm kUseHalfWarp is forced to false, so shfl_xor32's HIP branch is never compiled in. Ignoring the mask on HIP is the same thing warp::reduce already does, so there is no behavioural difference. Prefer the existing warp helpers (or one shared helper in warp.cuh, which would also cover shfl_xor_in_warp in deepseek_v4/main_norm_rope.cuh) over a second per-file shuffle abstraction.

Who is affected: Maintainers of the qknorm_rope JIT kernel; no runtime effect on CUDA or ROCm.

Checked against the code by an independent pass.

Related past reviews: #4302 on sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu, #27947 on python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh

_{Shared from a pilot of PR Pundit, an agentic code reviewer that builds and tests the change and reads the project's past reviews. Reviewed at 22289e7e; findings checked and posted by @peymanr; points already raised on this PR are left out.}

[fengshaoyi-amd]

@peymanr Thanks. By number:

1. Leaving the default inline: on ROCm it depends on the model, since only Qwen3 attention takes the fused path, so resolving None to True in platform_hook.py would report the fusion as on for every other model too.

2. Leaving as is: the gate follows apply_qk_norm's fused path in the same file, which also checks only the env var, and server and Engine launches set that env var from the server arg before the model is built, as you note.

3. Kept local on purpose: on CUDA the wrappers forward the same intrinsics with the same masks, so the CUDA kernel is unchanged from main. Moving the loops onto warp::reduce_sum / warp::broadcast would change the CUDA path as well, which fits better in a follow-up together with [Fix] [AMD] kernels: Repair fused diffusion QK-norm JIT kernels #40585, which repairs the same shuffles.

[BBuf]

[BBuf Review Agent] @fengshaoyi-amd Keeping the model-specific default inline is fine. The help text explains it, and the ROCm gate honors explicit opt-out.

[BBuf]

LGTM.

[BBuf]

/rerun-failed-ci