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

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),
         )
+        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)
+        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

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()