diff --git a/python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py b/python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py index 56d442098..f84a07eaa 100644 --- a/python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py +++ b/python/sglang/kernels/ops/attention/fused_qk_rmsnorm_rope_gate.py @@ -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, ) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index 1948ff7b9..f56ff48e4 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -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, diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 2871de406..f1f8582a6 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -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) diff --git a/test/registered/kernels/ops/attention/test_fused_qk_rmsnorm_rope_gate.py b/test/registered/kernels/ops/attention/test_fused_qk_rmsnorm_rope_gate.py new file mode 100644 index 000000000..b9f3338f6 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_fused_qk_rmsnorm_rope_gate.py @@ -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() diff --git a/test/registered/rotary/test_mrope_axis_map.py b/test/registered/rotary/test_mrope_axis_map.py new file mode 100644 index 000000000..3898ee731 --- /dev/null +++ b/test/registered/rotary/test_mrope_axis_map.py @@ -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()