[CPU][Zen] Pass f32 weight scales to the zentorch INT8 MoE (#59434)
Signed-off-by: Ganesh R <Ganesh.R@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>
Ganesh R committed
Oct 1, 2026 at 13:21 UTC
5addde2a0afa1e27b6e0c478ba9b6995724f006d
2 files changed
+11
-8
tests/kernels/moe/test_zen_cpu_int8_moe.py
+3
-2
@@ -224,8 +224,9 @@ def test_zen_int8_dispatch_contract(
224
call = mock_zentorch_fused_moe["fused_moe"]
225
assert call["w13_scales"].shape == (EXPERT_NUM, 2 * INTERMEDIATE_SIZE)
226
assert call["w2_scales"].shape == (EXPERT_NUM, HIDDEN_SIZE)
227
- assert call["w13_scales"].dtype == torch.bfloat16
228
- assert call["w2_scales"].dtype == torch.bfloat16
227
+ # f32 scales keep the layer on zentorch's fast routed MoE executor.
228
+ assert call["w13_scales"].dtype == torch.float32
229
+ assert call["w2_scales"].dtype == torch.float32
230
assert call["act"] == act.value.lower()
231
assert call["skip_weighted"] is False
232
assert call["topk_weights"].dtype == torch.float32
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+8
-6
@@ -1502,15 +1502,17 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
1502
"ZenCPUExpertsInt8 requires per-channel weight scales on the layer."
1503
)
1504
num_experts = self.w1_scale.shape[0]
1505
- self._w1_scale_bf16 = (
1505
+ # ZenDNN's fast routed MoE kernel only accepts f32 expert scales; bf16
1506
+ # silently drops the layer onto the slower per-expert path.
1507
+ self._w1_scale_f32 = (
1508
self.w1_scale.detach()
1507
- .to(torch.bfloat16)
1509
+ .to(torch.float32)
1510
.reshape(num_experts, -1)
1511
.contiguous()
1512
)
1511
- self._w2_scale_bf16 = (
1513
+ self._w2_scale_f32 = (
1514
self.w2_scale.detach()
1513
- .to(torch.bfloat16)
1515
+ .to(torch.float32)
1516
.reshape(num_experts, -1)
1517
.contiguous()
1518
)
@@ -1637,8 +1639,8 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
1639
topk_ids.to(torch.int32).contiguous(),
1640
apply_router_weight_on_input, # skip_weighted
1641
str(activation.value).lower(),
1640
- self._w1_scale_bf16,
1641
- self._w2_scale_bf16,
1642
+ self._w1_scale_f32,
1643
+ self._w2_scale_f32,
1644
)
1645
1646