iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[CPU][Zen] Route BF16 MoE inference through zentorch on AMD (#44201)

Browse Source 
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Priyansh Jain
2026-08-10 23:20:05 -07:00
committed by GitHub
co-authored by Andreas Karatzas
parent b76cf758db
commit 2acb055ed7
4 changed files with 316 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
setup.py
+1 -1
View File
@@ -1273,7 +1273,7 @@ setup(
install_requires=get_requirements(),
extras_require={
# AMD Zen CPU optimizations via zentorch
"zen": ["zentorch==2.11.0.0"],
"zen": ["zentorch==2.13.0.0"],
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.3.3"],
tests/kernels/moe/test_zen_cpu_fused_moe.py
+221
View File
@@ -0,0 +1,221 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for ZenTorch CPU fused MoE dispatch and forward."""
from types import SimpleNamespace
import pytest
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
from tests.kernels.moe.test_cpu_fused_moe import (
UNALIGNED_INTERMEDIATE_DIM,
_make_moe_config,
_StubMoELayer,
ref_fused_moe,
)
from vllm.model_executor.kernels.linear.zentorch_utils import (
has_zentorch_op,
is_zentorch_moe_config_supported,
)
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
biased_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
CPUUnquantizedExperts,
select_experts,
)
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
if not current_platform.is_cpu() or not current_platform.is_zen_cpu():
pytest.skip("skipping non-Zen CPU tests", allow_module_level=True)
if not has_zentorch_op(["zentorch_fused_moe"]):
pytest.skip(
"skipping tests: zentorch_fused_moe op not available",
allow_module_level=True,
)
ZENTORCH_ACT = [
MoEActivation.SILU,
MoEActivation.GELU,
MoEActivation.GELU_TANH,
MoEActivation.SWIGLUOAI,
]
USE_BIAS = [False, True]
def _make_layer_and_config(
act: MoEActivation,
*,
hidden_size: int = 128,
intermediate_size: int = 128,
expert_num: int = 8,
topk_num: int = 4,
use_bias: bool = False,
dtype: torch.dtype = torch.bfloat16,
) -> tuple[_StubMoELayer, FusedMoEConfig]:
up_dim = 2 * intermediate_size
w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
0.5 * hidden_size**0.5
)
w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
0.5 * intermediate_size**0.5
)
w13_bias = None
w2_bias = None
if use_bias:
w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
0.5 * hidden_size**0.5
)
layer = _StubMoELayer(
w13,
w2,
act,
w13_bias,
w2_bias,
)
moe_config = _make_moe_config(
expert_num,
hidden_size,
intermediate_size,
topk_num,
dtype,
act,
use_bias,
)
return layer, moe_config
@pytest.mark.parametrize("act", ZENTORCH_ACT)
def test_zen_cpu_fused_moe_dispatches_to_zentorch(act: MoEActivation):
"""When zentorch MoE is supported, experts select the zentorch path."""
layer, moe_config = _make_layer_and_config(act)
w13_before = layer.w13_weight.detach().clone()
w2_before = layer.w2_weight.detach().clone()
experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
experts.process_weights_after_loading(layer)
assert experts._use_zentorch is True
torch.testing.assert_close(layer.w13_weight, w13_before)
torch.testing.assert_close(layer.w2_weight, w2_before)
def test_zen_cpu_fused_moe_config_supported_unaligned_intermediate():
"""Zentorch accepts shapes the grouped-gemm kernel rejects."""
act = MoEActivation.SILU
moe_config = _make_moe_config(
expert_num=8,
hidden_size=128,
intermediate_size=UNALIGNED_INTERMEDIATE_DIM,
topk_num=4,
dtype=torch.bfloat16,
act=act,
)
assert is_zentorch_moe_config_supported(moe_config)
supported, reason = CPUUnquantizedExperts.is_supported_config(
CPUUnquantizedExperts,
moe_config,
None,
None,
mk.FusedMoEActivationFormat.Standard,
)
assert supported, reason
@pytest.mark.parametrize("act", ZENTORCH_ACT)
@pytest.mark.parametrize("use_bias", USE_BIAS)
def test_zen_cpu_fused_moe_forward(
default_vllm_config,
act: MoEActivation,
use_bias: bool,
):
"""The zentorch forward runs end-to-end and matches the reference."""
set_random_seed(0)
expert_num = 8
hidden_size = 128
batch_size = 64
dtype = torch.bfloat16
topk_num = expert_num // 2
layer, moe_config = _make_layer_and_config(
act,
expert_num=expert_num,
topk_num=topk_num,
use_bias=use_bias,
)
input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
0.5 * hidden_size**0.5
)
router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
quant_config = (
biased_moe_quant_config(layer.w13_bias, layer.w2_bias)
if use_bias
else FusedMoEQuantConfig.make()
)
experts = CPUUnquantizedExperts(moe_config, quant_config)
experts.process_weights_after_loading(layer)
assert experts._use_zentorch is True
topk_weight, topk_ids = select_experts(
hidden_states=input,
router_logits=router_logits,
top_k=topk_num,
use_grouped_topk=False,
renormalize=False,
)
output = experts.apply(
hidden_states=input,
w1=layer.w13_weight,
w2=layer.w2_weight,
router_logits=router_logits,
activation=act,
global_num_experts=expert_num,
expert_map=None,
a1q_scale=None,
apply_router_weight_on_input=False,
)
ref_output = ref_fused_moe(
input,
layer.w13_weight,
layer.w2_weight,
getattr(layer, "w13_bias", None),
getattr(layer, "w2_bias", None),
topk_weight,
topk_ids,
act,
)
atol, rtol = get_default_atol(output), get_default_rtol(output)
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
def test_zen_cpu_fused_moe_skips_without_activation():
"""Layers without an activation attribute fall back from zentorch."""
layer = SimpleNamespace()
moe_config = _make_moe_config(
expert_num=8,
hidden_size=128,
intermediate_size=128,
topk_num=4,
dtype=torch.bfloat16,
act=MoEActivation.SILU,
)
experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
experts.process_weights_after_loading(layer)
assert experts._use_zentorch is False
vllm/model_executor/kernels/linear/zentorch_utils.py
+63 -1
View File
@@ -4,11 +4,32 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
import torch
from vllm.platforms import current_platform
__all__ = ["has_zentorch_op"]
if TYPE_CHECKING:
from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
__all__ = [
"has_zentorch_op",
"is_zentorch_moe_config_supported",
"is_zentorch_moe_supported",
"_ZENTORCH_MOE_ACTIVATIONS",
]
_ZENTORCH_MOE_ACTIVATIONS = frozenset({"gelu", "gelu_tanh", "silu", "swigluoai"})
def _moe_activation_to_str(activation: object) -> str:
"""Normalize activation to a lowercase string (enum-safe)."""
if hasattr(activation, "value"):
return str(activation.value).lower()
return str(activation).lower()
def has_zentorch_op(op_names: list[str]) -> bool:
@@ -21,3 +42,44 @@ def has_zentorch_op(op_names: list[str]) -> bool:
if ns is None:
return False
return all(hasattr(ns, op_name) for op_name in op_names)
def is_zentorch_moe_config_supported(moe_config: FusedMoEConfig) -> bool:
if not has_zentorch_op(["zentorch_fused_moe"]):
return False
if not moe_config.is_act_and_mul:
return False
act = _moe_activation_to_str(moe_config.activation)
return act in _ZENTORCH_MOE_ACTIVATIONS
def is_zentorch_moe_supported(layer: torch.nn.Module) -> bool:
if not has_zentorch_op(["zentorch_fused_moe"]):
logging.debug(
"Skipping zentorch fused-MoE: not a Zen CPU or "
"zentorch not loaded; using default MoE."
)
return False
moe_config = getattr(layer, "moe_config", None)
if moe_config is not None and not moe_config.is_act_and_mul:
logging.debug(
"Skipping zentorch fused-MoE: layer is not a gated "
"(act-and-mul, e.g. SwiGLU) MLP, the only structure supported."
)
return False
activation = getattr(layer, "activation", None)
if activation is None:
logging.debug(
"Skipping zentorch fused-MoE: layer has no 'activation' "
"attribute, so the activation can't be verified."
)
return False
act = str(getattr(activation, "value", activation)).lower()
if act not in _ZENTORCH_MOE_ACTIVATIONS:
logging.debug(
"Skipping zentorch fused-MoE: activation %r unsupported (supported: %s).",
act,
_ZENTORCH_MOE_ACTIVATIONS,
)
return False
return True
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+31 -1
View File
@@ -20,6 +20,10 @@ from vllm._custom_ops import (
cpu_prepack_moe_weight_int8,
fused_experts_cpu,
)
from vllm.model_executor.kernels.linear.zentorch_utils import (
is_zentorch_moe_config_supported,
is_zentorch_moe_supported,
)
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
@@ -214,6 +218,7 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
self.renormalize = False
self.scoring_func = "softmax"
self.custom_routing_function: Callable | None = None
self._use_zentorch = False
@property
def expects_unquantized_inputs(self) -> bool:
@@ -240,6 +245,8 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
)
if not supported:
return supported, reason
if is_zentorch_moe_config_supported(moe_config):
return True, None
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
return cpu_cls._supports_grouped_gemm(moe_config)
@@ -287,11 +294,14 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
return True
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self._pad_moe_intermediate(layer)
self.use_grouped_topk = layer.use_grouped_topk
self.renormalize = layer.renormalize
self.scoring_func = layer.scoring_func
self.custom_routing_function = layer.custom_routing_function
self._use_zentorch = is_zentorch_moe_supported(layer)
if self._use_zentorch:
return
self._pad_moe_intermediate(layer)
replace_parameter(
layer, "w13_weight", cpu_prepack_moe_weight(layer.w13_weight, self.isa)
)
@@ -388,6 +398,22 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
)
hidden_states.mul_(topk_weights.to(hidden_states.dtype))
if self._use_zentorch:
output = torch.empty_like(hidden_states)
torch.ops.zentorch.zentorch_fused_moe(
output,
hidden_states,
w1,
w2,
self.w1_bias,
self.w2_bias,
topk_weights,
topk_ids,
apply_router_weight_on_input,
str(activation.value).lower(),
)
return output
return cpu_fused_moe(
hidden_states,
w1,
@@ -430,6 +456,8 @@ class X86CPUUnquantizedExperts(CPUUnquantizedExperts):
)
if not supported:
return supported, reason
if is_zentorch_moe_config_supported(moe_config):
return True, None
if moe_config.in_dtype != torch.bfloat16:
return False, "kernel requires bfloat16 activations"
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
@@ -464,6 +492,8 @@ class ArmCPUUnquantizedExperts(CPUUnquantizedExperts):
)
if not supported:
return supported, reason
if is_zentorch_moe_config_supported(moe_config):
return True, None
if moe_config.in_dtype != torch.bfloat16:
return False, "kernel requires bfloat16 activations"
cpu_cls = cast(type[CPUUnquantizedExperts], cls)