Repository navigation
[Perf] [NVIDIA] qwen3: Reuse the table QK-norm/RoPE kernel for MRoPE decode - #47993
Conversation
…decode Qwen3-VL decode now takes the same opt-in fused QK-norm + RoPE path as dense text Qwen3. Decode MRoPE positions repeat one position on all three axes, so the layer passes positions[0] and the BF16 rotary table. With a BF16 table the kernel rotates in BF16 with triton_mrope_fused's expression, so the output is bit-identical to the split path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Decode CUDA graphs commonly capture batch sizes of 256–512, and the kernel launches a (T, heads) grid with num_warps=1. That could be slower than the split path at those sizes. Could you add rows at 256 and 512 to check whether this PR outperforms in those scenarios or not? |
| and self.k_norm.weight.dtype == torch.bfloat16 | ||
| and positions.ndim == self.table_rope_positions_ndim | ||
| ): | ||
| if positions.ndim == 2: |
There was a problem hiding this comment.
The fused path relies on _compute_mrope_positions_decode producing repeated positions. Its fallback path (has_precomputed_mrope) produces real spatial positions where the axes differ. If another model built on Qwen3Attention ever supplies precomputed decode positions, and the flag is on, it will produce silently wrong output, with no crash.
There was a problem hiding this comment.
Agreed. A silently wrong fallback isn't acceptable. Fixed in b39e69b and eee8d67 by removing the assumption instead of documenting it.
The kernel now takes the full [3, T] MRoPE positions plus MRotaryEmbedding.axis_map, and gathers each rotary lane's position from its own axis. triton_mrope_fused does the same with its section masks, and axis_map encodes those masks exactly, including the interleaved and GLM-interleaved layouts. The result is correct for any decode positions, including precomputed spatial ones.
Tests now use distinct positions on each axis:
- Kernel: bit-exact against JIT QK-norm +
triton_mrope_fusedfor the section, interleaved, and GLM-interleaved layouts. - Model: the dispatch and CUDA-graph replay tests feed 2D positions with different rows.
…E tokens The table kernel now takes [3, T] MRoPE positions with MRotaryEmbedding.axis_map and gathers each rotary lane's position from its axis, so it no longer assumes decode repeats one position on all axes. Precomputed spatial decode positions stay bit-identical to triton_mrope_fused. On B300 the table kernel loses to the split path above ~3k (token, head) rows (e.g. 0.42x at 512 tokens x 72 heads), for plain RoPE and MRoPE alike, so the fused path is now capped at 3072 // (num_heads + num_kv_heads) tokens. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… runner The decode graph runner passes mrope_positions as a [:, :bs] slice of a [3, max_bs] buffer, so require unit token stride only; the kernel already reads rows through pos_stride. Model tests now use row-strided positions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…drop the token cap The match_cuda table kernel ran one (token, head) per one-warp program, so its cost grew with tokens x heads and it lost to the split path at decode batches of 256-512. Above 1024 (token, head) rows it now handles eight heads per program over four warps, with a [heads, 32] partial-sum tile that keeps each head's reduction order. On B300 it beats the split path at every measured decode size (1.35-3.09x, T=1..512, 9-72 heads, plain RoPE and MRoPE), bit-exact, so Qwen3 no longer caps fused tokens. The rotation arithmetic is flattened to one shared multiply/add branch. The non-match path keeps one head per program and four warps, bit-identical to before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…MRoPE dims and the eight-head path - Graph replay's third MRoPE axis reached position 6153 against a 4096-row table, so both kernels read out of bounds. Positions now wrap at the table length, keeping the row-strided view. - Graph replay adds batch 128 (1280 rows), which runs the eight-heads-per-program path inside a captured graph. - MRoPE kernel tests cover head_dim 64 as well as 128. At head_dim 256 the BF16 rotation is not bit-exact with triton_mrope_fused (about 14% of elements differ by up to 1 ULP), so Qwen3 keeps MRoPE at head_dim 256 on the split path. - col0 keeps its q/k form: the equivalent heads * D changes Triton's codegen and breaks bit-exactness at head_dim 64. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/rerun-test registered/kernels/ops/attention/test_table_qk_norm_rope.py registered/unit/models/test_qwen3_qk_norm_rope.py |
|
🚀 🚀 |
…decode (sgl-project#47993) Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Motivation
#46908 fuses QK RMSNorm and RoPE into one Triton kernel (
table_qk_norm_rope_) for dense text Qwen3 decode, behind--enable-fused-qk-norm-rope. Qwen3-VL uses the sameQwen3Attention, but itsMRotaryEmbeddingkeeps it on the split path: JIT QK-norm, thentriton_mrope_fused. This PR extends the table kernel to MRoPE. It supersedes #43425, which added a separate CUDA kernel for the same case.Modifications
table_qk_norm_rope_[3, T]MRoPE positions (row-strided, as the decode graph runner passes them) plusMRotaryEmbedding.axis_map, and gathers each rotary lane's position from its own axis. Any positions are handled correctly, including precomputed spatial decode positions.triton_mrope_fused's expression, so its output is bit-identical to the split MRoPE path.[heads, 32]partial-sum tile keeps each head's reduction order, so the output stays bit-exact. With one head per program, the kernel's cost grew with tokens x heads and it lost to the split path at decode batches of 256-512; it now wins at every measured size (table below).match_cudapath keeps one head per program and four warps, and its output is bit-identical to before (checked at head_dim 64/128/256).Qwen3AttentionMRotaryEmbeddingwith anmrope_section, at head_dim 64 or 128. At head_dim 256, the table kernel's BF16 rotation isn't bit-exact withtriton_mrope_fused(about 14% of elements differ by up to 1 ULP), so MRoPE at 256 keeps the split path. Plain RoPE still covers 64, 128 and 256._match_cos_sin_cache_dtype, the same cast the split path makes). The cast happens during graph warmup, not after capture.triton_mrope_fused, bit for bit, at head_dim 64 and 128, for the section, interleaved, and GLM-interleaved layouts.The flag stays opt-in on CUDA.
Accuracy
The fused path is bit-identical to the split path, so no accuracy eval was run.
test_table_qk_norm_rope.pyandtest_qwen3_qk_norm_rope.py.Qwen/Qwen3-VL-2B-Instruct, TP1, default backends, decode CUDA graphs on, one B300. 6 concurrent greedy chat requests (2 with an image, 4 text-only), 256 tokens each. All 6 outputs are identical with the flag off and on. A call-recording hook on the flag-on server confirmed fused calls during capture of the bs=512 decode graph, with row-strided[3, 512]positions and a BF16 table.Benchmark
Kernel only, on one idle B300. Each value is the median per-call time in µs over 30 replays of a CUDA graph holding 100 calls. Setup: bf16, head_dim 128.
triton_mrope_fused.apply_rope_with_cos_sin_cache_inplace.Every row is bit-exact with the split path.
The full sweep (T = 1, 8, 32, 64, 128, 256, 512 with heads 8/1, 16/8, 32/8, 64/8) ranges from 1.35x to 3.09x. Below 1024 rows the kernel keeps the one-head layout; interleaved old-vs-new runs at T <= 24 match within noise.
No end-to-end latency was measured.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #41586246656
Latest PR Test (Extra): ❌ Run #41586246680
Latest PR Test (AMD ROCm 10): ❌ Run #41586246809
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: