chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU]grok2 model support (#17119)

Browse Source 
Co-authored-by: cy <chenyang08056032@163.com>
This commit is contained in:
KnightLTC
2026-03-03 10:24:10 +08:00
committed by GitHub
parent 8b4c387aa2
commit 1041f240c0
5 changed files with 172 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
View File

@@ -831,6 +831,7 @@ class AscendAttnBackend(AttentionBackend):
layer.qk_head_dim <= 128
and causal
and forward_batch.encoder_lens is None
and layer.logit_cap == 0
and not getattr(self, "use_native_sdpa", False)
):
if not self.use_alibi:
@@ -896,6 +897,8 @@ class AscendAttnBackend(AttentionBackend):
scaling=layer.scaling,
enable_gqa=use_gqa,
causal=causal,
logit_cap=layer.logit_cap,
logit_capping_method=layer.logit_capping_method,
)
attn_output = attn_output.view(
-1, layer.tp_q_head_num * layer.v_head_dim
@@ -1022,7 +1025,7 @@ class AscendAttnBackend(AttentionBackend):
layer.layer_id
)
kv_cache = torch.cat([k_cache, v_cache], dim=-1)
attn_output = self.native_attn._run_sdpa_forward_extend(
attn_output = self.native_attn.run_sdpa_forward_extend(
q,
attn_output,
kv_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
@@ -1525,7 +1528,7 @@ class AscendAttnBackend(AttentionBackend):
)
# there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens
# forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues
elif forward_batch.encoder_lens is None:
elif forward_batch.encoder_lens is None and layer.logit_cap == 0:
query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
num_tokens = query.shape[0]
if not self.use_alibi:
@@ -1585,6 +1588,8 @@ class AscendAttnBackend(AttentionBackend):
scaling=layer.scaling,
enable_gqa=use_gqa,
causal=False,
logit_cap=layer.logit_cap,
logit_capping_method=layer.logit_capping_method,
)
return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
else:

125
python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py
View File

