[zcnrex]

Motivation

Qwen3 attention runs Q/K RMSNorm and RoPE as two separate kernels. This PR fuses them into one JIT kernel on datacenter Blackwell (SM100) for BF16 Q/K with 128-element heads, full-width NeoX rotation, and 1–128 tokens. That covers:

• Dense text Qwen3 (plain RotaryEmbedding), in decode and in short extends.
• Dense Qwen3-VL (MRotaryEmbedding), in decode only.

Every other configuration keeps the existing split path.

Modifications

Kernel. FusedQKNormRoPEWithCache<kUsePDL, CacheT> in elementwise/fused_qknorm_rope.cuh reuses the warp kernel from qknorm.cuh (load, reduction, store) through an epilogue hook, and applies RoPE from the model's cos_sin_cache before the store. The rotation arithmetic follows the cache dtype, so each path matches its split counterpart:

• FP32 cache (plain RoPE): the same FP32 expressions as rope.cuh, which apply_rope_with_cos_sin_cache_inplace uses.
• BF16 cache (MRoPE): the BF16 FMA order of triton_mrope_fused.

It accepts strided Q/K (views into packed QKV) and 1D int64 positions, and schedules one warp per head. Ordinary QK normalization keeps its persistent scheduling, and the head_dim-128 half-warp kernel from #45214 is unaffected.

API. fused_qk_norm_rope_with_cache(q, k, q_weight, k_weight, cos_sin_cache, positions, eps) in sglang.kernels.ops.attention.fused_qknorm_rope, registered as attention.fused_qk_norm_rope_with_cache.

Qwen3 dispatch (models/qwen3.py):

• At init, the fused path requires: CUDA SM100, head_dim == rotary_dim == 128, NeoX style, a rotary embedding of exactly RotaryEmbedding or MRotaryEmbedding with mrope_section, no forced-native rotary, no deterministic inference, and no Inductor prefill.
• Per call, it requires: BF16 Q and norm weights, 1–128 tokens, and 1D int64 positions.
• MRoPE uses the fused path in decode only. Decode MRoPE positions repeat one position on all three axes (ForwardBatch._compute_mrope_positions_decode), so positions[0] is exact. MRoPE extend and speculative steps keep the Triton path.
• forward_prepare_native now takes forward_batch so it can read the forward mode.

Out of scope: KV-cache write fusion, QKV projection fusion, and MoE Qwen3 (qwen3_moe.py).

Merged with upstream main at cb55d77d37.

Validation

On commit 9e62fc7a69, one NVIDIA B300 (SM 10.3). Another job was running on that GPU, so these are correctness results only.

• 24/24 cases match the split path bit for bit. The new test is test/registered/kernels/ops/attention/test_fused_qk_norm_rope_with_cache.py, and it compares against:

  • FP32 cache: JIT qknorm, then apply_rope_with_cos_sin_cache_inplace.
  • BF16 cache: JIT qknorm, then triton_mrope_fused with three identical position rows.

  Coverage: 1, 17 and 128 tokens; Q/K heads 8/2 and 32/8; packed QKV with and without 64 elements of row padding. V and padding are byte-identical to the split path.

• Sanity check: the kernel changes about 99.9% of Q/K values, leaves V untouched, and its output depends on position.

The earlier results on 70b63a4575 (the MRoPE-only version on GB300) predate both this generalization and the merge of #45214, and have not been rerun. They covered 168 synthetic MRoPE cases, 27 CUDA-graph replays, 60 base-vs-PR QK-norm comparisons, and the 560 existing test_qknorm.py cases.

Kernel Performance (stale, rerun pending)

Stale. These numbers were measured on 70b63a4575, before this PR was extended to plain RoPE and before #45214 made the split path's QK normalization faster at head_dim 128. Expect the split-path baseline to be lower now. The plain-RoPE path has not been measured. A rerun against current main is pending.

Kernel-only latency in microseconds, median of three CUDA-graph runs (MRoPE, BF16, head_dim 128, GB300):

Tokens	Q/K heads	Split QK norm + MRoPE	This PR	Existing Triton fusion	Split / PR
1	8/2	2.751	1.304	1.700	2.11x
1	32/8	3.452	1.311	1.778	2.63x
16	8/2	3.257	1.372	1.866	2.37x
16	32/8	4.231	1.469	2.090	2.88x
128	8/2	3.524	1.648	2.634	2.14x
128	32/8	5.121	3.082	6.577	1.66x

Validation Limits

These are kernel-level correctness results only. No end-to-end accuracy, serving-throughput, or decode-latency measurements have been made, and no serving improvement is claimed yet.

Checklist

• Repository pre-commit checks pass.
• Automated regression coverage for the fused path (plain RoPE and decode MRoPE).
• GPU correctness on the exact PR commit.
• Rerun kernel performance against current main, including plain RoPE.
• End-to-end decode latency and accuracy for Qwen3 (text) and Qwen3-VL.
• Upstream CI and review.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #40563264384
Latest PR Test (Extra): ❌ Run #40563263945
Latest PR Test (AMD ROCm 10): ❌ Run #40563264343

[zcnrex]

Superseded by #47993. It reuses #46908's Triton table kernel for Qwen3-VL decode MRoPE, and the output stays bit-identical to the split path. On B300 the separate CUDA kernel here was only 0.2-0.35 us/layer faster than the table route at decode batch <= 32. That's too small a gain to justify maintaining another CUDA kernel, and #46908 already covers the plain-RoPE half.