[Perf] [NVIDIA] qwen3: Reuse the table QK-norm/RoPE kernel for MRoPE decode

python/sglang/kernels/ops/attention/table_qk_norm_rope.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import triton
 import triton.language as tl
@@ -12,55 +14,77 @@ def _table_qk_norm_rope_kernel(
     k_weight_ptr,
     cos_sin_ptr,
     pos_ptr,
+    axis_map_ptr,
+    pos_stride,
     row_stride,
     q_size,
     NHQ: tl.constexpr,
+    NH: tl.constexpr,
     D: tl.constexpr,
     EPS: tl.constexpr,
     MATCH_CUDA: tl.constexpr,
+    BF16_TABLE: tl.constexpr,
+    MROPE: tl.constexpr,
+    HEADS_PER_PROGRAM: tl.constexpr,
 ):
     t = tl.program_id(0).to(tl.int64)
-    h = tl.program_id(1)
-    pos = tl.load(pos_ptr + t).to(tl.int64)
+    heads = tl.program_id(1) * HEADS_PER_PROGRAM + tl.arange(0, HEADS_PER_PROGRAM)
 
     HALF: tl.constexpr = D // 2
-    half_ar = tl.arange(0, HALF)
-    d_ar = tl.arange(0, D)
-    cos = tl.load(cos_sin_ptr + pos * D + half_ar).to(tl.float32)
-    sin = tl.load(cos_sin_ptr + pos * D + HALF + half_ar).to(tl.float32)
+    half_ar = tl.arange(0, HALF)[None, :]
+    d_ar = tl.arange(0, D)[None, :]
+    if MROPE:
+        axis = tl.load(axis_map_ptr + half_ar)
+        pos = tl.load(pos_ptr + axis * pos_stride + t).to(tl.int64)
+    else:
+        pos = tl.load(pos_ptr + t).to(tl.int64)
+    cos = tl.load(cos_sin_ptr + pos * D + half_ar)
+    sin = tl.load(cos_sin_ptr + pos * D + HALF + half_ar)
 
-    is_q = h < NHQ
-    col0 = tl.where(is_q, h * D, q_size + (h - NHQ) * D).to(tl.int64)
-    w_ptr = tl.where(is_q, q_weight_ptr.to(tl.int64), k_weight_ptr.to(tl.int64)).to(
-        tl.pointer_type(tl.bfloat16)
-    )
+    is_q = (heads < NHQ)[:, None]
+    # Equal to heads * D, but that form changes Triton's codegen and breaks bit-exactness.
+    col0 = tl.where(heads < NHQ, heads * D, q_size + (heads - NHQ) * D).to(tl.int64)
+    row = (qkv_ptr + t * row_stride + col0)[:, None]
+    mask = (heads < NH)[:, None]
 
-    row = qkv_ptr + t * row_stride + col0
     if MATCH_CUDA:
