minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[rotary] Fix the fused Qwen3.5 RoPE kernel discarding mrope height and width (#34446)

Browse Source
This commit is contained in:
Jason Wiemels
2026-08-30 15:37:40 +08:00
committed by GitHub
parent d249672ad3
commit e635577431
5 changed files with 374 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py
+30 -5
View File
@@ -1,7 +1,7 @@
"""Fused Q/K GemmaRMSNorm + NeoX RoPE + gate deinterleave (Triton).
Single kernel launch fusing per-head GemmaRMSNorm, partial NeoX RoPE,
and gate deinterleave for Qwen3.5's interleaved Q+Gate layout.
Single kernel launch fusing per-head GemmaRMSNorm, partial NeoX RoPE over 1-D or
mrope positions, and gate deinterleave for Qwen3.5's interleaved Q+Gate layout.
2D grid (T, num_q_heads + num_kv_heads) — each program handles one
(token, head) pair. Q programs also copy the gate slice.
@@ -39,12 +39,14 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
k_weight_ptr,
cos_sin_cache_ptr,
positions_ptr,
mrope_axis_map_ptr,
stride_qg_t,
stride_k_t,
stride_qo_t,
stride_ko_t,
stride_gate_t,
stride_cos_t,
stride_pos_axis,
NUM_Q_HEADS: tl.constexpr,
NUM_KV_HEADS: tl.constexpr,
HEAD_DIM: tl.constexpr,
@@ -56,6 +58,7 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
FP16: tl.constexpr,
HAS_PASS: tl.constexpr,
HAS_GATE: tl.constexpr,
MROPE: tl.constexpr,
ENABLE_PDL: tl.constexpr,
):
token = tl.program_id(0)
@@ -104,8 +107,14 @@ def _fused_qk_rmsnorm_rope_gate_kernel(
xr1 = (xr1 * inv_rms * (wr1 + 1.0)).to(out_dtype).to(tl.float32)
xr2 = (xr2 * inv_rms * (wr2 + 1.0)).to(out_dtype).to(tl.float32)
pos = tl.load(positions_ptr + token).to(tl.int64)
cache_off = pos * stride_cos_t
if MROPE:
axis = tl.load(mrope_axis_map_ptr + rot_offs, mask=rot_mask, other=0)
pos = tl.load(
positions_ptr + axis * stride_pos_axis + token, mask=rot_mask, other=0
)
else:
pos = tl.load(positions_ptr + token)
cache_off = pos.to(tl.int64) * stride_cos_t
cos = tl.load(
cos_sin_cache_ptr + cache_off + rot_offs, mask=rot_mask, other=0.0
).to(tl.float32)
@@ -141,6 +150,7 @@ def fused_qk_gemma_rmsnorm_rope_gate(
head_dim: int,
rotary_dim: int,
has_gate: bool = True,
mrope_axis_map: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
"""Fused QK GemmaRMSNorm + NeoX RoPE + gate deinterleave.
@@ -149,8 +159,20 @@ def fused_qk_gemma_rmsnorm_rope_gate(
k: [T, num_kv_heads * head_dim]
q_weight, k_weight: [head_dim] — raw GemmaRMSNorm weights (kernel adds +1.0)
cos_sin_cache: [max_seq_len, rotary_dim] — [cos..., sin...]
positions: [T] — token positions
positions: [T] token positions, or [3, T] mrope rows (temporal, height, width)
mrope_axis_map: [rotary_dim // 2] — the axis owning each rotary lane, from
MRotaryEmbedding
"""
assert positions.dim() in (1, 2), f"want [T] or [3, T], got {positions.shape}"
mrope = positions.dim() == 2
assert mrope == (mrope_axis_map is not None), "mrope_axis_map needs [3, T]"
if mrope:
assert positions.shape[0] == 3 and positions.stride(1) == 1, (
f"want [3, T] contiguous over T, got {positions.shape} "
f"stride {positions.stride()}"
)
lanes = rotary_dim // 2
assert mrope_axis_map.shape == (lanes,), f"want one axis per lane ({lanes})"
T = q_gate.shape[0]
q_size = num_q_heads * head_dim
kv_size = num_kv_heads * head_dim
@@ -178,12 +200,14 @@ def fused_qk_gemma_rmsnorm_rope_gate(
k_weight,
cos_sin_cache,
positions,
mrope_axis_map,
q_gate.stride(0),
k.stride(0),
q_out.stride(0),
k_out.stride(0),
gate_out.stride(0),
cos_sin_cache.stride(0),
positions.stride(0),
NUM_Q_HEADS=num_q_heads,
NUM_KV_HEADS=num_kv_heads,
HEAD_DIM=head_dim,
@@ -195,6 +219,7 @@ def fused_qk_gemma_rmsnorm_rope_gate(
FP16=q_gate.dtype == torch.float16,
HAS_PASS=rotary_dim < head_dim,
HAS_GATE=has_gate,
MROPE=mrope,
ENABLE_PDL=_ENABLE_PDL,
)
python/sglang/srt/layers/rotary_embedding/mrope.py
+39 -33
View File
@@ -100,40 +100,42 @@ class MRotaryEmbedding(RotaryEmbedding):
f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})"
)
# MRoPE axis_map interleaving pattern depends on mrope_section sizes.
# The algorithm cycles through axes [0(T), 1(H), 2(W)] round-robin,
# skipping any axis that has exhausted its allocated pairs.
#
# For GLM-V (mrope_section=[8,12,12]):
# T(8) < H(12) = W(12), so T exhausts first at pair 24.
# Result: [0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 0,1,2, 1,1,2, 1,1,2, 2,2]
# After T runs out, only H and W fill the remaining slots.
#
# For Qwen3-VL (mrope_section=[24,20,20]):
# T(24) > H(20) = W(20), so H and W exhaust first near the tail.
# Result: [0,1,2, 0,1,2, ...repeated evenly..., 0,1, 0,1, 0,0]
# After H/W run out, T fills the remaining slots.
if self.mrope_interleaved_glm:
num_pairs = rotary_dim // 2
axis_map = torch.empty(num_pairs, dtype=torch.long)
assert sum(self.mrope_section) == num_pairs
counts = [0, 0, 0]
current_ax = 0
for i in range(num_pairs):
current_ax = i % 3
while counts[current_ax] >= self.mrope_section[current_ax]:
current_ax = (current_ax + 1) % 3
axis_map[i] = current_ax
counts[current_ax] += 1
self.register_buffer("axis_map", axis_map, persistent=False)
else:
self.axis_map = None
self.register_buffer("axis_map", self._build_axis_map(), persistent=False)
if self._force_native:
self._forward_method = self.forward_native
def _build_axis_map(self) -> Optional[torch.Tensor]:
"""Which of the temporal, height and width axes owns each rotary lane."""
if not self.mrope_section:
return None
section = self.mrope_section
num_pairs = self.rotary_dim // 2
assert (
len(section) == 3 and sum(section) == num_pairs
), f"mrope_section {section} must be three axes summing to {num_pairs}"
if self.mrope_interleaved_glm:
axes = []
spent = [0, 0, 0]
for lane in range(num_pairs):
axis = lane % 3
while spent[axis] >= section[axis]:
axis = (axis + 1) % 3
spent[axis] += 1
axes.append(axis)
elif self.mrope_interleaved:
axes = [0] * num_pairs
for axis in (1, 2):
for lane in range(axis, min(3 * section[axis], num_pairs), 3):
axes[lane] = axis
else:
axes = [axis for axis, size in enumerate(section) for _ in range(size)]
return torch.tensor(axes, dtype=torch.long, device=self.cos_sin_cache.device)
@property
def _legacy_axis_map(self) -> Optional[torch.Tensor]:
"""The map only where the older rope kernels read it; one is out of tree."""
return self.axis_map if self.mrope_interleaved_glm else None
def get_cos_sin_with_position(self, positions):
if positions.ndim == 1:
return super().get_cos_sin_with_position(positions)
@@ -269,7 +271,7 @@ class MRotaryEmbedding(RotaryEmbedding):
self.mrope_interleaved,
self.mrope_interleaved_glm,
self.is_neox_style,
self.axis_map,
self._legacy_axis_map,
)
return query, key
@@ -319,7 +321,7 @@ class MRotaryEmbedding(RotaryEmbedding):
self.mrope_interleaved,
self.mrope_interleaved_glm,
self.is_neox_style,
self.axis_map,
self._legacy_axis_map,
)
return query, key
return self.forward_native(positions, query, key, fused_set_kv_buffer_arg)
@@ -515,6 +517,10 @@ class Ernie4_5_VLRotaryEmbedding(MRotaryEmbedding):
)
self._apply_rotary_emb_wrapped = torch.compile(dynamic=True)(apply_rotary_emb)
def _build_axis_map(self) -> Optional[torch.Tensor]:
"""No map: the shared builder reads mrope_section as t, h, w, Ernie as h, w, t."""
return None
def forward_native(
self,
positions: torch.Tensor,
python/sglang/srt/models/qwen3_5.py
+1
View File
@@ -1244,6 +1244,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
self.head_dim,
self.rotary_emb.rotary_dim,
has_gate=self.attn_output_gate,
mrope_axis_map=(self.rotary_emb.axis_map if positions.dim() == 2 else None),
)
seq_len = hidden_states.shape[0]
q = q_out.view(seq_len, -1)
test/registered/kernels/ops/attention/test_fused_qk_rmsnorm_rope_gate.py
+195
View File
@@ -0,0 +1,195 @@
import unittest
import torch
from sglang.kernels.ops.attention.fused_qk_rmsnorm_rope_gate import (
fused_qk_gemma_rmsnorm_rope_gate,
)
from sglang.srt.layers.rotary_embedding.mrope import MRotaryEmbedding
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large")
def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
dtype = x.dtype
x = x.float()
x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
return (x * (1.0 + weight.float())).to(dtype)
def neox_rope(
x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, rotary_dim: int
) -> torch.Tensor:
half = rotary_dim // 2
rotated, passthrough = x[..., :rotary_dim], x[..., rotary_dim:]
first, second = rotated[..., :half], rotated[..., half:]
cos = cos.unsqueeze(1).to(x.dtype)
sin = sin.unsqueeze(1).to(x.dtype)
return torch.cat(
[first * cos - second * sin, second * cos + first * sin, passthrough], dim=-1
)
class TestFusedQKRMSNormRoPEGate(CustomTestCase):
def setUp(self):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
torch.manual_seed(0)
self.tokens = 17
self.num_q_heads = 8
self.num_kv_heads = 2
self.head_dim = 128
self.rotary_dim = 64
self.eps = 1e-6
device, dtype = "cuda", torch.bfloat16
self.q_gate = torch.randn(
self.tokens,
self.num_q_heads * 2 * self.head_dim,
device=device,
dtype=dtype,
)
self.k = torch.randn(
self.tokens, self.num_kv_heads * self.head_dim, device=device, dtype=dtype
)
self.q_weight = torch.randn(self.head_dim, device=device, dtype=dtype)
self.k_weight = torch.randn(self.head_dim, device=device, dtype=dtype)
inv_freq = 10000 ** (
-torch.arange(self.rotary_dim // 2, device=device).float()
* 2
/ self.rotary_dim
)
angles = torch.arange(512, device=device).float().unsqueeze(1) * inv_freq
self.cos_sin_cache = torch.cat([angles.cos(), angles.sin()], dim=-1).to(dtype)
def call(self, positions, cos_sin_cache=None, mrope_axis_map=None, rotary_dim=None):
if cos_sin_cache is None:
cos_sin_cache = self.cos_sin_cache
return fused_qk_gemma_rmsnorm_rope_gate(
self.q_gate,
self.k,
self.q_weight,
self.k_weight,
cos_sin_cache,
positions,
self.eps,
self.num_q_heads,
self.num_kv_heads,
self.head_dim,
rotary_dim or self.rotary_dim,
has_gate=True,
mrope_axis_map=mrope_axis_map,
)
def build_mrope(self, mrope_section, interleaved):
return MRotaryEmbedding(
head_size=self.head_dim,
rotary_dim=2 * sum(mrope_section),
max_position_embeddings=512,
base=10000,
is_neox_style=True,
dtype=torch.bfloat16,
mrope_section=mrope_section,
mrope_interleaved=interleaved,
).to("cuda")
def graph_buffer_positions(self):
buffer = torch.zeros(3, 4 * self.tokens, dtype=torch.int64, device="cuda")
buffer[:, : self.tokens] = torch.stack(
[
torch.arange(self.tokens, device="cuda") % 7,
torch.arange(self.tokens, device="cuda") % 5 + 3,
torch.arange(self.tokens, device="cuda") % 3 + 11,
]
)
return buffer[:, : self.tokens]
def test_matches_reference_for_1d_positions(self):
positions = torch.arange(self.tokens, device="cuda", dtype=torch.int64)
q_out, k_out, gate_out = self.call(positions)
packed = self.q_gate.view(self.tokens, self.num_q_heads, 2 * self.head_dim)
cos, sin = self.cos_sin_cache[positions].chunk(2, dim=-1)
want_q = neox_rope(
gemma_rmsnorm(packed[..., : self.head_dim], self.q_weight, self.eps),
cos,
sin,
self.rotary_dim,
)
want_k = neox_rope(
gemma_rmsnorm(
self.k.view(self.tokens, self.num_kv_heads, self.head_dim),
self.k_weight,
self.eps,
),
cos,
sin,
self.rotary_dim,
)
torch.testing.assert_close(
q_out.view_as(want_q).float(), want_q.float(), atol=2e-2, rtol=2e-2
)
torch.testing.assert_close(
k_out.view_as(want_k).float(), want_k.float(), atol=2e-2, rtol=2e-2
)
torch.testing.assert_close(
gate_out.float(), packed[..., self.head_dim :].float(), atol=0, rtol=0
)
def test_mrope_matches_the_rotary_module(self):
"""With t == h == w every layout agrees, so only distinct rows catch a wrong
axis. Interleaved [11, 11, 10] is what Qwen3.6-35B-A3B ships, and [24, 20, 20]
fills the head dimension, leaving no pass-through tail.
"""
for section, interleaved in (
([11, 11, 10], False),
([11, 11, 10], True),
([24, 20, 20], True),
):
with self.subTest(section=section, interleaved=interleaved):
rope = self.build_mrope(section, interleaved)
positions = self.graph_buffer_positions()
q_in = self.q_gate.view(
self.tokens, self.num_q_heads, 2 * self.head_dim
)[..., : self.head_dim].reshape(self.tokens, -1)
want_q, want_k = rope.forward_native(
positions,
gemma_rmsnorm(
q_in.view(self.tokens, self.num_q_heads, self.head_dim),
self.q_weight,
self.eps,
).reshape(self.tokens, -1),
gemma_rmsnorm(
self.k.view(self.tokens, self.num_kv_heads, self.head_dim),
self.k_weight,
self.eps,
).reshape(self.tokens, -1),
)
q_out, k_out, _gate = self.call(
positions,
cos_sin_cache=rope.cos_sin_cache,
mrope_axis_map=rope.axis_map,
rotary_dim=rope.rotary_dim,
)
torch.testing.assert_close(
q_out.float(), want_q.float(), atol=2e-2, rtol=2e-2
)
torch.testing.assert_close(
k_out.float(), want_k.float(), atol=2e-2, rtol=2e-2
)
def test_rejects_positions_and_map_apart(self):
flat = torch.arange(self.tokens, device="cuda", dtype=torch.int64)
axis_map = self.build_mrope([11, 11, 10], interleaved=True).axis_map
cases = ((flat.unsqueeze(0).repeat(3, 1), None), (flat, axis_map))
for positions, axis_map_passed in cases:
with self.subTest(mrope_positions=positions.dim() == 2):
with self.assertRaises(AssertionError):
self.call(positions, mrope_axis_map=axis_map_passed)
if __name__ == "__main__":
unittest.main()
test/registered/rotary/test_mrope_axis_map.py
+109
View File
@@ -0,0 +1,109 @@
import unittest
from unittest.mock import patch
import torch
from sglang.srt.layers.rotary_embedding.mrope import (
Ernie4_5_VLRotaryEmbedding,
MRotaryEmbedding,
apply_interleaved_rope,
)
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def build_mrope(
mrope_section: list[int],
rotary_dim: int,
interleaved: bool = False,
glm: bool = False,
) -> MRotaryEmbedding:
return MRotaryEmbedding(
head_size=rotary_dim,
rotary_dim=rotary_dim,
max_position_embeddings=64,
base=10000,
is_neox_style=True,
dtype=torch.float32,
mrope_section=mrope_section,
mrope_interleaved=interleaved,
mrope_interleaved_glm=glm,
)
def select_by_axis(table: torch.Tensor, axis_map: torch.Tensor) -> torch.Tensor:
"""Take every lane from the axis the map names, the way the kernel does."""
lanes = torch.arange(table.shape[2])
return table[axis_map, :, lanes].T
class TestMRopeAxisMap(CustomTestCase):
def setUp(self):
cpu_patch = patch("sglang.srt.layers.rotary_embedding.base._is_cpu", True)
cpu_patch.start()
self.addCleanup(cpu_patch.stop)
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
torch.manual_seed(0)
def test_interleaved_matches_apply_interleaved_rope(self):
# Under [1, 1, 30] the reference places only 10 of the 30 lanes it was asked
# for, so this pins that the map reproduces that loss.
for section, rotary_dim in (
([24, 20, 20], 128),
([11, 11, 10], 64),
([1, 1, 30], 64),
):
with self.subTest(section=section):
rope = build_mrope(section, rotary_dim, interleaved=True)
table = torch.randn(3, 7, rotary_dim // 2)
torch.testing.assert_close(
select_by_axis(table, rope.axis_map),
apply_interleaved_rope(table, section),
atol=0,
rtol=0,
)
def test_contiguous_matches_section_split(self):
section, rotary_dim = [24, 20, 20], 128
rope = build_mrope(section, rotary_dim)
table = torch.randn(3, 7, rotary_dim // 2)
torch.testing.assert_close(
select_by_axis(table, rope.axis_map),
torch.cat(
[m[i] for i, m in enumerate(table.split(section, dim=-1))], dim=-1
),
atol=0,
rtol=0,
)
def test_glm_map_keeps_its_round_robin_order(self):
"""GLM's kernel is out of tree, so the order is pinned rather than compared."""
rope = build_mrope([8, 12, 12], 64, interleaved=True, glm=True)
want = [0, 1, 2] * 8 + [1, 1, 2, 1, 1, 2, 2, 2]
self.assertEqual(rope.axis_map.tolist(), want)
def test_only_glm_reaches_the_older_kernels(self):
self.assertIsNone(
build_mrope([24, 20, 20], 128, interleaved=True)._legacy_axis_map
)
glm = build_mrope([8, 12, 12], 64, interleaved=True, glm=True)
torch.testing.assert_close(glm._legacy_axis_map, glm.axis_map, atol=0, rtol=0)
def test_ernie_has_no_map(self):
ernie = Ernie4_5_VLRotaryEmbedding(
head_size=128,
rotary_dim=128,
max_position_embeddings=64,
base=10000,
is_neox_style=True,
dtype=torch.float32,
mrope_section=[16, 16, 32],
)
self.assertIsNone(ernie.axis_map)
if __name__ == "__main__":
unittest.main()