From 4a8200565e1c659100fb3910bcd3be194e5506f2 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Thu, 25 Jun 2026 15:58:54 +0800 Subject: [PATCH] Fused QK GemmaRMSNorm + RoPE + gate kernel for Qwen3.5 (#28320) --- .../srt/layers/fused_qk_rmsnorm_rope_gate.py | 201 ++++++++++++++++++ python/sglang/srt/models/qwen3_5.py | 41 +++- 2 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/layers/fused_qk_rmsnorm_rope_gate.py diff --git a/python/sglang/srt/layers/fused_qk_rmsnorm_rope_gate.py b/python/sglang/srt/layers/fused_qk_rmsnorm_rope_gate.py new file mode 100644 index 000000000..56d442098 --- /dev/null +++ b/python/sglang/srt/layers/fused_qk_rmsnorm_rope_gate.py @@ -0,0 +1,201 @@ +"""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. + +2D grid (T, num_q_heads + num_kv_heads) — each program handles one +(token, head) pair. Q programs also copy the gate slice. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +def _pdl_supported() -> bool: + """Check if Programmatic Dependent Launch is supported (NVIDIA SM >= 90).""" + if not torch.cuda.is_available(): + return False + try: + major, _ = torch.cuda.get_device_capability() + return major >= 9 + except Exception: + return False + + +_ENABLE_PDL = _pdl_supported() + + +@triton.jit +def _fused_qk_rmsnorm_rope_gate_kernel( + q_gate_ptr, + k_ptr, + q_out_ptr, + k_out_ptr, + gate_out_ptr, + q_weight_ptr, + k_weight_ptr, + cos_sin_cache_ptr, + positions_ptr, + stride_qg_t, + stride_k_t, + stride_qo_t, + stride_ko_t, + stride_gate_t, + stride_cos_t, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + ROTARY_DIM: tl.constexpr, + HALF_ROTARY: tl.constexpr, + HEAD_BLOCK: tl.constexpr, + ROT_HALF_BLOCK: tl.constexpr, + EPS: tl.constexpr, + FP16: tl.constexpr, + HAS_PASS: tl.constexpr, + HAS_GATE: tl.constexpr, + ENABLE_PDL: tl.constexpr, +): + token = tl.program_id(0) + head = tl.program_id(1) + is_k = head >= NUM_Q_HEADS + local_head = tl.where(is_k, head - NUM_Q_HEADS, head) + out_dtype = tl.float16 if FP16 else tl.bfloat16 + + if is_k: + in_base = k_ptr + token * stride_k_t + local_head * HEAD_DIM + w_ptr = k_weight_ptr + out_base = k_out_ptr + token * stride_ko_t + local_head * HEAD_DIM + else: + if HAS_GATE: + in_base = q_gate_ptr + token * stride_qg_t + local_head * 2 * HEAD_DIM + else: + in_base = q_gate_ptr + token * stride_qg_t + local_head * HEAD_DIM + w_ptr = q_weight_ptr + out_base = q_out_ptr + token * stride_qo_t + local_head * HEAD_DIM + + # Full load -> RMSNorm variance + head_offs = tl.arange(0, HEAD_BLOCK) + head_mask = head_offs < HEAD_DIM + x = tl.load(in_base + head_offs, mask=head_mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + head_offs, mask=head_mask, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=0) / HEAD_DIM + inv_rms = tl.rsqrt(var + EPS) + x_norm = (x * inv_rms * (w + 1.0)).to(out_dtype).to(tl.float32) + + # Pass-through tail [rotary_dim, head_dim) + if HAS_PASS: + pass_mask = head_mask & (head_offs >= ROTARY_DIM) + tl.store(out_base + head_offs, x_norm, mask=pass_mask) + + # Reload rotary portion from L1 -> re-norm -> RoPE + rot_offs = tl.arange(0, ROT_HALF_BLOCK) + rot_mask = rot_offs < HALF_ROTARY + xr1 = tl.load(in_base + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + xr2 = tl.load(in_base + HALF_ROTARY + rot_offs, mask=rot_mask, other=0.0).to( + tl.float32 + ) + wr1 = tl.load(w_ptr + rot_offs, mask=rot_mask, other=0.0).to(tl.float32) + wr2 = tl.load(w_ptr + HALF_ROTARY + rot_offs, mask=rot_mask, other=0.0).to( + tl.float32 + ) + 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 + cos = tl.load( + cos_sin_cache_ptr + cache_off + rot_offs, mask=rot_mask, other=0.0 + ).to(tl.float32) + sin = tl.load( + cos_sin_cache_ptr + cache_off + HALF_ROTARY + rot_offs, mask=rot_mask, other=0.0 + ).to(tl.float32) + tl.store(out_base + rot_offs, (xr1 * cos - xr2 * sin), mask=rot_mask) + tl.store(out_base + HALF_ROTARY + rot_offs, (xr2 * cos + xr1 * sin), mask=rot_mask) + + # Gate copy (Q heads only) + if HAS_GATE and not is_k: + gate_in = in_base + HEAD_DIM + gate_out = gate_out_ptr + token * stride_gate_t + local_head * HEAD_DIM + g = tl.load(gate_in + head_offs, mask=head_mask, other=0.0) + tl.store(gate_out + head_offs, g, mask=head_mask) + + # PDL: signal dependent kernels (attention/allreduce) can start early. + # Only available on NVIDIA Hopper+ (sm_90+); guarded for AMD/other backends. + if ENABLE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +def fused_qk_gemma_rmsnorm_rope_gate( + q_gate: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_dim: int, + has_gate: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Fused QK GemmaRMSNorm + NeoX RoPE + gate deinterleave. + + Args: + q_gate: [T, num_q_heads * (1 + has_gate) * head_dim] — interleaved Q+Gate if has_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 + """ + T = q_gate.shape[0] + q_size = num_q_heads * head_dim + kv_size = num_kv_heads * head_dim + + q_out = torch.empty(T, q_size, dtype=q_gate.dtype, device=q_gate.device) + k_out = torch.empty(T, kv_size, dtype=k.dtype, device=k.device) + gate_out = ( + torch.empty(T, num_q_heads, head_dim, dtype=q_gate.dtype, device=q_gate.device) + if has_gate + else q_out + ) + + half_rotary = rotary_dim // 2 + head_block = triton.next_power_of_2(head_dim) + rot_half_block = triton.next_power_of_2(half_rotary) + + grid = (T, num_q_heads + num_kv_heads) + _fused_qk_rmsnorm_rope_gate_kernel[grid]( + q_gate, + k, + q_out, + k_out, + gate_out, + q_weight, + k_weight, + cos_sin_cache, + positions, + q_gate.stride(0), + k.stride(0), + q_out.stride(0), + k_out.stride(0), + gate_out.stride(0), + cos_sin_cache.stride(0), + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + ROTARY_DIM=rotary_dim, + HALF_ROTARY=half_rotary, + HEAD_BLOCK=head_block, + ROT_HALF_BLOCK=rot_half_block, + EPS=eps, + FP16=q_gate.dtype == torch.float16, + HAS_PASS=rotary_dim < head_dim, + HAS_GATE=has_gate, + ENABLE_PDL=_ENABLE_PDL, + ) + + return q_out, k_out, gate_out if has_gate else None diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 8d99affe2..8b1f8052e 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -137,6 +137,11 @@ def _disable_shared_experts_fusion() -> bool: return get_global_server_args().disable_shared_experts_fusion +if _is_cuda: + from sglang.srt.layers.fused_qk_rmsnorm_rope_gate import ( + fused_qk_gemma_rmsnorm_rope_gate, + ) + if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( split_qkvgate_gemma_rmsnorm_rope, @@ -887,6 +892,35 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): k = k_by_head.view(k.shape) return q, k + def forward_prepare_cuda_fused(self, positions, hidden_states): + """Fused QK GemmaRMSNorm + NeoX RoPE + gate deinterleave.""" + qkv, _ = self.qkv_proj(hidden_states) + if self.attn_output_gate: + q_gate, k, v = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + else: + q_gate, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q_out, k_out, gate_out = fused_qk_gemma_rmsnorm_rope_gate( + q_gate, + k, + self.q_norm.weight.data, + self.k_norm.weight.data, + self.rotary_emb.cos_sin_cache, + positions, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.rotary_emb.rotary_dim, + has_gate=self.attn_output_gate, + ) + seq_len = hidden_states.shape[0] + q = q_out.view(seq_len, -1) + k = k_out.view(seq_len, -1) + gate = gate_out.view(seq_len, -1) if gate_out is not None else None + return q, k, v, gate + def forward_prepare_native(self, positions, hidden_states): qkv, _ = self.qkv_proj(hidden_states) if self.attn_output_gate: @@ -960,7 +994,12 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): forward_batch: ForwardBatch, ) -> torch.Tensor: """Full attention forward pass.""" - if (_is_hip or _is_xpu) and self.attn_output_gate: + if _is_cuda and self.attn_output_gate: + q, k, v, gate = self.forward_prepare_cuda_fused( + positions=positions, + hidden_states=hidden_states, + ) + elif (_is_hip or _is_xpu) and self.attn_output_gate: q, k, v, gate = self.forward_prepare_fused_gate( positions=positions, hidden_states=hidden_states,