diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index 75e34e6766a..9e99a274fb3 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 1a0acff058c..a957aca9eca 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -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"