chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU]Optimization of forward_npu for UnquantizedFusedMoEMethod (#13158)

Browse Source
This commit is contained in:
Vladimir221
2026-02-25 03:55:20 +03:00
committed by GitHub
parent d7a03c7ebf
commit f1088beb6a
1 changed files with 20 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -543,7 +543,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
import torch_npu
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
@@ -556,35 +555,33 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
topk_ids = topk_ids.to(torch.int32)
num_experts = layer.num_experts
top_k = layer.top_k
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_npu.npu_moe_init_routing(
x, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
w13 = layer.w13_weight
w2 = layer.w2_weight
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(
x,
topk_ids,
active_num=num_tokens * top_k,
expert_num=num_experts,
expert_tokens_num_type=1,
expert_tokens_num_flag=True,
active_expert_range=[0, num_experts],
quant_mode=-1,
)
)
expert_tokens = torch_npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
)
expert_tokens = expert_tokens.to(torch.int64)
w13_bias = [layer.w13_weight_bias] if self.with_bias else None
w2_bias = [layer.w2_weight_bias] if self.with_bias else None
# gmm1: gate_up_proj
hidden_states = torch_npu.npu_grouped_matmul(
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight],
bias=w13_bias,
split_item=2,
group_list_type=0,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
@@ -596,25 +593,25 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
hidden_states = swiglu_oai(layer, hidden_states)
elif self.moe_runner_config.activation == "silu":
hidden_states = torch_npu.npu_swiglu(hidden_states)
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_npu.npu_grouped_matmul(
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight],
bias=w2_bias,
split_item=2,
group_list_type=0,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
final_hidden_states = torch_npu.npu_moe_finalize_routing(
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
@@ -622,6 +619,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
drop_pad_mode=2,
)
return StandardCombineInput(hidden_states=final_hidden_states)