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

[NPU]use triton split_qkvgate_gemma_rmsnorm_rope for Qwen3.5 and Qwen3_next (#23925)

Browse Source
This commit is contained in:
Liwansi
2026-05-20 20:22:10 +08:00
committed by GitHub
parent 34d3e23232
commit 55ba03db6a
3 changed files with 194 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/rotary_embedding/mrope.py
+93 -3
View File
@@ -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))],
python/sglang/srt/models/qwen3_5.py
+52 -8
View File
@@ -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:
python/sglang/srt/models/qwen3_next.py
+49 -8
View File
@@ -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)