[Bugfix][CPU] Zero-pad MoE intermediate size for grouped-gemm TP alignment (#49591)
Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Li, Jiang committed
Jul 27, 2026 at 16:32 UTC
eb290ab673c8f3ff87648cd8e4cd60f50fe7b301
2 files changed
+211
-14
tests/kernels/moe/test_cpu_fused_moe.py
+130
@@ -14,6 +14,7 @@ from vllm._custom_ops import (
14
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15
from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
16
_CPU_MOE_ACT_FN,
17
+ CPUFusedMOE,
18
)
19
from vllm.platforms import CpuArchEnum, current_platform
20
from vllm.utils.torch_utils import set_random_seed
@@ -346,3 +347,132 @@ def test_cpu_fused_moe_int8(
347
)
348
349
torch.testing.assert_close(output, ref_output, atol=2e-2, rtol=2e-2)
350
+
351
+
352
+# moe_intermediate_size not a multiple of 32, e.g. what tensor-parallel
353
+# sharding of moe_intermediate_size=704 at tp=4 produces (704 // 4 == 176).
354
+UNALIGNED_INTERMEDIATE_DIM = 176
355
+
356
+
357
+class _StubMoELayer(torch.nn.Module):
358
+ """Minimal stand-in for the real MoE layer module, exposing just what
359
+ CPUFusedMOE reads/replaces (w13_weight, w2_weight, activation, and
360
+ optionally w13_bias/w2_bias)."""
361
+
362
+ def __init__(
363
+ self,
364
+ w13_weight: torch.Tensor,
365
+ w2_weight: torch.Tensor,
366
+ activation: MoEActivation,
367
+ w13_bias: torch.Tensor | None = None,
368
+ w2_bias: torch.Tensor | None = None,
369
+ ):
370
+ super().__init__()
371
+ self.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
372
+ self.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
373
+ self.activation = activation
374
+ if w13_bias is not None:
375
+ self.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
376
+ if w2_bias is not None:
377
+ self.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
378
+
379
+
380
+@pytest.mark.parametrize("expert_num", EXPERT_NUM)
381
+@pytest.mark.parametrize("hidden_size", HIDDEN_DIM)
382
+@pytest.mark.parametrize("use_bias", USE_BIAS)
383
+@pytest.mark.parametrize("dtype", DTYPE)
384
+@pytest.mark.parametrize(
385
+ "act",
386
+ [MoEActivation.SILU, MoEActivation.GELU, MoEActivation.GELU_TANH],
387
+)
388
+def test_cpu_fused_moe_unaligned_intermediate_size(
389
+ default_vllm_config,
390
+ expert_num: int,
391
+ hidden_size: int,
392
+ use_bias: bool,
393
+ dtype: torch.dtype,
394
+ act: MoEActivation,
395
+):
396
+ """An unaligned per-partition moe_intermediate_size must still hit the
397
+ grouped-gemm fast path via automatic zero-padding, with numerically
398
+ correct output, instead of silently falling back to the much slower
399
+ per-expert torch loop."""
400
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
401
+ pytest.skip("padding is only applied on the x86 AMX/vector kernels")
402
+
403
+ set_random_seed(0)
404
+ batch_size = 64
405
+ intermediate_size = UNALIGNED_INTERMEDIATE_DIM
406
+ topk_num = max(expert_num // 2, 1)
407
+ up_dim = 2 * intermediate_size
408
+
409
+ input = torch.randn((batch_size, hidden_size), dtype=dtype) / (
410
+ 0.5 * hidden_size**0.5
411
+ )
412
+ w13 = torch.randn((expert_num, up_dim, hidden_size), dtype=dtype) / (
413
+ 0.5 * hidden_size**0.5
414
+ )
415
+ w2 = torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype) / (
416
+ 0.5 * intermediate_size**0.5
417
+ )
418
+ router_logits = torch.randn((batch_size, expert_num), dtype=dtype)
419
+ w13_bias = None
420
+ w2_bias = None
421
+ if use_bias:
422
+ w13_bias = torch.randn((expert_num, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
423
+ w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
424
+ 0.5 * hidden_size**0.5
425
+ )
426
+ score = torch.softmax(router_logits, dim=-1, dtype=torch.float32)
427
+ topk_weight, topk_ids = torch.topk(score, topk_num)
428
+ topk_ids = topk_ids.to(torch.int32)
429
+
430
+ ref_output = ref_fused_moe(
431
+ input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act
432
+ )
433
+
434
+ layer = _StubMoELayer(
435
+ w13.clone(),
436
+ w2.clone(),
437
+ act,
438
+ w13_bias.clone() if w13_bias is not None else None,
439
+ w2_bias.clone() if w2_bias is not None else None,
440
+ )
441
+
442
+ cpu_moe = CPUFusedMOE(layer)
443
+ assert cpu_moe.forward_method == cpu_moe.forward_grouped_gemm, (
444
+ "expected the padded intermediate size to hit the grouped-gemm fast path"
445
+ )
446
+
447
+ output = cpu_moe.forward_method(
448
+ layer, input, topk_weight, topk_ids, act, expert_num, False
449
+ )
450
+
451
+ atol, rtol = get_default_atol(output), get_default_rtol(output)
452
+ torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
453
+
454
+
455
+def test_cpu_fused_moe_unaligned_intermediate_size_swigluoai(default_vllm_config):
456
+ """swigluoai's interleaved gate/up layout isn't padded automatically. On
457
+ AMX-capable CPUs this must raise rather than silently falling back to
458
+ the (correct but much slower) per-expert torch loop; elsewhere it should
459
+ fall back exactly as before."""
460
+ set_random_seed(0)
461
+ expert_num = 8
462
+ hidden_size = 128
463
+ intermediate_size = UNALIGNED_INTERMEDIATE_DIM
464
+ dtype = torch.bfloat16
465
+ up_dim = 2 * intermediate_size
466
+
467
+ layer = _StubMoELayer(
468
+ torch.randn((expert_num, up_dim, hidden_size), dtype=dtype),
469
+ torch.randn((expert_num, hidden_size, intermediate_size), dtype=dtype),
470
+ MoEActivation.SWIGLUOAI,
471
+ )
472
+
473
+ if torch.cpu._is_amx_tile_supported():
474
+ with pytest.raises(RuntimeError):
475
+ CPUFusedMOE(layer)
476
+ else:
477
+ cpu_moe = CPUFusedMOE(layer)
478
+ assert cpu_moe.forward_method == cpu_moe.forward_torch
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+81
-14
@@ -20,6 +20,11 @@ from vllm.platforms import CpuArchEnum, current_platform
20
from vllm.utils.torch_utils import direct_register_custom_op
21
22
_CPU_MOE_LAYER_CACHE = {}
23
+# The CPU grouped-gemm MoE kernels (AMX and vector) tile the expert
24
+# intermediate ("N") dimension in blocks of this size and have no tail/
25
+# remainder handling, so a shard is only eligible for the fast path when
26
+# its per-partition intermediate size is a multiple of it.
27
+_MOE_GROUPED_GEMM_N_TILE = 32
28
29
30
def _swigluoai_forward_native(
@@ -230,6 +235,7 @@ class CPUFusedMOE:
235
"""CPU-based fused MoE implementation."""
236
237
def __init__(self, layer: torch.nn.Module) -> None:
238
+ self._pad_moe_intermediate_for_grouped_gemm(layer)
239
use_grouped_gemm, isa = self.check_grouped_gemm(layer)
240
self.isa = isa
241
if use_grouped_gemm:
@@ -284,6 +290,69 @@ class CPUFusedMOE:
290
apply_router_weight_on_input,
291
)
292
293
+ def _pad_moe_intermediate_for_grouped_gemm(self, layer: torch.nn.Module) -> None:
294
+ """Zero-pad the per-partition MoE intermediate dim up to a multiple
295
+ of _MOE_GROUPED_GEMM_N_TILE, so the AMX/vector grouped-gemm kernels
296
+ can be used even when TP sharding (moe_intermediate_size // tp_size)
297
+ lands on an unaligned value (e.g. moe_intermediate_size=704 at tp=4
298
+ -> 176). Only applies to the x86 (AMX/vec) kernels and half-split
299
+ gate/up activations; interleaved layouts (swigluoai) are left
300
+ untouched.
301
+ """
302
+ if not hasattr(torch.ops._C, "prepack_moe_weight"):
303
+ return
304
+ if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
305
+ return
306
+
307
+ intermediate_size = layer.w2_weight.size(2)
308
+ remainder = intermediate_size % _MOE_GROUPED_GEMM_N_TILE
309
+ if remainder == 0:
310
+ return
311
+ if layer.activation == MoEActivation.SWIGLUOAI:
312
+ return
313
+
314
+ pad = _MOE_GROUPED_GEMM_N_TILE - remainder
315
+ padded_size = intermediate_size + pad
316
+ num_experts, _, hidden_size = layer.w13_weight.shape
317
+
318
+ new_w13 = layer.w13_weight.new_zeros(num_experts, 2 * padded_size, hidden_size)
319
+ new_w13[:, :intermediate_size] = layer.w13_weight[:, :intermediate_size]
320
+ new_w13[:, padded_size : padded_size + intermediate_size] = layer.w13_weight[
321
+ :, intermediate_size:
322
+ ]
323
+ replace_parameter(layer, "w13_weight", new_w13)
324
+
325
+ new_w2 = layer.w2_weight.new_zeros(num_experts, hidden_size, padded_size)
326
+ new_w2[:, :, :intermediate_size] = layer.w2_weight
327
+ replace_parameter(layer, "w2_weight", new_w2)
328
+
329
+ if hasattr(layer, "w13_bias"):
330
+ new_bias = layer.w13_bias.new_zeros(num_experts, 2 * padded_size)
331
+ new_bias[:, :intermediate_size] = layer.w13_bias[:, :intermediate_size]
332
+ new_bias[:, padded_size : padded_size + intermediate_size] = layer.w13_bias[
333
+ :, intermediate_size:
334
+ ]
335
+ replace_parameter(layer, "w13_bias", new_bias)
336
+
337
+ def _grouped_gemm_alignment_error(self, layer: torch.nn.Module) -> str:
338
+ # w2's input size is the per-partition MoE intermediate size (the
339
+ # dimension TP-sharding splits), and it's what most commonly breaks
340
+ # alignment, e.g. moe_intermediate_size=704 at tp=4 gives
341
+ # 704 // 4 == 176, which isn't a multiple of 32.
342
+ intermediate_size_per_partition = layer.w2_weight.size(2)
343
+ return (
344
+ "CPU fused-MoE AMX grouped-gemm kernel cannot be used for a "
345
+ f"layer with w13 shape {tuple(layer.w13_weight.shape)} / w2 "
346
+ f"shape {tuple(layer.w2_weight.shape)}: the per-partition MoE "
347
+ f"intermediate size ({intermediate_size_per_partition}) is not "
348
+ f"a multiple of {_MOE_GROUPED_GEMM_N_TILE}, and automatic "
349
+ "zero-padding could not resolve this (typically because the "
350
+ "activation uses an interleaved gate/up layout, e.g. "
351
+ "swigluoai). vLLM refuses to silently fall back to the much "
352
+ "slower per-expert torch loop on AMX-capable CPUs; consider a "
353
+ "different --tensor-parallel-size."
354
+ )
355
+
356
def check_grouped_gemm(
357
self,
358
layer: torch.nn.Module,
@@ -297,20 +366,19 @@ class CPUFusedMOE:
366
w2_input_size = layer.w2_weight.size(2)
367
w2_output_size = layer.w2_weight.size(1)
368
300
- if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
301
- return False, "none"
302
-
369
supports_amx = torch.cpu._is_amx_tile_supported()
304
-
305
- if (
306
- supports_amx
307
- and dtype == torch.bfloat16
308
- and w13_input_size % 32 == 0
309
- and w2_input_size % 32 == 0
310
- ):
311
- return True, "amx"
312
-
370
if supports_amx:
371
+ if (
372
+ dtype == torch.bfloat16
373
+ and w13_output_size % 32 == 0
374
+ and w2_output_size % 32 == 0
375
+ and w13_input_size % 32 == 0
376
+ and w2_input_size % 32 == 0
377
+ ):
378
+ return True, "amx"
379
+ raise RuntimeError(self._grouped_gemm_alignment_error(layer))
380
+
381
+ if not (w13_output_size % 32 == 0 and w2_output_size % 32 == 0):
382
return False, "none"
383
384
supports_neon = current_platform.get_cpu_architecture() == CpuArchEnum.ARM
@@ -321,8 +389,7 @@ class CPUFusedMOE:
389
and w2_input_size % 4 == 0
390
):
391
return True, "neon"
324
- else:
325
- return False, "none"
392
+ return False, "none"
393
394
return True, "vec"
395