minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 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)