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

[CPU][Llama4] Fix Llama4 MoE inputs with "apply_router_weight_on_input" (#7889)

Browse Source
This commit is contained in:
jianan-gu
2025-07-18 12:43:25 +08:00
committed by GitHub
parent 8aa5ae6b04
commit 48c1fa7bb6
5 changed files with 35 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -344,9 +344,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
) -> torch.Tensor:
assert activation == "silu", f"activation = {activation} is not supported."
if use_intel_amx_backend(layer) and not apply_router_weight_on_input:
if use_intel_amx_backend(layer):
from sglang.srt.layers.moe.topk import select_experts
from sglang.srt.layers.moe.topk import (
select_experts,
apply_topk_weights_cpu,
)
topk_weights, topk_ids = select_experts(
hidden_states=x,
@@ -361,8 +364,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
correction_bias=correction_bias,
routed_scaling_factor=routed_scaling_factor,
)
x, topk_weights = apply_topk_weights_cpu(
apply_router_weight_on_input, topk_weights, x
)
# TODO: support apply_router_weight_on_input in the fused_experts_cpu kernel
return torch.ops.sgl_kernel.fused_experts_cpu(
x,
layer.w13_weight,