diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 5fa3f64d8..db97e2958 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -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,