mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-02 03:48:06 +00:00
[Bugfix][CPU] Zero-pad MoE intermediate size for grouped-gemm TP alignment (#49591)
Signed-off-by: jiang1.li <[email protected]> Co-authored-by: Claude Sonnet 5 <[email protected]>
This commit is contained in:
Li, Jiang
co-authored by
Claude Sonnet 5
Claude Sonnet 5
parent
cbc3a87200
commit
eb290ab673
@@ -14,6 +14,7 @@ from vllm._custom_ops import (
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
|
||||
_CPU_MOE_ACT_FN,
|
||||
CPUFusedMOE,
|
||||
)
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
@@ -346,3 +347,132 @@ def test_cpu_fused_moe_int8(
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
# moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel
|
||||
# sharding of moe_intermediate_size=704 at tp=4 produces (704 // 4 == 176).
|
||||
UNALIGNED_INTERMEDIATE_DIM = 176
|
||||
|
||||
|
||||
class _StubMoELayer(torch.nn.Module):
|
||||
"""Minimal stand-in for the real MoE layer module, exposing just what
|
||||
CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and
|
||||
optionally w13_bias/w2_bias)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
w13_weight: torch.Tensor,
|
||||
w2_weight: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
w13_bias: torch.Tensor | None = None,
|
||||
w2_bias: torch.Tensor | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
|
||||
self.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
|
||||
self.activation = activation
|
||||
if w13_bias is not None:
|
||||
self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
|
||||
if w2_bias is not None:
|
||||
self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expert_num", EXPERT_NUM)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
|
||||
@pytest.mark.parametrize("use_bias", USE_BIAS)
|
||||
@pytest.mark.parametrize("dtype", DTYPE)
|
||||
@pytest.mark.parametrize(
|
||||
"act",
|
||||
[MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH],
|
||||
)
|
||||
def test_cpu_fused_moe_unaligned_intermediate_size(
|
||||
default_vllm_config,
|
||||
expert_num: int,
|
||||
hidden_size: int,
|
||||
use_bias: bool,
|
||||
dtype: torch.dtype,
|
||||
act: MoEActivation,
|
||||
):
|
||||
"""An unaligned per-partition moe_intermediate_size must still hit the
|
||||
grouped-gemm fast path via automatic zero-padding, with numerically
|
||||
correct output, instead of silently falling back to the much slower
|
||||
per-expert torch loop."""
|
||||
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
|
||||
pytest.skip("padding is only applied on the x86 AMX/vector kernels")
|
||||
|
||||
set_random_seed(0)
|
||||
batch_size = 64
|
||||
intermediate_size = UNALIGNED_INTERMEDIATE_DIM
|
||||
topk_num = max(expert_num // 2, 1)
|
||||
up_dim = 2 * intermediate_size
|
||||
|
||||
input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
|
||||
0.5 * hidden_size**0.5
|
||||
)
|
||||
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
|
||||
)
|
||||
router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
|
||||
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
|
||||
)
|
||||
score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk_num)
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
ref_output = ref_fused_moe(
|
||||
input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act
|
||||
)
|
||||
|
||||
layer = _StubMoELayer(
|
||||
w13.clone(),
|
||||
w2.clone(),
|
||||
act,
|
||||
w13_bias.clone() if w13_bias is not None else None,
|
||||
w2_bias.clone() if w2_bias is not None else None,
|
||||
)
|
||||
|
||||
cpu_moe = CPUFusedMOE(layer)
|
||||
assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, (
|
||||
"expected the padded intermediate size to hit the grouped-gemm fast path"
|
||||
)
|
||||
|
||||
output = cpu_moe.forward_method(
|
||||
layer, input, topk_weight, topk_ids, act, expert_num, False
|
||||
)
|
||||
|
||||
atol, rtol = get_default_atol(output), get_default_rtol(output)
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config):
|
||||
"""swigluoai's interleaved gate/up layout isn't padded automatically. On
|
||||
AMX-capable CPUs this must raise rather than silently falling back to
|
||||
the (correct but much slower) per-expert torch loop; elsewhere it should
|
||||
fall back exactly as before."""
|
||||
set_random_seed(0)
|
||||
expert_num = 8
|
||||
hidden_size = 128
|
||||
intermediate_size = UNALIGNED_INTERMEDIATE_DIM
|
||||
dtype = torch.bfloat16
|
||||
up_dim = 2 * intermediate_size
|
||||
|
||||
layer = _StubMoELayer(
|
||||
torch.randn((expert_num, up_dim, hidden_size), dtype=dtype),
|
||||
torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype),
|
||||
MoEActivation.SWIGLUOAI,
|
||||
)
|
||||
|
||||
if torch.cpu._is_amx_tile_supported():
|
||||
with pytest.raises(RuntimeError):
|
||||
CPUFusedMOE(layer)
|
||||
else:
|
||||
cpu_moe = CPUFusedMOE(layer)
|
||||
assert cpu_moe.forward_method == cpu_moe.forward_torch
|
||||
|
||||
@@ -20,6 +20,11 @@ from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
_CPU_MOE_LAYER_CACHE = {}
|
||||
# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert
|
||||
# intermediate ("N") dimension in blocks of this size and have no tail/
|
||||
# remainder handling, so a shard is only eligible for the fast path when
|
||||
# its per-partition intermediate size is a multiple of it.
|
||||
_MOE_GROUPED_GEMM_N_TILE = 32
|
||||
|
||||
|
||||
def _swigluoai_forward_native(
|
||||
@@ -230,6 +235,7 @@ class CPUFusedMOE:
|
||||
"""CPU-based fused MoE implementation."""
|
||||
|
||||
def __init__(self, layer: torch.nn.Module) -> None:
|
||||
self._pad_moe_intermediate_for_grouped_gemm(layer)
|
||||
use_grouped_gemm, isa = self.check_grouped_gemm(layer)
|
||||
self.isa = isa
|
||||
if use_grouped_gemm:
|
||||
@@ -284,6 +290,69 @@ class CPUFusedMOE:
|
||||
apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None:
|
||||
"""Zero-pad the per-partition MoE intermediate dim up to a multiple
|
||||
of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels
|
||||
can be used even when TP sharding (moe_intermediate_size // tp_size)
|
||||
lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4
|
||||
-> 176). Only applies to the x86 (AMX/vec) kernels and half-split
|
||||
gate/up activations; interleaved layouts (swigluoai) are left
|
||||
untouched.
|
||||
"""
|
||||
if not hasattr(torch.ops._C, "prepack_moe_weight"):
|
||||
return
|
||||
if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
|
||||
return
|
||||
|
||||
intermediate_size = layer.w2_weight.size(2)
|
||||
remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE
|
||||
if remainder == 0:
|
||||
return
|
||||
if layer.activation == MoEActivation.SWIGLUOAI:
|
||||
return
|
||||
|
||||
pad = _MOE_GROUPED_GEMM_N_TILE - remainder
|
||||
padded_size = intermediate_size + pad
|
||||
num_experts, _, hidden_size = layer.w13_weight.shape
|
||||
|
||||
new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size)
|
||||
new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size]
|
||||
new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[
|
||||
:, intermediate_size:
|
||||
]
|
||||
replace_parameter(layer, "w13_weight", new_w13)
|
||||
|
||||
new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size)
|
||||
new_w2[:, :, :intermediate_size] = layer.w2_weight
|
||||
replace_parameter(layer, "w2_weight", new_w2)
|
||||
|
||||
if hasattr(layer, "w13_bias"):
|
||||
new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size)
|
||||
new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size]
|
||||
new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[
|
||||
:, intermediate_size:
|
||||
]
|
||||
replace_parameter(layer, "w13_bias", new_bias)
|
||||
|
||||
def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str:
|
||||
# w2's input size is the per-partition MoE intermediate size (the
|
||||
# dimension TP-sharding splits), and it's what most commonly breaks
|
||||
# alignment, e.g. moe_intermediate_size=704 at tp=4 gives
|
||||
# 704 // 4 == 176, which isn't a multiple of 32.
|
||||
intermediate_size_per_partition = layer.w2_weight.size(2)
|
||||
return (
|
||||
"CPU fused-MoE AMX grouped-gemm kernel cannot be used for a "
|
||||
f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 "
|
||||
f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE "
|
||||
f"intermediate size ({intermediate_size_per_partition}) is not "
|
||||
f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic "
|
||||
"zero-padding could not resolve this (typically because the "
|
||||
"activation uses an interleaved gate/up layout, e.g. "
|
||||
"swigluoai). vLLM refuses to silently fall back to the much "
|
||||
"slower per-expert torch loop on AMX-capable CPUs; consider a "
|
||||
"different --tensor-parallel-size."
|
||||
)
|
||||
|
||||
def check_grouped_gemm(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -297,20 +366,19 @@ class CPUFusedMOE:
|
||||
w2_input_size = layer.w2_weight.size(2)
|
||||
w2_output_size = layer.w2_weight.size(1)
|
||||
|
||||
if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
|
||||
return False, "none"
|
||||
|
||||
supports_amx = torch.cpu._is_amx_tile_supported()
|
||||
|
||||
if (
|
||||
supports_amx
|
||||
and dtype == torch.bfloat16
|
||||
and w13_input_size % 32 == 0
|
||||
and w2_input_size % 32 == 0
|
||||
):
|
||||
return True, "amx"
|
||||
|
||||
if supports_amx:
|
||||
if (
|
||||
dtype == torch.bfloat16
|
||||
and w13_output_size % 32 == 0
|
||||
and w2_output_size % 32 == 0
|
||||
and w13_input_size % 32 == 0
|
||||
and w2_input_size % 32 == 0
|
||||
):
|
||||
return True, "amx"
|
||||
raise RuntimeError(self._grouped_gemm_alignment_error(layer))
|
||||
|
||||
if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
|
||||
return False, "none"
|
||||
|
||||
supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
|
||||
@@ -321,8 +389,7 @@ class CPUFusedMOE:
|
||||
and w2_input_size % 4 == 0
|
||||
):
|
||||
return True, "neon"
|
||||
else:
|
||||
return False, "none"
|
||||
return False, "none"
|
||||
|
||||
return True, "vec"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user