[CPU][Zen] Route BF16 MoE inference through zentorch on AMD (#44201)
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Priyansh Jain committed
Aug 11, 2026 at 11:50 UTC
2acb055ed768f07cdf22e718d41b3b3252719dc2
4 files changed
+316
-3
setup.py
+1
-1
@@ -1273,7 +1273,7 @@ setup(
1273
install_requires=get_requirements(),
1274
extras_require={
1275
# AMD Zen CPU optimizations via zentorch
1276
- "zen": ["zentorch==2.11.0.0"],
1276
+ "zen": ["zentorch==2.13.0.0"],
1277
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
1278
"tensorizer": ["tensorizer==2.10.1"],
1279
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
tests/kernels/moe/test_zen_cpu_fused_moe.py
new
+221
@@ -0,0 +1,221 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for ZenTorch CPU fused MoE dispatch and forward."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from tests.kernels.allclose_default import get_default_atol, get_default_rtol
12
+from tests.kernels.moe.test_cpu_fused_moe import (
13
+ UNALIGNED_INTERMEDIATE_DIM,
14
+ _make_moe_config,
15
+ _StubMoELayer,
16
+ ref_fused_moe,
17
+)
18
+from vllm.model_executor.kernels.linear.zentorch_utils import (
19
+ has_zentorch_op,
20
+ is_zentorch_moe_config_supported,
21
+)
22
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
23
+from vllm.model_executor.layers.fused_moe.config import (
24
+ FusedMoEConfig,
25
+ FusedMoEQuantConfig,
26
+ biased_moe_quant_config,
27
+)
28
+from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
29
+ CPUUnquantizedExperts,
30
+ select_experts,
31
+)
32
+from vllm.platforms import current_platform
33
+from vllm.utils.torch_utils import set_random_seed
34
+
35
+if not current_platform.is_cpu() or not current_platform.is_zen_cpu():
36
+ pytest.skip("skipping non-Zen CPU tests", allow_module_level=True)
37
+
38
+if not has_zentorch_op(["zentorch_fused_moe"]):
39
+ pytest.skip(
40
+ "skipping tests: zentorch_fused_moe op not available",
41
+ allow_module_level=True,
42
+ )
43
+
44
+ZENTORCH_ACT = [
45
+ MoEActivation.SILU,
46
+ MoEActivation.GELU,
47
+ MoEActivation.GELU_TANH,
48
+ MoEActivation.SWIGLUOAI,
49
+]
50
+USE_BIAS = [False, True]
51
+
52
+
53
+def _make_layer_and_config(
54
+ act: MoEActivation,
55
+ *,
56
+ hidden_size: int = 128,
57
+ intermediate_size: int = 128,
58
+ expert_num: int = 8,
59
+ topk_num: int = 4,
60
+ use_bias: bool = False,
61
+ dtype: torch.dtype = torch.bfloat16,
62
+) -> tuple[_StubMoELayer, FusedMoEConfig]:
63
+ up_dim = 2 * intermediate_size
64
+ w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
65
+ 0.5 * hidden_size**0.5
66
+ )
67
+ w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
68
+ 0.5 * intermediate_size**0.5
69
+ )
70
+ w13_bias = None
71
+ w2_bias = None
72
+ if use_bias:
73
+ w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
74
+ w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
75
+ 0.5 * hidden_size**0.5
76
+ )
77
+
78
+ layer = _StubMoELayer(
79
+ w13,
80
+ w2,
81
+ act,
82
+ w13_bias,
83
+ w2_bias,
84
+ )
85
+ moe_config = _make_moe_config(
86
+ expert_num,
87
+ hidden_size,
88
+ intermediate_size,
89
+ topk_num,
90
+ dtype,
91
+ act,
92
+ use_bias,
93
+ )
94
+ return layer, moe_config
95
+
96
+
97
+@pytest.mark.parametrize("act", ZENTORCH_ACT)
98
+def test_zen_cpu_fused_moe_dispatches_to_zentorch(act: MoEActivation):
99
+ """When zentorch MoE is supported, experts select the zentorch path."""
100
+ layer, moe_config = _make_layer_and_config(act)
101
+ w13_before = layer.w13_weight.detach().clone()
102
+ w2_before = layer.w2_weight.detach().clone()
103
+
104
+ experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
105
+ experts.process_weights_after_loading(layer)
106
+
107
+ assert experts._use_zentorch is True
108
+ torch.testing.assert_close(layer.w13_weight, w13_before)
109
+ torch.testing.assert_close(layer.w2_weight, w2_before)
110
+
111
+
112
+def test_zen_cpu_fused_moe_config_supported_unaligned_intermediate():
113
+ """Zentorch accepts shapes the grouped-gemm kernel rejects."""
114
+ act = MoEActivation.SILU
115
+ moe_config = _make_moe_config(
116
+ expert_num=8,
117
+ hidden_size=128,
118
+ intermediate_size=UNALIGNED_INTERMEDIATE_DIM,
119
+ topk_num=4,
120
+ dtype=torch.bfloat16,
121
+ act=act,
122
+ )
123
+ assert is_zentorch_moe_config_supported(moe_config)
124
+
125
+ supported, reason = CPUUnquantizedExperts.is_supported_config(
126
+ CPUUnquantizedExperts,
127
+ moe_config,
128
+ None,
129
+ None,
130
+ mk.FusedMoEActivationFormat.Standard,
131
+ )
132
+ assert supported, reason
133
+
134
+
135
+@pytest.mark.parametrize("act", ZENTORCH_ACT)
136
+@pytest.mark.parametrize("use_bias", USE_BIAS)
137
+def test_zen_cpu_fused_moe_forward(
138
+ default_vllm_config,
139
+ act: MoEActivation,
140
+ use_bias: bool,
141
+):
142
+ """The zentorch forward runs end-to-end and matches the reference."""
143
+ set_random_seed(0)
144
+
145
+ expert_num = 8
146
+ hidden_size = 128
147
+ batch_size = 64
148
+ dtype = torch.bfloat16
149
+ topk_num = expert_num // 2
150
+
151
+ layer, moe_config = _make_layer_and_config(
152
+ act,
153
+ expert_num=expert_num,
154
+ topk_num=topk_num,
155
+ use_bias=use_bias,
156
+ )
157
+
158
+ input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
159
+ 0.5 * hidden_size**0.5
160
+ )
161
+ router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
162
+
163
+ quant_config = (
164
+ biased_moe_quant_config(layer.w13_bias, layer.w2_bias)
165
+ if use_bias
166
+ else FusedMoEQuantConfig.make()
167
+ )
168
+ experts = CPUUnquantizedExperts(moe_config, quant_config)
169
+ experts.process_weights_after_loading(layer)
170
+ assert experts._use_zentorch is True
171
+
172
+ topk_weight, topk_ids = select_experts(
173
+ hidden_states=input,
174
+ router_logits=router_logits,
175
+ top_k=topk_num,
176
+ use_grouped_topk=False,
177
+ renormalize=False,
178
+ )
179
+
180
+ output = experts.apply(
181
+ hidden_states=input,
182
+ w1=layer.w13_weight,
183
+ w2=layer.w2_weight,
184
+ router_logits=router_logits,
185
+ activation=act,
186
+ global_num_experts=expert_num,
187
+ expert_map=None,
188
+ a1q_scale=None,
189
+ apply_router_weight_on_input=False,
190
+ )
191
+
192
+ ref_output = ref_fused_moe(
193
+ input,
194
+ layer.w13_weight,
195
+ layer.w2_weight,
196
+ getattr(layer, "w13_bias", None),
197
+ getattr(layer, "w2_bias", None),
198
+ topk_weight,
199
+ topk_ids,
200
+ act,
201
+ )
202
+
203
+ atol, rtol = get_default_atol(output), get_default_rtol(output)
204
+ torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
205
+
206
+
207
+def test_zen_cpu_fused_moe_skips_without_activation():
208
+ """Layers without an activation attribute fall back from zentorch."""
209
+ layer = SimpleNamespace()
210
+ moe_config = _make_moe_config(
211
+ expert_num=8,
212
+ hidden_size=128,
213
+ intermediate_size=128,
214
+ topk_num=4,
215
+ dtype=torch.bfloat16,
216
+ act=MoEActivation.SILU,
217
+ )
218
+ experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
219
+ experts.process_weights_after_loading(layer)
220
+
221
+ assert experts._use_zentorch is False
vllm/model_executor/kernels/linear/zentorch_utils.py
+63
-1
@@ -4,11 +4,32 @@
4
5
from __future__ import annotations
6
7
+import logging
8
+from typing import TYPE_CHECKING
9
+
10
import torch
11
12
from vllm.platforms import current_platform
13
11
-__all__ = ["has_zentorch_op"]
14
+if TYPE_CHECKING:
15
+ from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
16
+
17
+__all__ = [
18
+ "has_zentorch_op",
19
+ "is_zentorch_moe_config_supported",
20
+ "is_zentorch_moe_supported",
21
+ "_ZENTORCH_MOE_ACTIVATIONS",
22
+]
23
+
24
+
25
+_ZENTORCH_MOE_ACTIVATIONS = frozenset({"gelu", "gelu_tanh", "silu", "swigluoai"})
26
+
27
+
28
+def _moe_activation_to_str(activation: object) -> str:
29
+ """Normalize activation to a lowercase string (enum-safe)."""
30
+ if hasattr(activation, "value"):
31
+ return str(activation.value).lower()
32
+ return str(activation).lower()
33
34
35
def has_zentorch_op(op_names: list[str]) -> bool:
@@ -21,3 +42,44 @@ def has_zentorch_op(op_names: list[str]) -> bool:
42
if ns is None:
43
return False
44
return all(hasattr(ns, op_name) for op_name in op_names)
45
+
46
+
47
+def is_zentorch_moe_config_supported(moe_config: FusedMoEConfig) -> bool:
48
+ if not has_zentorch_op(["zentorch_fused_moe"]):
49
+ return False
50
+ if not moe_config.is_act_and_mul:
51
+ return False
52
+ act = _moe_activation_to_str(moe_config.activation)
53
+ return act in _ZENTORCH_MOE_ACTIVATIONS
54
+
55
+
56
+def is_zentorch_moe_supported(layer: torch.nn.Module) -> bool:
57
+ if not has_zentorch_op(["zentorch_fused_moe"]):
58
+ logging.debug(
59
+ "Skipping zentorch fused-MoE: not a Zen CPU or "
60
+ "zentorch not loaded; using default MoE."
61
+ )
62
+ return False
63
+ moe_config = getattr(layer, "moe_config", None)
64
+ if moe_config is not None and not moe_config.is_act_and_mul:
65
+ logging.debug(
66
+ "Skipping zentorch fused-MoE: layer is not a gated "
67
+ "(act-and-mul, e.g. SwiGLU) MLP, the only structure supported."
68
+ )
69
+ return False
70
+ activation = getattr(layer, "activation", None)
71
+ if activation is None:
72
+ logging.debug(
73
+ "Skipping zentorch fused-MoE: layer has no 'activation' "
74
+ "attribute, so the activation can't be verified."
75
+ )
76
+ return False
77
+ act = str(getattr(activation, "value", activation)).lower()
78
+ if act not in _ZENTORCH_MOE_ACTIVATIONS:
79
+ logging.debug(
80
+ "Skipping zentorch fused-MoE: activation %r unsupported (supported: %s).",
81
+ act,
82
+ _ZENTORCH_MOE_ACTIVATIONS,
83
+ )
84
+ return False
85
+ return True
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+31
-1
@@ -20,6 +20,10 @@ from vllm._custom_ops import (
20
cpu_prepack_moe_weight_int8,
21
fused_experts_cpu,
22
)
23
+from vllm.model_executor.kernels.linear.zentorch_utils import (
24
+ is_zentorch_moe_config_supported,
25
+ is_zentorch_moe_supported,
26
+)
27
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
28
from vllm.model_executor.layers.fused_moe.config import (
29
FusedMoEConfig,
@@ -214,6 +218,7 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
218
self.renormalize = False
219
self.scoring_func = "softmax"
220
self.custom_routing_function: Callable | None = None
221
+ self._use_zentorch = False
222
223
@property
224
def expects_unquantized_inputs(self) -> bool:
@@ -240,6 +245,8 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
245
)
246
if not supported:
247
return supported, reason
248
+ if is_zentorch_moe_config_supported(moe_config):
249
+ return True, None
250
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
251
return cpu_cls._supports_grouped_gemm(moe_config)
252
@@ -287,11 +294,14 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
294
return True
295
296
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
290
- self._pad_moe_intermediate(layer)
297
self.use_grouped_topk = layer.use_grouped_topk
298
self.renormalize = layer.renormalize
299
self.scoring_func = layer.scoring_func
300
self.custom_routing_function = layer.custom_routing_function
301
+ self._use_zentorch = is_zentorch_moe_supported(layer)
302
+ if self._use_zentorch:
303
+ return
304
+ self._pad_moe_intermediate(layer)
305
replace_parameter(
306
layer, "w13_weight", cpu_prepack_moe_weight(layer.w13_weight, self.isa)
307
)
@@ -388,6 +398,22 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
398
)
399
hidden_states.mul_(topk_weights.to(hidden_states.dtype))
400
401
+ if self._use_zentorch:
402
+ output = torch.empty_like(hidden_states)
403
+ torch.ops.zentorch.zentorch_fused_moe(
404
+ output,
405
+ hidden_states,
406
+ w1,
407
+ w2,
408
+ self.w1_bias,
409
+ self.w2_bias,
410
+ topk_weights,
411
+ topk_ids,
412
+ apply_router_weight_on_input,
413
+ str(activation.value).lower(),
414
+ )
415
+ return output
416
+
417
return cpu_fused_moe(
418
hidden_states,
419
w1,
@@ -430,6 +456,8 @@ class X86CPUUnquantizedExperts(CPUUnquantizedExperts):
456
)
457
if not supported:
458
return supported, reason
459
+ if is_zentorch_moe_config_supported(moe_config):
460
+ return True, None
461
if moe_config.in_dtype != torch.bfloat16:
462
return False, "kernel requires bfloat16 activations"
463
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
@@ -464,6 +492,8 @@ class ArmCPUUnquantizedExperts(CPUUnquantizedExperts):
492
)
493
if not supported:
494
return supported, reason
495
+ if is_zentorch_moe_config_supported(moe_config):
496
+ return True, None
497
if moe_config.in_dtype != torch.bfloat16:
498
return False, "kernel requires bfloat16 activations"
499
cpu_cls = cast(type[CPUUnquantizedExperts], cls)