[AMD][Perf] Fuse QK RMSNorm + gate extraction Triton kernel for Qwen3.5 on HIP

[yichiche]

Motivation

In Qwen3.5's attention layers on HIP, the forward_prepare path for models with attn_output_gate runs three separate kernels to deinterleave q/gate from the QKV projection output and normalize q/k:

1. elementwise_kernel -- copy q from interleaved buffer (4.5 us)
2. elementwise_kernel -- copy gate from interleaved buffer (4.8 us)
3. _fused_qk_gemma_rmsnorm_kernel -- normalize q and k (4.5 us)

This PR fuses all three operations into a single Triton kernel (_fused_qk_gemma_rmsnorm_gate_kernel) that reads q and gate from the interleaved [q_h0, gate_h0, q_h1, gate_h1, ...] layout, normalizes q/k with GemmaRMSNorm, and writes contiguous q, k, and gate outputs -- eliminating 2 kernel launches per attention layer.

Modifications

• python/sglang/srt/models/utils.py: Add _fused_qk_gemma_rmsnorm_gate_kernel Triton kernel and fused_qk_gemma_rmsnorm_with_gate() Python wrapper. The kernel handles the interleaved q_gate layout with stride-based access (no intermediate copies), performs GemmaRMSNorm on q and k, and extracts gate to a contiguous output.
• python/sglang/srt/models/qwen3_5.py: Add forward_prepare_hip() method to Qwen3_5AttentionDecoderLayer that uses the fused kernel. Dispatch to it when _is_hip and self.attn_output_gate in the forward_attention() method.

Accuracy Tests

Model: Qwen3.5-397B-A17B-MXFP4, TP=2, MI355x

Benchmark	Score	Threshold
GSM8K (2000 questions, parallel=2000)	0.911	0.880

Benchmarking and Profiling

Kernel-level (per attention layer, decode, TP=2, MI355x)

Kernel	Before (us)	After (us)	Notes
elementwise_kernel (q deinterleave)	4.5	--	eliminated
elementwise_kernel (gate deinterleave)	4.8	--	eliminated
_fused_qk_gemma_rmsnorm_kernel	4.5	--	eliminated
_fused_qk_gemma_rmsnorm_gate_kernel	--	4.6	new fused kernel
Attention layer total	91.0	81.6	-9.4 us (-10.3%)

15 attention layers x 9.4 us = ~141 us savings per decode iteration.

E2E benchmark (Qwen3.5-397B-A17B-MXFP4, TP=2, ISL=8192, OSL=1024, conc=4, MI355x)

Metric	Before	After	Delta
Total throughput (tok/s)	2920.1	2951.4	+1.1%
Output throughput (tok/s)	324.8	328.3	+1.1%
Median TTFT (ms)	393.4	378.5	-3.8%
Median ITL (ms)	10.59	10.47	-1.1%
Median TPOT (ms)	11.66	11.53	-1.1%

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.

---

CI States

Latest PR Test (Base): 🚫 Run #27216146908
Latest PR Test (Extra): ❌ Run #27216145634

[gemini-code-assist]

Code Review

This pull request updates the documentation to feature a new blog post on MoRI and AMD Instinct MI355X, and introduces a fused Triton kernel (fused_qk_gemma_rmsnorm_with_gate) to optimize the attention forward pass for Qwen 3.5 on AMD GPUs (HIP). The reviewer provided valuable performance optimization feedback to avoid an expensive memory copy of the non-contiguous k tensor by passing it as a 3D tensor and utilizing stride-based indexing directly inside the Triton kernel.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[yichiche]

@amd-bot ci-status

[amd-bot]

@yichiche

CI Status for PR #27656

Merge verdict: No failures are caused by this PR — all 6 leaf-job failures are in unrelated code paths (aiter all-reduce, NPU diffusion/quant perf, a CPU missing-module, and a HuggingFace 429 rate-limit). The PR's new Triton kernel helper is well-tested and green on AMD (113 unit cases passed). ⚠️ However, the model-level HIP wiring (forward_prepare_hip in qwen3_5.py) is not exercised by any PR-CI test, so green does not verify the end-to-end Qwen3.5 forward pass with this fusion.

Warning

