diff --git a/python/sglang/srt/models/dbrx.py b/python/sglang/srt/models/dbrx.py index 12d67d5d8..e37827565 100644 --- a/python/sglang/srt/models/dbrx.py +++ b/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, diff --git a/test/registered/ascend/llm_models/test_ascend_dbrx_instruct.py b/test/registered/ascend/llm_models/test_ascend_dbrx_instruct.py new file mode 100644 index 000000000..189e9fe90 --- /dev/null +++ b/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()