@@ -1,5 +1,7 @@
from __future__ import annotations
import math
import torch
from torch.nn.functional import scaled_dot_product_attention
@@ -8,6 +10,49 @@ class AscendTorchNativeAttnBackend:
def __init__(self):
pass
def scaled_dot_product_attention_with_softcapping(
self,
query,
key,
value,
attn_mask=None,
is_causal=False,
scale=None,
enable_gqa=False,
logit_cap=0.0,
logit_capping_method="tanh",
) -> torch.Tensor:
L, S = query.size(-2), key.size(-2)
scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
if is_causal:
assert attn_mask is None
temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(
diagonal=0
)
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
attn_bias.to(query.dtype)
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
else:
attn_bias = attn_mask + attn_bias
if enable_gqa:
key = key.repeat_interleave(query.size(-3) // key.size(-3), -3)
value = value.repeat_interleave(query.size(-3) // value.size(-3), -3)
attn_weight = query @ key.transpose(-2, -1) * scale_factor
if logit_cap > 0:
if logit_capping_method == "tanh":
attn_weight = logit_cap * torch.tanh(attn_weight / logit_cap)
attn_weight += attn_bias
attn_weight = torch.softmax(attn_weight, dim=-1)
return attn_weight @ value
def run_sdpa_forward_extend(
self,
query: torch.Tensor,
@@ -24,6 +69,8 @@ class AscendTorchNativeAttnBackend:
scaling=None,
enable_gqa=False,
causal=False,
logit_cap: float = 0.0,
logit_capping_method: str = "tanh",
):
"""Run the extend forward by using torch native sdpa op.
@@ -94,18 +141,34 @@ class AscendTorchNativeAttnBackend:
per_req_key = per_req_key.to(per_req_query.dtype)
per_req_value = per_req_value.to(per_req_query.dtype)
per_req_out_redudant = (
scaled_dot_product_attention(
per_req_query_redudant.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
if logit_cap > 0:
per_req_out_redudant = (
self.scaled_dot_product_attention_with_softcapping(
per_req_query_redudant.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
logit_cap=logit_cap,
logit_capping_method=logit_capping_method,
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
else:
per_req_out_redudant = (
scaled_dot_product_attention(
per_req_query_redudant.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :]
start_q, start_kv = end_q, end_kv
return output
@@ -124,6 +187,8 @@ class AscendTorchNativeAttnBackend:
scaling=None,
enable_gqa=False,
causal=False,
logit_cap: float = 0.0,
logit_capping_method: str = "tanh",
):
"""Run the decode forward by using torch native sdpa op.
@@ -180,18 +245,34 @@ class AscendTorchNativeAttnBackend:
per_req_key = per_req_key.to(per_req_query.dtype)
per_req_value = per_req_value.to(per_req_query.dtype)
per_req_out = (
scaled_dot_product_attention(
per_req_query.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
if logit_cap > 0:
per_req_out = (
self.scaled_dot_product_attention_with_softcapping(
per_req_query.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
logit_cap=logit_cap,
logit_capping_method=logit_capping_method,
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
else:
per_req_out = (
scaled_dot_product_attention(
per_req_query.unsqueeze(0),
per_req_key.unsqueeze(0),
per_req_value.unsqueeze(0),
enable_gqa=enable_gqa,
scale=scaling,
is_causal=causal,
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
.squeeze(0)
.movedim(query.dim() - 2, 0)
)
output[start_q:end_q, :, :] = per_req_out
start_q, start_kv = end_q, end_kv

34
python/sglang/srt/models/grok.py
View File

@@ -18,6 +18,7 @@ import math
from typing import Iterable, Optional, Tuple
import torch
import torch.nn.functional as F
from torch import nn
from transformers import PretrainedConfig
@@ -59,7 +60,9 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.loader import DefaultModelLoader
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.utils import add_prefix
from sglang.srt.utils import add_prefix, is_npu
_is_npu = is_npu()
logger = logging.getLogger(__name__)
@@ -143,7 +146,7 @@ class Grok1MoE(nn.Module):
top_k=top_k,
renormalize=False,
layer_id=layer_id,
custom_routing_function=custom_routing_function,
custom_routing_function=None if _is_npu else custom_routing_function,
)
self.experts = FusedMoE(
@@ -162,8 +165,21 @@ class Grok1MoE(nn.Module):
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
topk_output = self.topk(hidden_states, self.gate.weight)
return self.experts(hidden_states, topk_output)
if not _is_npu:
topk_output = self.topk(hidden_states, self.gate.weight)
return self.experts(hidden_states, topk_output)
else:
orig_shape = hidden_states.shape
hidden_states = hidden_states.view(-1, self.hidden_size)
router_logits, _ = self.gate(hidden_states)
router_logits = self.router_logit_softcapping * F.tanh(
router_logits / self.router_logit_softcapping
)
topk_output = self.topk(hidden_states, router_logits)
final_hidden_states = self.experts(hidden_states, topk_output)
return final_hidden_states.view(orig_shape)
def _yarn_linear_ramp_mask(
@@ -228,6 +244,8 @@ class ScalingRotaryEmbedding(RotaryEmbedding):
self.attn_factor = attn_factor
self.beta_fast = beta_fast
self.beta_slow = beta_slow
if _is_npu:
dtype = torch.float32
# Get n-d magnitude scaling corrected for interpolation
self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor)
super().__init__(
@@ -396,6 +414,7 @@ class Grok1Attention(nn.Module):
max_position=max_position,
base=int(self.rope_theta),
is_neox_style=True,
dtype=torch.float32 if _is_npu else None,
)
pos_encoding_mode = "NONE"
@@ -425,7 +444,12 @@ class Grok1Attention(nn.Module):
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self.rotary_emb(positions, q, k)
if not _is_npu:
q, k = self.rotary_emb(positions, q, k)
else:
odtype = q.dtype
q, k = self.rotary_emb(positions, q.to(torch.float32), k.to(torch.float32))
q, k = q.to(odtype), k.to(odtype)
attn_output = self.attn(q, k, v, forward_batch)