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

[NPU] solve accuracy problem for stablelm-2-1-6b for npu (#17470)

Browse Source
This commit is contained in:
McZyWu
2026-01-24 08:27:38 +08:00
committed by GitHub
parent 8a5ed2434f
commit b4a611fb33
2 changed files with 51 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
python/sglang/srt/models/stablelm.py
View File

@@ -42,7 +42,9 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
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()
class StablelmMLP(nn.Module):
@@ -137,12 +139,21 @@ class StablelmAttention(nn.Module):
quant_config=quant_config,
prefix=add_prefix("o_proj", prefix),
)
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.rotary_ndims,
max_position=self.config.max_position_embeddings,
base=self.config.rope_theta,
)
if not _is_npu:
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.rotary_ndims,
max_position=self.config.max_position_embeddings,
base=self.config.rope_theta,
)
else:
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.rotary_ndims,
max_position=self.config.max_position_embeddings,
base=self.config.rope_theta,
dtype=torch.float32,
)
self.attn = RadixAttention(
self.num_heads,
self.head_dim,
@@ -161,7 +172,12 @@ class StablelmAttention(nn.Module):
) -> torch.Tensor:
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)
output, _ = self.o_proj(attn_output)
return output