From 4f25a48d7a4718e6aac7e29d7363e88d8719674e Mon Sep 17 00:00:00 2001
From: fy <hit_fuyong@163.com>
Date: Tue, 24 Feb 2026 14:11:20 +0800
Subject: [PATCH] support xverse_moe on npu

Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
---
 .../npu/quantization/fused_moe_method_npu.py  | 79 +++++++++++++++++++
 python/sglang/srt/models/xverse_moe.py        |  8 +-
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
index b3bd7c215..fbad373c4 100644
--- a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
+++ b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
@@ -139,6 +139,85 @@ def npu_fused_moe_without_routing_weights_bf16(
     return hidden_states
 
 
+def fused_moe_npu(
+    x,
+    w1,
+    w2,
+    topk_output,
+    moe_runner_config,
+):
+    # TODO: reuse the codes of UnquantizedFusedMoEMethod-forward_npu
+    topk_weights, topk_ids, _ = topk_output
+    original_dtype = x.dtype
+    num_tokens = x.shape[0]
+    topk_weights = topk_weights.to(x.dtype)
+    topk_ids = topk_ids.to(torch.int32)
+    num_experts = w1.shape[0]
+    top_k = topk_weights.shape[-1]
+    row_idx_len = num_tokens * top_k
+    row_idx = (
+        torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device)
+        .view(top_k, -1)
+        .permute(1, 0)
+        .contiguous()
+    )
+
+    hidden_states, expanded_row_idx, expanded_expert_idx = (
+        torch.ops.npu.npu_moe_init_routing(
+            x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
+        )
+    )
+
+    expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
+        expanded_expert_idx, num_experts
+    )
+
+    expert_tokens = expert_tokens.to(torch.int64)
+
+    # gmm1: gate_up_proj
+    hidden_states = torch.ops.npu.npu_grouped_matmul(
+        x=[hidden_states],
+        weight=[w1.permute(0, 2, 1)],
+        bias=None,
+        split_item=2,
+        group_list_type=0,
+        group_type=0,
+        group_list=expert_tokens,
+        output_dtype=original_dtype,
+    )[0]
+
+    # act_fn:
+    if moe_runner_config.activation == "silu":
+        hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
+    else:
+        from sglang.srt.layers.activation import GeluAndMul
+
+        hidden_states = GeluAndMul()(hidden_states)
+
+    # gmm2: down_proj
+    hidden_states = torch.ops.npu.npu_grouped_matmul(
+        x=[hidden_states],
+        weight=[w2.permute(0, 2, 1)],
+        bias=None,
+        split_item=2,
+        group_list_type=0,
+        group_type=0,
+        group_list=expert_tokens,
+        output_dtype=original_dtype,
+    )[0]
+
+    final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
+        hidden_states,
+        skip1=None,
+        skip2=None,
+        bias=None,
+        scales=topk_weights,
+        expanded_src_to_dst_row=expanded_row_idx,
+        export_for_source_row=topk_ids,
+    )
+    return final_hidden_states
+
+
 class _NPUFusedMoEMethodBase(FusedMoEMethodBase):
 
     def __init__(
diff --git a/python/sglang/srt/models/xverse_moe.py b/python/sglang/srt/models/xverse_moe.py
index 6067acec6..e4489055f 100644
--- a/python/sglang/srt/models/xverse_moe.py
+++ b/python/sglang/srt/models/xverse_moe.py
@@ -24,6 +24,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.activation import SiluAndMul
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
@@ -45,7 +48,7 @@ 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
 
 
 class XverseMLP(nn.Module):
@@ -147,6 +150,7 @@ class XverseMoE(nn.Module):
                 reduce_results=False,
                 prefix=add_prefix("shared_experts", prefix),
             )
+        self.fused_moe_method = fused_moe if not is_npu() else fused_moe_npu
 
     def pack_params(self):
         w1 = []
@@ -175,7 +179,7 @@ class XverseMoE(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.w1,
             self.w2,