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

use sgl_kernel_npu rmsrope accelerate llada2 (#27127)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
wenxuewuhd
2026-07-20 22:46:17 +03:00
committed by GitHub
co-authored by gemini-code-assist[bot] ronnie_zheng
parent ff6c755952
commit e856eae921
1 changed files with 60 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/llada2.py
+60 -26
View File
@@ -96,6 +96,15 @@ logger = logging.getLogger(__name__)
_is_cuda = is_cuda()
_is_npu = is_npu()
split_qkv_rmsnorm_rope_pos_cache_half_npu = None
if _is_npu:
try:
from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope_pos_cache_half_npu import (
split_qkv_rmsnorm_rope_pos_cache_half_npu,
)
except (ImportError, OSError):
pass
class LLaDA2MoeMLP(nn.Module):
def __init__(
@@ -521,34 +530,59 @@ class LLaDA2MoeAttention(nn.Module):
if hidden_states.shape[0] == 0:
return hidden_states
qkv, _ = self.query_key_value(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
if self.use_qk_norm:
q, k = apply_qk_norm(
q=q,
k=k,
q_norm=self.query_layernorm,
k_norm=self.key_layernorm,
head_dim=self.head_dim,
alt_stream=self.alt_stream,
if _is_npu and split_qkv_rmsnorm_rope_pos_cache_half_npu is not None:
q, k, v = split_qkv_rmsnorm_rope_pos_cache_half_npu(
qkv,
positions,
self.rotary_emb.cos_sin_cache,
self.q_size,
self.kv_size,
self.head_dim,
eps=self.query_layernorm.variance_epsilon if self.use_qk_norm else None,
q_weight=self.query_layernorm.weight if self.use_qk_norm else None,
k_weight=self.key_layernorm.weight if self.use_qk_norm else None,
q_bias=(
getattr(self.query_layernorm, "bias", None)
if self.use_qk_norm
else None
),
k_bias=(
getattr(self.key_layernorm, "bias", None)
if self.use_qk_norm
else None
),
rope_dim=self.rotary_dim,
)
can_fuse_set_kv = (
self.head_dim == self.rotary_emb.rotary_dim
and enable_fused_set_kv_buffer(forward_batch)
)
q, k = self.rotary_emb(
positions,
q,
k,
fused_set_kv_buffer_arg=(
create_fused_set_kv_buffer_arg(
value=v,
layer=self.attn,
forward_batch=forward_batch,
can_fuse_set_kv = False
else:
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
if self.use_qk_norm:
q, k = apply_qk_norm(
q=q,
k=k,
q_norm=self.query_layernorm,
k_norm=self.key_layernorm,
head_dim=self.head_dim,
alt_stream=self.alt_stream,
)
if can_fuse_set_kv
else None
),
)
can_fuse_set_kv = (
self.head_dim == self.rotary_emb.rotary_dim
and enable_fused_set_kv_buffer(forward_batch)
)
q, k = self.rotary_emb(
positions,
q,
k,
fused_set_kv_buffer_arg=(
create_fused_set_kv_buffer_arg(
value=v,
layer=self.attn,
forward_batch=forward_batch,
)
if can_fuse_set_kv
else None
),
)
context_layer = self.attn(
q,
k,