Coverage is partial. The kernel math in models/utils.py (fused_qk_gemma_rmsnorm_with_gate) is fully verified — the new jit-kernel-unit-test-amd suite ran 113 parametrized cases against a PyTorch reference and all passed (job). But the production integration in qwen3_5.py (forward_prepare_hip, gated by _is_hip and self.attn_output_gate) is reachable only by running a Qwen3.5 model on AMD — and the three registered Qwen3.5 tests (test_lora_qwen3_5_4b_logprob_diff, test_lora_qwen3_5_35b_a3b_logprob_diff, test_qwen3_5_packed_weight_loader) register CUDA-only / CPU-only, never AMD. Before merge, run a Qwen3.5 (with attn_output_gate) accuracy/logprob check on an AMD MI325 to confirm the fused path matches the native path end-to-end.

Changed files: jit_kernel/tests/test_fused_qk_gemma_rmsnorm_gate.py (+151), srt/models/qwen3_5.py (+37/-3), srt/models/utils.py (+127/-4)

AMD: 2 failures (0 related) · Others: 4 failures (0 related) — all 169 checks completed, none pending.

AMD CI Failures

Job	Test File	Test Function	Error	Related?	Why
stage-c-test-large-8-gpu-amd (3)	test/registered/ops/test_aiter_allreduce_fusion_amd.py	test_fused_ar_rms_residual_accuracy	AssertionError: Residual accuracy check failed (max_diff=6.25e-2 for all shapes)	🟢	aiter custom all-reduce + RMS fusion path; PR touches QK-norm only, no all-reduce code. Uniform failure across every shape ⇒ aiter/env issue, not this diff.
stage-b-test-1-gpu-large-amd (0)	test/registered/spec/eagle/test_eagle_constrained_decoding.py	setUpClass	HfHubHTTPError: 429 Too Many Requests (meta-llama/Llama-2-7b-chat-hf)	🟢	HuggingFace download rate-limit during model fetch — infra/network, unrelated to code.

Other CI Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b-test-1-npu-a2 (0)	test/registered/ascend/basic_function/quant/test_npu_w8a8_quantization.py	(throughput check)	AssertionError: 487.2 not >= 700 (throughput)	🟢	NPU perf-threshold flake on Qwen2.5-0.5B w8a8; PR is HIP-only, doesn't touch NPU or quant.
multimodal-gen-test-1-npu-a3	multimodal_gen/.../ascend/test_server_1_npu.py	test_diffusion_generation[wan2_1_t2v_1.3b_1_npu]	Failed: Diffusion testcase failed 1 check(s)	🟢	NPU diffusion (Wan2.1) generation check — unrelated subsystem.
multimodal-gen-test-2-npu-a3	multimodal_gen/.../ascend/test_server_2_npu.py	test_diffusion_generation[flux_2_image_t2i_2npu], [wan2_2_t2v_14b_w8a8_2npu]	Failed: Diffusion testcase failed 1 check(s)	🟢	NPU diffusion (Flux/Wan2.2) generation checks — unrelated subsystem.
build-test (xeon-gnr, base-b-test-cpu)	test/registered/unit/bench/test_mixed_prefix_gsm8k.py	(import)	ModuleNotFoundError: No module named 'sglang.test.simple_eval_mixed_prefix_gsm8k'	🟢	Missing module on CPU build — pre-existing infra/packaging issue, not in this diff.

The 4 *-finish / call-gate / pr-gate failures are rollup/cascade jobs reflecting the leaf failures above, not independent failures.

Details / what to do before merge

• None of the 6 failures block this PR — they are infra (HF 429, CPU missing module), or unrelated subsystems (aiter all-reduce, NPU diffusion/quant perf). The maintainer may re-run the AMD stage-b eagle job to clear the transient 429.
• Close the coverage gap (the one real action item): the kernel helper is verified, but the forward_prepare_hip model path has zero AMD coverage. Either add register_amd_ci(...) to a Qwen3.5 logprob/accuracy test (matching the existing CUDA registration), or manually run a Qwen3.5-with-attn_output_gate accuracy check on AMD MI325 and confirm output parity with the native (forward_prepare_native) path before merging.

---

Generated by amd-bot using Claude Code CLI

[mmangkad]

This broke lint on main