-        lane = tl.arange(0, 32)
-        part = tl.full((32,), 0, tl.float32)
+        lane = tl.arange(0, 32)[None, :]
+        part = tl.full((HEADS_PER_PROGRAM, 32), 0, tl.float32)
         for i in tl.static_range(D // 32):
-            x = tl.load(row + lane * (D // 32) + i).to(tl.float32)
+            x = tl.load(row + lane * (D // 32) + i, mask=mask, other=0)
+            x = x.to(tl.float32)
             part = tl.fma(x, x, part)
-        inv = cuda_rsqrtf(tl.fma(tl.sum(part, 0), 1.0 / D, EPS))
+        inv = cuda_rsqrtf(tl.fma(tl.sum(part, 1), 1.0 / D, EPS))[:, None]
     else:
-        x = tl.load(row + d_ar).to(tl.float32)
-        ms = tl.sum(x * x, 0) / D
-        inv = 1.0 / tl.sqrt(ms + EPS)
-    w1 = tl.load(w_ptr + half_ar).to(tl.float32)
-    w2 = tl.load(w_ptr + HALF + half_ar).to(tl.float32)
-    x1 = tl.load(row + half_ar).to(tl.float32) * inv * w1
-    x2 = tl.load(row + HALF + half_ar).to(tl.float32) * inv * w2
-    x1 = x1.to(tl.bfloat16).to(tl.float32)
-    x2 = x2.to(tl.bfloat16).to(tl.float32)
-    if MATCH_CUDA:
+        x = tl.load(row + d_ar, mask=mask, other=0).to(tl.float32)
+        ms = tl.sum(x * x, 1) / D
+        inv = (1.0 / tl.sqrt(ms + EPS))[:, None]
+    w1 = tl.where(
+        is_q, tl.load(q_weight_ptr + half_ar), tl.load(k_weight_ptr + half_ar)
+    ).to(tl.float32)
+    w2 = tl.where(
+        is_q,
+        tl.load(q_weight_ptr + HALF + half_ar),
+        tl.load(k_weight_ptr + HALF + half_ar),
+    ).to(tl.float32)
+    x1 = tl.load(row + half_ar, mask=mask, other=0).to(tl.float32) * inv * w1
+    x2 = tl.load(row + HALF + half_ar, mask=mask, other=0).to(tl.float32) * inv * w2
+    x1 = x1.to(tl.bfloat16)
+    x2 = x2.to(tl.bfloat16)
+    if not BF16_TABLE:
+        x1 = x1.to(tl.float32)
+        x2 = x2.to(tl.float32)
+        cos = cos.to(tl.float32)
+        sin = sin.to(tl.float32)
+    if MATCH_CUDA and not BF16_TABLE:
         o1 = tl.fma(x1, cos, -mul_rn_f32(x2, sin))
         o2 = tl.fma(x2, cos, mul_rn_f32(x1, sin))
     else:
+        # A BF16 table uses triton_mrope_fused's expression, so its codegen matches.
         o1 = x1 * cos - x2 * sin
         o2 = x2 * cos + x1 * sin
-    tl.store(row + half_ar, o1.to(tl.bfloat16))
-    tl.store(row + HALF + half_ar, o2.to(tl.bfloat16))
+    tl.store(row + half_ar, o1.to(tl.bfloat16), mask=mask)
+    tl.store(row + HALF + half_ar, o2.to(tl.bfloat16), mask=mask)
 
 
 def table_qk_norm_rope_(
@@ -75,6 +99,7 @@ def table_qk_norm_rope_(
     eps: float,
     *,
     match_cuda: bool = False,
+    axis_map: Optional[torch.Tensor] = None,
 ) -> None:
     """In-place QK RMSNorm + table-lookup neox RoPE on the fused QKV tensor.
 
@@ -85,8 +110,12 @@ def table_qk_norm_rope_(
     ``match_cuda`` preserves the reduction and rounding of the separate CUDA
     QK-norm and RoPE kernels for head dimensions 64/128/256. It requires BF16
     QKV with contiguous columns, contiguous BF16 weights, a contiguous FP32
-    rotary table, and contiguous int32/int64 positions on the same CUDA device.
-    Row-strided QKV is supported.
+    or BF16 rotary table, and int32/int64 positions with unit token stride on
+    the same CUDA device. Row-strided QKV and MRoPE positions are supported.
+
+    MRoPE passes [3, T] positions and ``MRotaryEmbedding.axis_map``, the axis
+    each rotary lane reads its position from. A BF16 table rotates in BF16 to
+    match ``triton_mrope_fused``.
     """
     if match_cuda:
         assert head_dim in (64, 128, 256)
@@ -96,28 +125,50 @@ def table_qk_norm_rope_(
         assert q_weight.shape == k_weight.shape == (head_dim,)
         assert q_weight.is_contiguous() and k_weight.is_contiguous()
         assert cos_sin_cache.ndim == 2 and cos_sin_cache.shape[1] == head_dim
-        assert cos_sin_cache.dtype == torch.float32 and cos_sin_cache.is_contiguous()
-        assert positions.shape == (qkv.shape[0],) and positions.is_contiguous()
+        assert cos_sin_cache.dtype in (torch.float32, torch.bfloat16)
+        assert cos_sin_cache.is_contiguous()
+        num_axes = () if axis_map is None else (3,)
+        assert positions.shape == (*num_axes, qkv.shape[0])
+        assert positions.stride(-1) == 1
         assert positions.dtype in (torch.int32, torch.int64)
+        if axis_map is not None:
+            assert axis_map.shape == (head_dim // 2,) and axis_map.is_contiguous()
+            assert axis_map.dtype == torch.int64
         assert qkv.is_cuda and all(
             t.device == qkv.device
-            for t in (positions, q_weight, k_weight, cos_sin_cache)
+            for t in (positions, q_weight, k_weight, cos_sin_cache, axis_map)
+            if t is not None
         )
     T = qkv.shape[0]
     if T == 0:
         return
-    grid = (T, num_q_heads + num_k_heads)
+    num_heads = num_q_heads + num_k_heads
+    # Measured on B300: one head per program wins up to ~1k (token, head) rows,
+    # eight heads over four warps above that.
+    if not match_cuda:
+        heads_per_program, num_warps = 1, 4
+    elif T * num_heads <= 1024:
+        heads_per_program, num_warps = 1, 1
+    else:
+        heads_per_program, num_warps = 8, 4
+    grid = (T, triton.cdiv(num_heads, heads_per_program))
     _table_qk_norm_rope_kernel[grid](
         qkv,
         q_weight,
         k_weight,
         cos_sin_cache,
         positions,
+        axis_map,
+        positions.stride(0),
         qkv.stride(0),
         num_q_heads * head_dim,
         NHQ=num_q_heads,
+        NH=num_heads,
         D=head_dim,
         EPS=eps,
         MATCH_CUDA=match_cuda,
-        num_warps=1 if match_cuda else 4,
+        BF16_TABLE=cos_sin_cache.dtype == torch.bfloat16,
+        MROPE=axis_map is not None,
+        HEADS_PER_PROGRAM=heads_per_program,
+        num_warps=num_warps,
     )

python/sglang/srt/models/qwen3.py

@@ -262,9 +262,20 @@ def __init__(
             and not envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
             and get_exec().deterministic.rl_on_policy_target is None
             and get_exec().graph.cuda_graph_config.prefill.tc_compiler != "inductor"
-            and type(self.rotary_emb) is RotaryEmbedding
+            and (
+                type(self.rotary_emb) is RotaryEmbedding
+                or (
+                    # At head_dim 256, triton_mrope_fused's BF16 rounding differs.
+                    type(self.rotary_emb) is MRotaryEmbedding
+                    and self.rotary_emb.mrope_section is not None
+                    and self.head_dim in (64, 128)
+                )
+            )
             and self.head_dim in (64, 128, 256)
         )
+        self.table_rope_positions_ndim = (
+            2 if type(self.rotary_emb) is MRotaryEmbedding else 1
+        )
 
         self.use_fused_qk_norm_mrope = (
             _has_fused_qk_norm_mrope
@@ -309,7 +320,12 @@ def forward_prepare_native(self, positions, hidden_states, forward_batch):
             and qkv.dtype == torch.bfloat16
             and self.q_norm.weight.dtype == torch.bfloat16
             and self.k_norm.weight.dtype == torch.bfloat16
+            and positions.ndim == self.table_rope_positions_ndim
         ):
+            axis_map = None
+            if positions.ndim == 2:
+                axis_map = self.rotary_emb.axis_map
+                self.rotary_emb._match_cos_sin_cache_dtype(qkv)
             table_qk_norm_rope_(
                 qkv,
                 positions,
@@ -321,6 +337,7 @@ def forward_prepare_native(self, positions, hidden_states, forward_batch):
                 self.head_dim,
                 self.q_norm.variance_epsilon,
                 match_cuda=True,
+                axis_map=axis_map,
             )
         else:
             q, k = apply_qk_norm(

test/registered/kernels/ops/attention/test_table_qk_norm_rope.py

@@ -1,11 +1,14 @@
 import sys
+from types import SimpleNamespace
 
 import pytest
 import torch
 
 from sglang.kernels.ops.attention.rope import apply_rope_inplace
+from sglang.kernels.ops.attention.rotary_triton import triton_mrope_fused
 from sglang.kernels.ops.attention.table_qk_norm_rope import table_qk_norm_rope_
 from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -52,6 +55,7 @@ def inputs(d, batch, hq, hk, seed, padded=False, position_dtype=torch.int64):
         (1, 32, 8, False, torch.int64),
         (4, 8, 2, True, torch.int32),
         (32, 32, 8, False, torch.int64),
+        (129, 8, 2, True, torch.int64),
         (257, 7, 1, True, torch.int64),
     ],
 )
@@ -71,6 +75,59 @@ def test_cuda_sequence_is_bit_exact(d, batch, hq, hk, padded, pos_dtype):
         assert_same_bits(storage[:, -8:], before[:, -8:])
 
 
+@pytest.mark.parametrize("d,section", [(64, [16, 8, 8]), (128, [24, 20, 20])])
+@pytest.mark.parametrize("layout", ["sections", "interleaved", "interleaved_glm"])
+@pytest.mark.parametrize(
+    "batch,hq,hk", [(1, 32, 8), (32, 8, 1), (129, 12, 2), (257, 64, 8)]
+)
+def test_bf16_mrope_matches_triton_mrope(d, section, batch, hq, hk, layout):
+    """A BF16 table with per-axis positions must reproduce the split MRoPE path."""
+    _, x, pos, qw, kw, cache = inputs(d, batch, hq, hk, 29)
+    cache = cache.bfloat16()
+    # Distinct axes, as in prefill or precomputed spatial decode positions.
+    positions = torch.stack([pos, (pos + 7) % len(cache), (pos * 3) % len(cache)])
+    interleaved, glm = layout != "sections", layout == "interleaved_glm"
+    axis_map = MRotaryEmbedding._build_axis_map(
+        SimpleNamespace(
+            mrope_section=section,
+            rotary_dim=d,
+            mrope_interleaved=interleaved,
+            mrope_interleaved_glm=glm,
+            cos_sin_cache=cache,
+        )
+    )
+    expected = x.clone()
+    q, k, _ = expected.split([hq * d, hk * d, hk * d], dim=-1)
+    fused_inplace_qknorm(q.view(batch, hq, d), k.view(batch, hk, d), qw, kw, 1e-6)
+    triton_mrope_fused(
+        q,
+        k,
+        cache,
+        positions,
+        section,
+        d,
+        d,
+        interleaved,
+        glm,
+        True,
+        axis_map if glm else None,
+    )
+    table_qk_norm_rope_(
+        x,
+        positions,
+        qw,
+        kw,
+        cache,
+        hq,
+        hk,
+        d,
+        1e-6,
+        match_cuda=True,
+        axis_map=axis_map,
+    )
+    assert_same_bits(x, expected)
+
+
 @pytest.mark.parametrize("d", [64, 128, 256])
 @pytest.mark.parametrize("eps", [1e-6, 1e-5])
 def test_extreme_values_and_rounding(d, eps):
@@ -115,7 +172,7 @@ def test_cuda_mode_rejects_invalid_inputs(invalid):
     elif invalid == "qkv_stride":
         x = x[:, ::2]
     elif invalid == "table":
-        cache = cache.bfloat16()
+        cache = cache.half()
     else:
         pos = pos.float()
     with pytest.raises(AssertionError):