[AMD] Optimize Qwen3-VL decode - fuse QK-norm + 3D mRoPE + KV cache write

[yctseng0211]

Motivation

Use aiter's fused_qk_norm_mrope_3d_cache_pts_quant_shuffle kernel to replace 4 separate kernel launches (QKV split, QK RMSNorm, 3D mRoPE, KV cache write) with a single HIP kernel on the ROCm decode path.

Guarded by SGLANG_USE_AITER + is_hip() + MRotaryEmbedding detection;

Modifications

Accuracy Tests

python3  sglang/benchmark/gsm8k/bench_sglang.py \
                --port 30000 \
                --parallel 32 \
                --num-questions 200 \
                --num-shots 5

Benchmarking and Profiling

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 Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[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!

[kkHuang-amd]

I don't suggest to copy whole attention block processing logic into one function.

It will not follow the sglang attention processing logic.

forward_prepare -> forward_core.

It will not be easily to maintain

[yctseng0211]

I don't suggest to copy whole attention block processing logic into one function.

It will not follow the sglang attention processing logic.

forward_prepare -> forward_core.

It will not be easily to maintain

@kkHuang-amd
I've refactored to follow the forward_prepare -> forward_core pattern. The fused kernel now lives in forward_prepare_fused_mrope() which returns (q, None, None), and forward() uses the standard prepare -> attn -> o_proj flow with save_kv_cache=False.

A separate forward_prepare variant is needed because the aiter fused kernel combines both the prepare stage (split + QK-norm + mRoPE) and part of the core stage (KV cache write) into a single kernel launch, so it can't be cleanly inserted into forward_prepare_native which has no knowledge of the KV cache.

[kkHuang-amd]

LGTM

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #21458

1 possibly related · 4 unrelated

Workflow	Job	Error	Related?	Log
PR Test (AMD)	stage-c-test-large-8-gpu-amd-mi35x (linux-mi35x-gpu-8, 0)	RuntimeError: invalid argument for batch_prefill	🟡 Possibly -- This is an AMD/ROCm runner where the PR's new aiter fused kernel import and MRotaryEmbedding logic would be exercised if SGLANG_USE_AITER is set, so the failure could be caused by the PR's changes to the Qwen3 attention forward path.	View
PR Test	stage-c-test-4-gpu-h100 (0)	Fast-fail: skipping â�� root cause job(s): stage-c-test-4-gpu-b200 (1)	🟢 Unlikely -- The failure is in a health-check infrastructure step, not in any test execution, and the PR only modifies Qwen3 model code with ROCm/aiter-specific paths that wouldn't affect H100 GPU machines.	View
PR Test	stage-c-test-8-gpu-h200 (0)	AssertionError: Tests failed:	🟢 Unlikely -- The PR changes are specific to ROCm/HIP with aiter fused kernels and minor control-flow changes in Qwen3 attention; H200 tests failing likely reflect a pre-existing or infrastructure issue unrelated to these changes.	View
PR Test	stage-c-test-4-gpu-b200 (1)	Cleaning up orphan processes	🟢 Unlikely -- B200 is an NVIDIA GPU and the PR's new code paths are gated behind _is_hip and SGLANG_USE_AITER, so the fused path cannot activate; the failure is likely infrastructure or a pre-existing flaky test.	View
PR Test	stage-c-test-4-gpu-b200 (0)	The action 'Run test' has timed out after 30 minutes.	🟢 Unlikely -- Same reasoning as B200 (1): the PR's functional changes are ROCm-only and the minor refactor (save_kv_cache flag) preserves existing behavior on NVIDIA GPUs, so this is likely an unrelated failure.	View

---

Generated by amd-bot

[yctseng0211]

The guard in Line:274
is needed because there's a downstream k.to(torch.bfloat16) cast in the RL on-policy path, without the guard, the fused prepare would return k=None and that .to() call would crash.

[bingxche]

@amd-bot ci-status

[amd-bot]

@bingxche

Now I have all the information needed. Let me compile the analysis.

CI Status for PR #21458

PR: [AMD] Optimize Qwen3-VL decode - fuse QK-norm + 3D mRoPE + KV cache write
Changed file: python/sglang/srt/models/qwen3.py (101 additions, 3 deletions)
Scope: Adds a fused HIP kernel path (aiter fused_qk_norm_mrope_3d_cache_pts_quant_shuffle) for Qwen3 decode on ROCm, guarded by SGLANG_USE_AITER + is_hip() + MRotaryEmbedding.

Job	Error	Related to PR?	Explanation
stage-c-test-4-gpu-b200 (1)	test_nvidia_nemotron_3_super_nvfp4.py: Server exited with code 1 (DeepSeek V3.2 model type warning)	Unlikely	NVFP4 Nemotron test on B200 GPU — completely unrelated model/codepath. PR only touches Qwen3 attention with ROCm guards.
stage-c-test-4-gpu-b200 (0)	test_gpt_oss_4gpu.py: accuracy 0.586 < 0.6 threshold on test_mxfp4_120b, then 30-min timeout	Unlikely	GPT OSS MXFP4 accuracy test — different model, different code path. The accuracy miss + timeout is not in Qwen3 code.
stage-c-test-8-gpu-h200 (0)	test_ring_2_5_1t.py: Server failed to start within timeout (missing model shards for inclusionAI/Ring-2.5-1T)	Unlikely	Ring-2.5-1T model download/cache issue — missing shard 160 of 160. Infrastructure/cache problem, no relation to Qwen3 changes.
stage-c-test-4-gpu-h100 (0)	Fast-fail cascade from stage-c-test-4-gpu-b200 (1)	Unlikely	Not a real failure — CI health-check cascade triggered by the B200 job failing first.
stage-c-test-large-8-gpu-amd-mi35x (0)	test_qwen3_coder_next_8gpu.py: mha_batch_prefill exception in aiter, server killed (code -9)	Possibly	This test runs a Qwen3 variant on MI35X with aiter. The crash is in aiter/csrc/py_itfs_ck/mha_batch_prefill_kernels.cu, not in the fused kernel this PR adds. However, the PR changes the Qwen3Attention forward path on ROCm and the test model is Qwen3-based. See details below.
pr-test-finish	Reports stage-c-test-4-gpu-h100: failure	Unlikely	Gate job — just aggregates results from the above failures.
pr-test-amd-finish	Reports stage-c-test-large-8-gpu-amd-mi35x: failure	Unlikely	Gate job — aggregates AMD test results.

Details

stage-c-test-large-8-gpu-amd-mi35x (0) — Possibly related

The failing test test_qwen3_coder_next_8gpu.py runs Qwen/Qwen3-Coder-Next on MI35X with --attention-backend aiter --kv-cache-dtype fp8_e4m3. The crash trace shows:

Exception raised from mha_batch_prefill at aiter/csrc/py_itfs_ck/mha_batch_prefill_kernels.cu:820

This is in the attention prefill kernel, not the fused QK-norm+mRoPE kernel the PR introduces. The PR's fused path is only activated during decode (forward_batch.forward_mode.is_decode()), and requires both SGLANG_USE_AITER=1 and a model with MRotaryEmbedding. Qwen3-Coder-Next may or may not use MRotaryEmbedding (standard Qwen3 uses regular RoPE, not mRoPE). Even if the fused path were active, the crash is in mha_batch_prefill which runs before decode.

The crash is more likely an aiter/MI35X infrastructure issue or model compatibility problem, but given that the PR modifies Qwen3Attention.forward() control flow (adding a new branch and changing save_kv_cache plumbing), a side effect cannot be 100% ruled out.

All NVIDIA failures — Unlikely related

The PR code is entirely guarded by _is_hip and SGLANG_USE_AITER. On NVIDIA GPUs, is_hip() returns False, so _has_fused_qk_norm_mrope is always False and use_fused is always False. The NVIDIA code path is unchanged — the only modification visible to NVIDIA is the trivially-added save_kv_cache=True default that preserves existing behavior.

All other AMD jobs — Passed

All MI325X AMD jobs (1-GPU, 2-GPU, 4-GPU, 8-GPU) passed successfully, including stage-c-test-4-gpu-amd and stage-c-test-large-8-gpu-amd. Only the MI35X partition (0) failed.

Generated by amd-bot using Claude Code CLI