From 55ba03db6a46a676470096a9da1fdae5ad7d0e6f Mon Sep 17 00:00:00 2001 From: Liwansi <62291011+Liwansi@users.noreply.github.com> Date: Wed, 20 May 2026 20:22:10 +0800 Subject: [PATCH] [NPU]use triton split_qkvgate_gemma_rmsnorm_rope for Qwen3.5 and Qwen3_next (#23925) --- .../srt/layers/rotary_embedding/mrope.py | 96 ++++++++++++++++++- python/sglang/srt/models/qwen3_5.py | 60 ++++++++++-- python/sglang/srt/models/qwen3_next.py | 57 +++++++++-- 3 files changed, 194 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index ade5bd5a8..775fcc0c7 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -19,7 +19,7 @@ from sglang.srt.layers.rotary_embedding.yarn import ( yarn_linear_ramp_mask, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import cpu_has_amx_support, is_cuda, is_npu +from sglang.srt.utils import cpu_has_amx_support, is_cuda, is_npu, support_triton _is_cuda = is_cuda() _is_npu = is_npu() @@ -32,6 +32,92 @@ if _is_npu: import torch_npu +import triton +import triton.language as tl + + +@triton.jit +def apply_interleaved_rope_kernel( + x_ptr, + out_ptr, + S: tl.constexpr, + D: tl.constexpr, + stride_x_m, + stride_x_s, + stride_out_s, + section_1_end, + section_2_end, + BLOCK_S: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + start_s = tl.program_id(0) * BLOCK_S + s_offsets = start_s + tl.arange(0, BLOCK_S) + + dim_offset = tl.program_id(1) * BLOCK_SIZE + dim_indices = dim_offset + tl.arange(0, BLOCK_SIZE) + + mask_s = s_offsets < S + mask_d = dim_indices < D + mask = mask_s[:, None] & mask_d[None, :] + + val_ptr = ( + x_ptr + 0 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] + ) + val = tl.load(val_ptr, mask=mask, other=0.0) + + cond_a = (dim_indices[None, :] % 3 == 1) & ( + dim_indices[None, :] < section_1_end * 3 + ) + val_a_ptr = ( + x_ptr + 1 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] + ) + val_a = tl.load(val_a_ptr, mask=mask & cond_a, other=0.0) + + cond_b = (dim_indices[None, :] % 3 == 2) & ( + dim_indices[None, :] < section_2_end * 3 + ) + val_b_ptr = ( + x_ptr + 2 * stride_x_m + s_offsets[:, None] * stride_x_s + dim_indices[None, :] + ) + val_b = tl.load(val_b_ptr, mask=mask & cond_b, other=0.0) + + val = tl.where(cond_a, val_a, val) + val = tl.where(cond_b, val_b, val) + + out_ptr = out_ptr + s_offsets[:, None] * stride_out_s + dim_indices[None, :] + tl.store(out_ptr, val, mask=mask) + + +def apply_interleaved_rope_triton(x: torch.Tensor, mrope_section: list) -> torch.Tensor: + x = x.contiguous() + M, S, D = x.shape + + out = torch.empty((S, D), dtype=x.dtype, device=x.device) + + BLOCK_S = 64 + BLOCK_SIZE = 128 + + grid = (triton.cdiv(S, BLOCK_S), triton.cdiv(D, BLOCK_SIZE)) + + section_1_end = mrope_section[1] + section_2_end = mrope_section[2] + + apply_interleaved_rope_kernel[grid]( + x, + out, + S, + D, + x.stride(0), + x.stride(1), + out.stride(0), + section_1_end, + section_2_end, + BLOCK_S=BLOCK_S, + BLOCK_SIZE=BLOCK_SIZE, + ) + return out + + def apply_interleaved_rope(x: torch.Tensor, mrope_section: list) -> torch.Tensor: x_t = x[0].clone() x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3] @@ -131,8 +217,12 @@ class MRotaryEmbedding(RotaryEmbedding): last_dim = cos_sin.size()[-1] cos, sin = cos_sin.chunk(2, dim=-1) if self.mrope_interleaved: - cos = apply_interleaved_rope(cos, self.mrope_section) - sin = apply_interleaved_rope(sin, self.mrope_section) + if support_triton(get_global_server_args().attention_backend): + cos = apply_interleaved_rope_triton(cos, self.mrope_section) + sin = apply_interleaved_rope_triton(sin, self.mrope_section) + else: + cos = apply_interleaved_rope(cos, self.mrope_section) + sin = apply_interleaved_rope(sin, self.mrope_section) else: cos = torch.cat( [m[i] for i, m in enumerate(cos.split(self.mrope_section, dim=-1))], diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 8f1f32b95..9a95dd422 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -109,6 +109,11 @@ _is_amx_available = cpu_has_amx_support() cached_get_processor = lru_cache(get_processor) +if _is_npu: + from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( + split_qkvgate_gemma_rmsnorm_rope, + ) + class Qwen3_5GatedDeltaNet(nn.Module): def __init__( @@ -841,15 +846,8 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): k = k_by_head.view(k.shape) return q, k - def self_attention( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - ) -> torch.Tensor: - """Full attention forward pass.""" + def forward_prepare_native(self, positions, hidden_states): 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 @@ -861,9 +859,55 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): gate = gate.reshape(*orig_shape, -1) else: q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + gate = None q, k = self._apply_qk_norm(q, k) q, k = self.rotary_emb(positions, q, k) + return q, k, v, gate + + def forward_prepare_npu(self, positions, hidden_states, forward_batch): + qkv, _ = self.qkv_proj(hidden_states) + # Calculate first full attention layer ID based on config + if self.attn.layer_id == (self.config.full_attention_interval - 1): + self.rotary_emb.get_cos_sin_with_position(positions) + + q, k, v, gate = split_qkvgate_gemma_rmsnorm_rope( + qkv, + self.rotary_emb.position_sin, + self.rotary_emb.position_cos, + self.q_size, + self.kv_size, + self.head_dim, + int(self.head_dim * self.partial_rotary_factor), + eps=self.q_norm.variance_epsilon, + q_weight=self.q_norm.weight, + k_weight=self.k_norm.weight, + ) + return q, k, v, gate + + def self_attention( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + """Full attention forward pass.""" + if ( + not _is_npu + or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() + or not self.attn_output_gate + ): + q, k, v, gate = self.forward_prepare_native( + positions=positions, + hidden_states=hidden_states, + ) + else: + q, k, v, gate = self.forward_prepare_npu( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + attn_output = self.attn(q, k, v, forward_batch) if self.attn_output_gate: diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index a3214a0bd..5b79a6b9f 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -69,6 +69,9 @@ if _is_npu: from sgl_kernel_npu.fla.utils import ( fused_qkvzba_split_reshape_cat as fused_qkvzba_split_reshape_cat_npu, ) + from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( + split_qkvgate_gemma_rmsnorm_rope, + ) fused_qkvzba_split_reshape_cat = fused_qkvzba_split_reshape_cat_npu @@ -751,14 +754,8 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): k = k_by_head.view(k.shape) return q, k - def self_attention( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - ) -> torch.Tensor: + def forward_prepare_native(self, positions, hidden_states): 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 @@ -770,10 +767,54 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): gate = gate.reshape(*orig_shape, -1) else: q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + gate = None q, k = self._apply_qk_norm(q, k) - q, k = self.rotary_emb(positions, q, k) + return q, k, v, gate + + def forward_prepare_npu(self, positions, hidden_states, forward_batch): + qkv, _ = self.qkv_proj(hidden_states) + # Calculate first full attention layer ID based on config + if self.attn.layer_id == (self.config.full_attention_interval - 1): + self.rotary_emb.get_cos_sin_with_position(positions) + + q, k, v, gate = split_qkvgate_gemma_rmsnorm_rope( + qkv, + self.rotary_emb.position_sin, + self.rotary_emb.position_cos, + self.q_size, + self.kv_size, + self.head_dim, + int(self.head_dim * self.partial_rotary_factor), + eps=self.q_norm.variance_epsilon, + q_weight=self.q_norm.weight, + k_weight=self.k_norm.weight, + ) + return q, k, v, gate + + def self_attention( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + """Full attention forward pass.""" + if ( + not _is_npu + or forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() + or not self.attn_output_gate + ): + q, k, v, gate = self.forward_prepare_native( + positions=positions, + hidden_states=hidden_states, + ) + else: + q, k, v, gate = self.forward_prepare_npu( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) attn_output = self.attn(q, k, v, forward_batch)