diff --git a/python/sglang/srt/models/stablelm.py b/python/sglang/srt/models/stablelm.py index 1566893ee..2adcfe92f 100644 --- a/python/sglang/srt/models/stablelm.py +++ b/python/sglang/srt/models/stablelm.py @@ -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 diff --git a/test/registered/ascend/llm_models/test_ascend_stablelm-2-1_6b.py b/test/registered/ascend/llm_models/test_ascend_stablelm-2-1_6b.py new file mode 100644 index 000000000..9261bcf42 --- /dev/null +++ b/test/registered/ascend/llm_models/test_ascend_stablelm-2-1_6b.py @@ -0,0 +1,27 @@ +import unittest + +from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.test_utils import CustomTestCase + +register_npu_ci(est_time=400, suite="nightly-1-npu-a3", nightly=True) + + +class TestStablelm(GSM8KAscendMixin, CustomTestCase): + model = "/root/.cache/modelscope/hub/models/stabilityai/stablelm-2-1_6b" + accuracy = 0.195 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.8", + "--attention-backend", + "ascend", + "--disable-cuda-graph", + "--tp-size", + 1, + "--enable-torch-compile", + ] + + +if __name__ == "__main__": + unittest.main()