dbrx instruct npu support (#17121)

Co-authored-by: McZyWu <zhuoyun.wu.23@ucl.ac.uk>

python/sglang/srt/models/dbrx.py

@@ -26,6 +26,9 @@ from sglang.srt.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
+    fused_moe_npu,
+)
 from sglang.srt.layers.linear import (
     QKVParallelLinear,
     ReplicatedLinear,
@@ -48,7 +51,9 @@ from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
 )
-from sglang.srt.utils import add_prefix, set_weight_attrs
+from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs
+
+_is_npu = is_npu()
 
 
 class DbrxRouter(nn.Module):
@@ -142,6 +147,7 @@ class DbrxExperts(nn.Module):
                 "weight_loader": self.weight_loader,
             },
         )
+        self.fused_moe_method = fused_moe if not _is_npu else fused_moe_npu
 
     def weight_loader(
         self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str
@@ -177,7 +183,7 @@ class DbrxExperts(nn.Module):
         # router_logits: (num_tokens, n_experts)
         router_logits = self.router(hidden_states)
         topk_output = self.topk(hidden_states, router_logits)
-        final_hidden_states = fused_moe(
+        final_hidden_states = self.fused_moe_method(
             hidden_states,
             self.ws,
             self.w2s,

test/registered/ascend/llm_models/test_ascend_dbrx_instruct.py

@@ -0,0 +1,26 @@
+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-8-npu-a3", nightly=True)
+
+
+class TestDbrx(GSM8KAscendMixin, CustomTestCase):
+    model = "/root/.cache/modelscope/hub/models/AI-ModelScope/dbrx-instruct"
+    accuracy = 0.735
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.8",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "8",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()