A high-throughput and memory-efficient inference and serving engine for LLMs
[CPU] Support CPU W4A16 INT4 MoE (#43409)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com>
Yuwen Zhou committed
Jun 12, 2026 at 15:12 UTC
0cd9b7af25cd3c47a84e2164392f755415c74fd2
8 files changed
+685
-22
.buildkite/hardware_tests/cpu.yaml
+1
-1
@@ -91,7 +91,7 @@ steps:
91
- tests/quantization/test_cpu_wna16.py
92
commands:
93
- |
94
- bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
94
+ bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
95
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
96
pytest -x -v -s tests/quantization/test_cpu_wna16.py"
97
tests/kernels/moe/test_cpu_quant_fused_moe.py
+253
@@ -496,5 +496,258 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
496
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
497
498
499
+# INT4 W4A16 group-quantized MoE
500
+
501
+
502
+def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor:
503
+ """Pack INT4 values [N, K] → [N, K//8] int32 along K dim (GPTQ format)."""
504
+ N, K = w_int4.shape
505
+ assert K % 8 == 0
506
+ w = w_int4.to(torch.int32)
507
+ w_packed = torch.zeros(N, K // 8, dtype=torch.int32)
508
+ for j in range(8):
509
+ w_packed |= (w[:, j::8] & 0xF) << (j * 4)
510
+ return w_packed
511
+
512
+
513
+def _pack_int4_awq(w_int4: torch.Tensor) -> torch.Tensor:
514
+ """Pack INT4 values [..., N] → [..., N//8] int32 along last dim (AWQ format)."""
515
+ # AWQ packing bitshifts: indices {0,4,1,5,2,6,3,7} * 4 bits each
516
+ _AWQ_BITSHIFTS = [0, 16, 4, 20, 8, 24, 12, 28]
517
+
518
+ N = w_int4.shape[-1]
519
+ assert N % 8 == 0
520
+ w = w_int4.to(torch.int32)
521
+ w_packed = torch.zeros(*w.shape[:-1], N // 8, dtype=torch.int32)
522
+ for j, shift in enumerate(_AWQ_BITSHIFTS):
523
+ w_packed |= (w[..., j::8] & 0xF) << shift
524
+ return w_packed
525
+
526
+
527
+def _ref_int4_moe(
528
+ a: torch.Tensor,
529
+ w1_int4: torch.Tensor,
530
+ w2_int4: torch.Tensor,
531
+ w1_zeros: torch.Tensor | None,
532
+ w2_zeros: torch.Tensor | None,
533
+ w1_s: torch.Tensor,
534
+ w2_s: torch.Tensor,
535
+ topk_weight: torch.Tensor,
536
+ topk_ids: torch.Tensor,
537
+ group_size: int,
538
+) -> torch.Tensor:
539
+ """Reference INT4 W4A16 group-quantized fused MoE in pure torch."""
540
+ B = a.shape[0]
541
+ topk = topk_ids.size(1)
542
+ K_out = a.shape[1]
543
+
544
+ out = torch.zeros(B, topk, K_out, dtype=torch.float32)
545
+ for b in range(B):
546
+ for t in range(topk):
547
+ eid = topk_ids[b, t].item()
548
+ x = a[b : b + 1].float()
549
+
550
+ # Dequantize w1: [K, 2*N], groups along K (input dim)
551
+ K_dim = w1_int4.shape[1]
552
+ w1_dq = torch.zeros(K_dim, w1_int4.shape[2], dtype=torch.float32)
553
+ for g in range(w1_s.shape[1]):
554
+ k_start = g * group_size
555
+ k_end = min((g + 1) * group_size, K_dim)
556
+ zp = w1_zeros[eid, g, :].float() if w1_zeros is not None else 8.0
557
+ w1_dq[k_start:k_end, :] = (
558
+ w1_int4[eid, k_start:k_end, :].float() - zp
559
+ ) * w1_s[eid, g, :].float()
560
+
561
+ ic = torch.matmul(x, w1_dq) # [1, K] @ [K, 2*N] → [1, 2*N]
562
+ ic = _silu_and_mul(ic) # [1, N]
563
+
564
+ # Dequantize w2: [N, K], groups along N (input dim)
565
+ N_dim = w2_int4.shape[1]
566
+ w2_dq = torch.zeros(N_dim, w2_int4.shape[2], dtype=torch.float32)
567
+ for g in range(w2_s.shape[1]):
568
+ n_start = g * group_size
569
+ n_end = min((g + 1) * group_size, N_dim)
570
+ zp = w2_zeros[eid, g, :].float() if w2_zeros is not None else 8.0
571
+ w2_dq[n_start:n_end, :] = (
572
+ w2_int4[eid, n_start:n_end, :].float() - zp
573
+ ) * w2_s[eid, g, :].float()
574
+
575
+ oc = torch.matmul(ic, w2_dq) # [1, N] @ [N, K] → [1, K]
576
+ out[b, t] = oc.squeeze(0)
577
+
578
+ return (out * topk_weight.unsqueeze(-1)).sum(dim=1).to(a.dtype)
579
+
580
+
581
+def _make_int4_moe_weights(E, N, K, group_size, quant_algo):
582
+ """Create INT4 MoE weights in GPTQ or AWQ packed format.
583
+
584
+ Canonical layout (input × output):
585
+ w1_int4: [E, K, 2*N] w2_int4: [E, N, K]
586
+
587
+ GPTQ packed (pack transposed weight along input/K dim):
588
+ w1_packed: [E, K//8, 2*N] w2_packed: [E, N//8, K]
589
+ zeros: actual int4 zero points, same packing as weights
590
+
591
+ AWQ packed (pack along output/N dim):
592
+ w1_packed: [E, K, 2*N//8] w2_packed: [E, N, K//8]
593
+ zeros: actual int4 zero points, same packing as weights
594
+
595
+ Returns:
596
+ w1_int4, w2_int4,
597
+ w1_packed, w2_packed,
598
+ w1_zeros, w2_zeros,
599
+ w1_zeros_packed, w2_zeros_packed,
600
+ w1_s, w2_s
601
+ """
602
+ w1_int4 = torch.randint(0, 16, (E, K, 2 * N), dtype=torch.int32)
603
+ w2_int4 = torch.randint(0, 16, (E, N, K), dtype=torch.int32)
604
+
605
+ num_groups_w1 = K // group_size
606
+ num_groups_w2 = N // group_size
607
+ w1_s = (
608
+ torch.randn(E, num_groups_w1, 2 * N, dtype=torch.bfloat16) * 0.01
609
+ ).abs() + 0.001
610
+ w2_s = (torch.randn(E, num_groups_w2, K, dtype=torch.bfloat16) * 0.01).abs() + 0.001
611
+
612
+ if quant_algo == ops.CPUQuantAlgo.GPTQ:
613
+ # Pack: canonical [E, K, 2*N] → transpose [E, 2*N, K] → GPTQ pack
614
+ # [E, 2*N, K//8] → transpose [E, K//8, 2*N]
615
+ w1_t = w1_int4.transpose(1, 2).contiguous() # [E, 2*N, K]
616
+ w1_packed = (
617
+ torch.stack([_pack_int4_gptq(w1_t[e]) for e in range(E)])
618
+ .transpose(1, 2)
619
+ .contiguous()
620
+ ) # [E, K//8, 2*N]
621
+ w2_t = w2_int4.transpose(1, 2).contiguous() # [E, K, N]
622
+ w2_packed = (
623
+ torch.stack([_pack_int4_gptq(w2_t[e]) for e in range(E)])
624
+ .transpose(1, 2)
625
+ .contiguous()
626
+ ) # [E, N//8, K]
627
+ w1_zeros = w2_zeros = None
628
+ w1_zeros_packed = torch.full(
629
+ (E, num_groups_w1, 2 * N // 8), 0x77777777, dtype=torch.int32
630
+ )
631
+ w2_zeros_packed = torch.full(
632
+ (E, num_groups_w2, K // 8), 0x77777777, dtype=torch.int32
633
+ )
634
+ else: # AWQ
635
+ # Asymmetric: actual zero points, packed along output dim.
636
+ w1_zeros = torch.randint(1, 15, (E, num_groups_w1, 2 * N), dtype=torch.int32)
637
+ w2_zeros = torch.randint(1, 15, (E, num_groups_w2, K), dtype=torch.int32)
638
+ w1_packed = torch.stack(
639
+ [_pack_int4_awq(w1_int4[e]) for e in range(E)]
640
+ ) # [E, K, 2*N//8]
641
+ w2_packed = torch.stack(
642
+ [_pack_int4_awq(w2_int4[e]) for e in range(E)]
643
+ ) # [E, N, K//8]
644
+ w1_zeros_packed = torch.stack(
645
+ [_pack_int4_awq(w1_zeros[e]) for e in range(E)]
646
+ ) # [E, K//gs, 2*N//8]
647
+ w2_zeros_packed = torch.stack(
648
+ [_pack_int4_awq(w2_zeros[e]) for e in range(E)]
649
+ ) # [E, N//gs, K//8]
650
+
651
+ return (
652
+ w1_int4,
653
+ w2_int4,
654
+ w1_packed,
655
+ w2_packed,
656
+ w1_zeros,
657
+ w2_zeros,
658
+ w1_zeros_packed,
659
+ w2_zeros_packed,
660
+ w1_s,
661
+ w2_s,
662
+ )
663
+
664
+
665
+INT4_MOE_CONFIGS = [
666
+ # (N, K, E, topk, group_size)
667
+ (256, 512, 8, 2, 128),
668
+ (512, 256, 8, 2, 128),
669
+ (512, 512, 8, 4, 128),
670
+ (768, 2048, 8, 2, 128),
671
+]
672
+
673
+
674
+@pytest.mark.parametrize("M", [1, 2, 64, 121])
675
+@pytest.mark.parametrize("N,K,E,topk,group_size", INT4_MOE_CONFIGS)
676
+@pytest.mark.parametrize("quant_algo", [ops.CPUQuantAlgo.GPTQ, ops.CPUQuantAlgo.AWQ])
677
+@pytest.mark.parametrize("seed", [0])
678
+def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed):
679
+ """Test fused_experts_cpu INT4 W4A16 for both GPTQ and AWQ quant formats."""
680
+ set_random_seed(seed)
681
+
682
+ a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
683
+ (
684
+ w1_int4,
685
+ w2_int4,
686
+ w1_packed,
687
+ w2_packed,
688
+ w1_zeros,
689
+ w2_zeros,
690
+ w1_zeros_packed,
691
+ w2_zeros_packed,
692
+ w1_s,
693
+ w2_s,
694
+ ) = _make_int4_moe_weights(E, N, K, group_size, quant_algo)
695
+
696
+ score = torch.randn(M, E, dtype=torch.bfloat16)
697
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
698
+ topk_weight, topk_ids = torch.topk(score, topk)
699
+ topk_ids = topk_ids.to(torch.int32)
700
+
701
+ ref_out = _ref_int4_moe(
702
+ a,
703
+ w1_int4,
704
+ w2_int4,
705
+ w1_zeros,
706
+ w2_zeros,
707
+ w1_s,
708
+ w2_s,
709
+ topk_weight,
710
+ topk_ids,
711
+ group_size,
712
+ )
713
+
714
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
715
+ prepare_int4_moe_layer_for_cpu,
716
+ )
717
+
718
+ (blocked_w1, blocked_w2, blocked_s1, blocked_s2, blocked_z1, blocked_z2) = (
719
+ prepare_int4_moe_layer_for_cpu(
720
+ w1_packed,
721
+ w2_packed,
722
+ w1_s,
723
+ w2_s,
724
+ quant_algo=quant_algo,
725
+ w13_zeros=w1_zeros_packed,
726
+ w2_zeros=w2_zeros_packed,
727
+ )
728
+ )
729
+
730
+ out = ops.fused_experts_cpu(
731
+ a.clone(),
732
+ blocked_w1,
733
+ blocked_w2,
734
+ topk_weight,
735
+ topk_ids,
736
+ False, # inplace
737
+ ops.CPUQuantMethod.INT4_W4A8,
738
+ blocked_s1,
739
+ blocked_s2,
740
+ blocked_z1,
741
+ blocked_z2,
742
+ None, # block_size
743
+ None, # w1_bias
744
+ None, # w2_bias
745
+ None, # alpha
746
+ None, # limit
747
+ True, # is_vnni
748
+ )
749
+ torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
750
+
751
+
752
if __name__ == "__main__":
753
sys.exit(pytest.main([__file__, "-v"]))
tests/quantization/test_cpu_wna16.py
+3
@@ -16,6 +16,9 @@ MODELS = [
16
"Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear
17
"Qwen/Qwen3-30B-A3B-FP8", # FP8 W8A16 block-quantized MoE
18
"openai/gpt-oss-20b", # MXFP4 W4A16
19
+ "QuixiAI/Qwen3-30B-A3B-AWQ", # AWQ W4A16 MoE
20
+ "Qwen/Qwen3-30B-A3B-GPTQ-Int4", # GPTQ W4A16 MoE
21
+ "RedHatAI/Qwen3-30B-A3B-quantized.w4a16", # compressed-tensors W4A16 MoE
22
]
23
DTYPE = ["bfloat16"]
24
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+213
-1
@@ -5,7 +5,12 @@
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
-from vllm._custom_ops import CPUQuantMethod, fused_experts_cpu
8
+from vllm._custom_ops import (
9
+ CPUQuantAlgo,
10
+ CPUQuantMethod,
11
+ convert_weight_packed_scale_zp,
12
+ fused_experts_cpu,
13
+)
14
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEConfig,
@@ -17,6 +22,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
QuantKey,
23
kFp8Dynamic128Sym,
24
kFp8Static128BlockSym,
25
+ kInt4Static,
26
kMxfp4Static,
27
)
28
from vllm.platforms import current_platform
@@ -318,3 +324,209 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
324
limit,
325
True, # is_vnni
326
)
327
+
328
+
329
+def prepare_int4_moe_layer_for_cpu(
330
+ w13_packed: torch.Tensor,
331
+ w2_packed: torch.Tensor,
332
+ w13_scale: torch.Tensor,
333
+ w2_scale: torch.Tensor,
334
+ quant_algo: CPUQuantAlgo = CPUQuantAlgo.GPTQ,
335
+ w13_zeros: torch.Tensor | None = None,
336
+ w2_zeros: torch.Tensor | None = None,
337
+) -> tuple[
338
+ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
339
+]:
340
+ """Repack INT4 MoE weights via convert_weight_packed_scale_zp for CPU.
341
+
342
+ Args:
343
+ w13_packed: [E, K//8, 2*I] int32 (packed int4)
344
+ w2_packed: [E, I//8, K] int32 (packed int4)
345
+ w13_scale: [E, num_groups, 2*I] float16/bf16
346
+ w2_scale: [E, num_groups, K] float16/bf16
347
+ quant_algo: CPUQuantAlgo.GPTQ or CPUQuantAlgo.AWQ
348
+ w13_zeros: optional [E, num_groups, N//8] int32 packed zeros.
349
+ If None, synthetic zeros are created for symmetric quant.
350
+ w2_zeros: optional [E, num_groups, N//8] int32 packed zeros.
351
+ If None, synthetic zeros are created for symmetric quant.
352
+
353
+ Returns:
354
+ (blocked_w13, blocked_w2, blocked_s13, blocked_s2,
355
+ blocked_z13, blocked_z2)
356
+ """
357
+ E = w13_packed.size(0)
358
+
359
+ # No qzeros are available in compressed-tensors symmetric checkpoints.
360
+ # The GPTQ unpack kernel (unpack_4bit_to_32bit_signed) adds +1 to stored zeros,
361
+ # so we store 7 per nibble: 0x77777777 → +1 → 8.
362
+ if w13_zeros is None:
363
+ num_groups_w13 = w13_scale.size(1)
364
+ N_w13 = w13_scale.size(2) # 2*I
365
+ _zp = 0x77777777
366
+ w13_zeros = torch.full(
367
+ (E, num_groups_w13, N_w13 // 8),
368
+ _zp,
369
+ dtype=torch.int32,
370
+ )
371
+
372
+ if w2_zeros is None:
373
+ num_groups_w2 = w2_scale.size(1)
374
+ N_w2 = w2_scale.size(2) # K
375
+ _zp = 0x77777777
376
+ w2_zeros = torch.full(
377
+ (E, num_groups_w2, N_w2 // 8),
378
+ _zp,
379
+ dtype=torch.int32,
380
+ )
381
+
382
+ blocked_w13, blocked_z13, blocked_s13 = convert_weight_packed_scale_zp(
383
+ w13_packed, w13_zeros, w13_scale, quant_algo
384
+ )
385
+ blocked_w2, blocked_z2, blocked_s2 = convert_weight_packed_scale_zp(
386
+ w2_packed, w2_zeros, w2_scale, quant_algo
387
+ )
388
+ return (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2)
389
+
390
+
391
+class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
392
+ """CPU INT4 W4A16 group-quantized monolithic MoE experts.
393
+
394
+ Weights are int4 (packed), activations are bf16/fp16.
395
+ Internally uses int8 compute via fused_experts_cpu with INT4_W4A8.
396
+ """
397
+
398
+ def __init__(
399
+ self,
400
+ moe_config: FusedMoEConfig,
401
+ quant_config: FusedMoEQuantConfig,
402
+ max_num_tokens: int | None = None,
403
+ num_dispatchers: int | None = None,
404
+ ):
405
+ super().__init__(
406
+ moe_config,
407
+ quant_config,
408
+ )
409
+
410
+ @property
411
+ def expects_unquantized_inputs(self) -> bool:
412
+ return True
413
+
414
+ @staticmethod
415
+ def activation_format() -> mk.FusedMoEActivationFormat:
416
+ return mk.FusedMoEActivationFormat.Standard
417
+
418
+ @staticmethod
419
+ def _supports_current_device() -> bool:
420
+ return current_platform.is_cpu()
421
+
422
+ @staticmethod
423
+ def _supports_no_act_and_mul() -> bool:
424
+ return False
425
+
426
+ @staticmethod
427
+ def _supports_activation(activation: MoEActivation) -> bool:
428
+ return activation == MoEActivation.SILU
429
+
430
+ @staticmethod
431
+ def _supports_parallel_config(
432
+ moe_parallel_config: FusedMoEParallelConfig,
433
+ ) -> bool:
434
+ return True
435
+
436
+ @staticmethod
437
+ def _supports_quant_scheme(
438
+ weight_key: QuantKey | None,
439
+ activation_key: QuantKey | None,
440
+ ) -> bool:
441
+ SUPPORTED_W_A = [
442
+ (kInt4Static, None),
443
+ ]
444
+ return (weight_key, activation_key) in SUPPORTED_W_A
445
+
446
+ @staticmethod
447
+ def _supports_routing_method(
448
+ routing_method: RoutingMethodType,
449
+ weight_key: QuantKey | None,
450
+ activation_key: QuantKey | None,
451
+ ) -> bool:
452
+ return routing_method in [
453
+ RoutingMethodType.Default,
454
+ RoutingMethodType.Renormalize,
455
+ RoutingMethodType.RenormalizeNaive,
456
+ ]
457
+
458
+ @staticmethod
459
+ def _supports_router_logits_dtype(
460
+ router_logits_dtype: torch.dtype | None,
461
+ routing_method: RoutingMethodType,
462
+ ) -> bool:
463
+ return True
464
+
465
+ def supports_expert_map(self) -> bool:
466
+ return False
467
+
468
+ def apply(
469
+ self,
470
+ hidden_states: torch.Tensor,
471
+ w1: torch.Tensor,
472
+ w2: torch.Tensor,
473
+ router_logits: torch.Tensor,
474
+ activation: MoEActivation,
475
+ global_num_experts: int,
476
+ expert_map: torch.Tensor | None,
477
+ a1q_scale: torch.Tensor | None,
478
+ apply_router_weight_on_input: bool,
479
+ # grouped topk + fused topk bias parameters
480
+ num_expert_group: int | None = None,
481
+ e_score_correction_bias: torch.Tensor | None = None,
482
+ routed_scaling_factor: float | None = None,
483
+ topk_group: int | None = None,
484
+ ) -> torch.Tensor:
485
+ if apply_router_weight_on_input:
486
+ raise NotImplementedError(
487
+ "CPUExpertsInt4 (W4A16) does not support "
488
+ "apply_router_weight_on_input=True. "
489
+ )
490
+
491
+ from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
492
+ select_experts,
493
+ )
494
+
495
+ topk_weights, topk_ids = select_experts(
496
+ hidden_states=hidden_states,
497
+ router_logits=router_logits,
498
+ use_grouped_topk=num_expert_group is not None,
499
+ top_k=self.moe_config.experts_per_token,
500
+ renormalize=self.moe_config.routing_method
501
+ in (
502
+ RoutingMethodType.Renormalize,
503
+ RoutingMethodType.RenormalizeNaive,
504
+ ),
505
+ topk_group=topk_group,
506
+ num_expert_group=num_expert_group,
507
+ scoring_func="softmax",
508
+ routed_scaling_factor=(
509
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
510
+ ),
511
+ e_score_correction_bias=e_score_correction_bias,
512
+ )
513
+
514
+ return fused_experts_cpu(
515
+ hidden_states,
516
+ w1,
517
+ w2,
518
+ topk_weights,
519
+ topk_ids,
520
+ False, # inplace
521
+ CPUQuantMethod.INT4_W4A8,
522
+ self.w1_scale,
523
+ self.w2_scale,
524
+ self.w1_zp,
525
+ self.w2_zp,
526
+ None, # block_size
527
+ None, # w1_bias
528
+ None, # w2_bias
529
+ None, # alpha
530
+ None, # limit
531
+ True, # is_vnni
532
+ )
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+140
-2
@@ -45,6 +45,7 @@ logger = init_logger(__name__)
45
class WNA16MoEBackend(Enum):
46
MARLIN = "MARLIN"
47
BATCHED_MARLIN = "BATCHED_MARLIN"
48
+ CPU = "CPU"
49
FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
50
XPU = "XPU"
51
@@ -65,6 +66,12 @@ def backend_to_kernel_cls(
66
)
67
68
return [XPUExpertsWNA16]
69
+ elif backend == WNA16MoEBackend.CPU:
70
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
71
+ CPUExpertsInt4,
72
+ )
73
+
74
+ return [CPUExpertsInt4]
75
else:
76
raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
77
@@ -73,6 +80,8 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
80
"""
81
Get available backends in priority order based on platform and config.
82
"""
83
+ if current_platform.is_cpu():
84
+ return [WNA16MoEBackend.CPU]
85
if current_platform.is_xpu():
86
return [WNA16MoEBackend.XPU]
87
@@ -210,17 +219,21 @@ def make_wna16_moe_kernel(
219
from vllm.model_executor.layers.fused_moe.all2all_utils import (
220
maybe_make_prepare_finalize,
221
)
222
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
223
+ CPUExpertsInt4,
224
+ )
225
from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
226
XPUExpertsWNA16,
227
)
228
217
- # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts
218
- # and BatchedMarlinExperts
229
+ # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts,
230
+ # BatchedMarlinExperts, XPUExpertsWNA16, and CPUExpertsInt4
231
assert experts_cls in (
232
MarlinExperts,
233
BatchedMarlinExperts,
234
TrtLlmMxint4ExpertsMonolithic,
235
XPUExpertsWNA16,
236
+ CPUExpertsInt4,
237
)
238
239
is_monolithic = experts_cls.is_monolithic()
@@ -683,6 +696,117 @@ def _process_awq_weights_marlin(
696
)
697
698
699
+def _process_weights_cpu(
700
+ quant_config: QuantizationConfig | QuantizationArgs | None,
701
+ w13: torch.Tensor,
702
+ w2: torch.Tensor,
703
+ w13_scale: torch.Tensor,
704
+ w2_scale: torch.Tensor,
705
+ w13_g_idx: torch.Tensor | None = None,
706
+ w2_g_idx: torch.Tensor | None = None,
707
+ w13_qzeros: torch.Tensor | None = None,
708
+ w2_qzeros: torch.Tensor | None = None,
709
+ w13_bias: torch.Tensor | None = None,
710
+ w2_bias: torch.Tensor | None = None,
711
+) -> tuple[
712
+ torch.Tensor, # w13_qweight
713
+ torch.Tensor, # w2_qweight
714
+ torch.Tensor, # w13_scales
715
+ torch.Tensor, # w2_scales
716
+ torch.Tensor | None, # w13_g_idx
717
+ torch.Tensor | None, # w2_g_idx
718
+ torch.Tensor | None, # w13_g_idx_sort_indices
719
+ torch.Tensor | None, # w2_g_idx_sort_indices
720
+ torch.Tensor | None, # w13_qzeros
721
+ torch.Tensor | None, # w2_qzeros
722
+ torch.Tensor | None, # w13_input_global_scale
723
+ torch.Tensor | None, # w2_input_global_scale
724
+ torch.Tensor | None, # w13_bias
725
+ torch.Tensor | None, # w2_bias
726
+]:
727
+ """CPU INT4 W4A16 weight post-processing."""
728
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
729
+ prepare_int4_moe_layer_for_cpu,
730
+ )
731
+ from vllm.model_executor.layers.quantization.auto_gptq import (
732
+ AutoGPTQConfig,
733
+ )
734
+ from vllm.model_executor.layers.quantization.awq_marlin import (
735
+ AWQMarlinConfig,
736
+ )
737
+
738
+ # Detect packing format.
739
+ # AWQ: qweight is [E, K, 2*N//8] (packed along output/N dim).
740
+ # GPTQ: qweight is [E, K//8, 2*N] (packed along input/K dim).
741
+ # compressed-tensors: qweight is [E, K//8, 2*N] (packed along input/K dim).
742
+ if isinstance(quant_config, AWQMarlinConfig):
743
+ # AWQ: K is stored unpacked in dim 1.
744
+ cpu_quant_algo = ops.CPUQuantAlgo.AWQ
745
+ elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
746
+ # GPTQ / compressed-tensors: K//8 is stored packed in dim 1.
747
+ if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
748
+ raise NotImplementedError(
749
+ "CPU WNA16 MoE backend does not support GPTQ with "
750
+ "desc_act=True. The fused MoE kernel has no g_idx "
751
+ "reordering support."
752
+ )
753
+ cpu_quant_algo = ops.CPUQuantAlgo.GPTQ
754
+ else:
755
+ raise TypeError(
756
+ "CPU WNA16 MoE backend requires AWQMarlinConfig, AutoGPTQConfig "
757
+ f"or QuantizationArgs, got {type(quant_config).__name__}."
758
+ )
759
+
760
+ # Determine zero points for repacking.
761
+ w13_zeros: torch.Tensor | None = None
762
+ w2_zeros: torch.Tensor | None = None
763
+ if w13_qzeros is not None:
764
+ w13_zeros = (
765
+ w13_qzeros.data.view(torch.int32)
766
+ if w13_qzeros.dtype != torch.int32
767
+ else w13_qzeros.data
768
+ )
769
+ if w2_qzeros is not None:
770
+ w2_zeros = (
771
+ w2_qzeros.data.view(torch.int32)
772
+ if w2_qzeros.dtype != torch.int32
773
+ else w2_qzeros.data
774
+ )
775
+
776
+ (
777
+ blocked_w13,
778
+ blocked_w2,
779
+ blocked_s13,
780
+ blocked_s2,
781
+ blocked_z13,
782
+ blocked_z2,
783
+ ) = prepare_int4_moe_layer_for_cpu(
784
+ w13,
785
+ w2,
786
+ w13_scale,
787
+ w2_scale,
788
+ quant_algo=cpu_quant_algo,
789
+ w13_zeros=w13_zeros,
790
+ w2_zeros=w2_zeros,
791
+ )
792
+ return (
793
+ blocked_w13,
794
+ blocked_w2,
795
+ blocked_s13,
796
+ blocked_s2,
797
+ w13_g_idx,
798
+ w2_g_idx,
799
+ None, # w13_g_idx_sort_indices (unused on CPU)
800
+ None, # w2_g_idx_sort_indices (unused on CPU)
801
+ blocked_z13,
802
+ blocked_z2,
803
+ None, # w13_input_global_scale
804
+ None, # w2_input_global_scale
805
+ w13_bias.to(torch.float32) if w13_bias is not None else None,
806
+ w2_bias.to(torch.float32) if w2_bias is not None else None,
807
+ )
808
+
809
+
810
def _process_weights_xpu(
811
layer: torch.nn.Module,
812
quant_config: QuantizationConfig,
@@ -857,6 +981,20 @@ def convert_to_wna16_moe_kernel_format(
981
w13_bias,
982
w2_bias,
983
)
984
+ elif backend == WNA16MoEBackend.CPU:
985
+ return _process_weights_cpu(
986
+ quant_config,
987
+ w13,
988
+ w2,
989
+ w13_scale,
990
+ w2_scale,
991
+ w13_g_idx,
992
+ w2_g_idx,
993
+ w13_qzeros,
994
+ w2_qzeros,
995
+ w13_bias,
996
+ w2_bias,
997
+ )
998
elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
999
return _process_weights_flashinfer(
1000
w13,
vllm/model_executor/layers/quantization/auto_gptq.py
+40
-12
@@ -16,6 +16,7 @@ from vllm.model_executor.kernels.linear import (
16
)
17
from vllm.model_executor.layers.fused_moe import (
18
FusedMoEConfig,
19
+ FusedMoEExpertsModular,
20
FusedMoEMethodBase,
21
FusedMoEQuantConfig,
22
FusedMoeWeightScaleSupported,
@@ -640,8 +641,11 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
641
layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
642
set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
643
643
- device = layer.w13_qweight.device
644
- layer.workspace = marlin_make_workspace_new(device, 4)
644
+ if self.experts_cls is not None and issubclass(
645
+ self.experts_cls, FusedMoEExpertsModular
646
+ ):
647
+ device = layer.w13_qweight.device
648
+ layer.workspace = marlin_make_workspace_new(device, 4)
649
650
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
651
is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1
@@ -660,8 +664,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
664
w2_g_idx,
665
w13_g_idx_sort_indices,
666
w2_g_idx_sort_indices,
663
- _w13_qzeros,
664
- _w2_qzeros,
667
+ w13_qzeros,
668
+ w2_qzeros,
669
w13_input_global_scale,
670
w2_input_global_scale,
671
w13_bias,
@@ -689,6 +693,10 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
693
replace_parameter(layer, "w2_g_idx", w2_g_idx)
694
replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
695
replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
696
+ if w13_qzeros is not None:
697
+ replace_parameter(layer, "w13_qzeros", w13_qzeros)
698
+ if w2_qzeros is not None:
699
+ replace_parameter(layer, "w2_qzeros", w2_qzeros)
700
if w13_input_global_scale is not None:
701
if hasattr(layer, "w13_input_global_scale"):
702
replace_parameter(
@@ -735,8 +743,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
743
is_k_full=self.is_k_full,
744
w13_g_idx=layer.w13_g_idx,
745
w2_g_idx=layer.w2_g_idx,
738
- w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
739
- w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
746
+ w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
747
+ w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
748
routing_tables=layer._expert_routing_tables(),
749
)
750
@@ -750,12 +758,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
758
w2_scale=layer.w2_scales,
759
weight_bits=self.quant_config.weight_bits,
760
group_size=self.quant_config.group_size,
753
- w1_zp=getattr(layer, "w13_qzeros", None)
754
- if not self.quant_config.is_sym
755
- else None,
756
- w2_zp=getattr(layer, "w2_qzeros", None)
757
- if not self.quant_config.is_sym
758
- else None,
761
+ w1_zp=getattr(layer, "w13_qzeros", None),
762
+ w2_zp=getattr(layer, "w2_qzeros", None),
763
w1_bias=getattr(layer, "w13_bias", None),
764
w2_bias=getattr(layer, "w2_bias", None),
765
)
@@ -794,3 +798,27 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
798
shared_experts=shared_experts,
799
shared_experts_input=shared_experts_input,
800
)
801
+
802
+ def apply_monolithic(
803
+ self,
804
+ layer: RoutedExperts,
805
+ x: torch.Tensor,
806
+ router_logits: torch.Tensor,
807
+ input_ids: torch.Tensor | None = None,
808
+ ) -> torch.Tensor:
809
+ assert self.is_monolithic
810
+ assert self.moe_kernel is not None
811
+ return self.moe_kernel.apply_monolithic(
812
+ hidden_states=x,
813
+ w1=layer.w13_qweight,
814
+ w2=layer.w2_qweight,
815
+ router_logits=router_logits,
816
+ activation=layer.activation,
817
+ global_num_experts=layer.global_num_experts,
818
+ expert_map=layer.expert_map,
819
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
820
+ num_expert_group=layer.num_expert_group,
821
+ topk_group=layer.topk_group,
822
+ e_score_correction_bias=layer.e_score_correction_bias,
823
+ routed_scaling_factor=layer.routed_scaling_factor,
824
+ )
vllm/model_executor/layers/quantization/awq_marlin.py
+26
-2
@@ -700,8 +700,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
700
is_k_full=self.is_k_full,
701
w13_g_idx=getattr(layer, "w13_g_idx", None),
702
w2_g_idx=getattr(layer, "w2_g_idx", None),
703
- w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
704
- w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
703
+ w13_g_idx_sort_indices=getattr(layer, "w13_g_idx_sort_indices", None),
704
+ w2_g_idx_sort_indices=getattr(layer, "w2_g_idx_sort_indices", None),
705
routing_tables=layer._expert_routing_tables(),
706
)
707
@@ -757,3 +757,27 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
757
shared_experts=shared_experts,
758
shared_experts_input=shared_experts_input,
759
)
760
+
761
+ def apply_monolithic(
762
+ self,
763
+ layer: RoutedExperts,
764
+ x: torch.Tensor,
765
+ router_logits: torch.Tensor,
766
+ input_ids: torch.Tensor | None = None,
767
+ ) -> torch.Tensor:
768
+ assert self.is_monolithic
769
+ assert self.moe_kernel is not None
770
+ return self.moe_kernel.apply_monolithic(
771
+ hidden_states=x,
772
+ w1=layer.w13_qweight,
773
+ w2=layer.w2_qweight,
774
+ router_logits=router_logits,
775
+ activation=layer.activation,
776
+ global_num_experts=layer.global_num_experts,
777
+ expert_map=layer.expert_map,
778
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
779
+ num_expert_group=layer.num_expert_group,
780
+ topk_group=layer.topk_group,
781
+ e_score_correction_bias=layer.e_score_correction_bias,
782
+ routed_scaling_factor=layer.routed_scaling_factor,
783
+ )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+9
-4
@@ -10,6 +10,7 @@ from compressed_tensors.quantization import (
10
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.fused_moe import (
13
+ FusedMoEExpertsModular,
14
RoutedExperts,
15
SharedExperts,
16
)
@@ -414,8 +415,9 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
415
replace_parameter(layer, "w13_weight_scale", w13_scales)
416
replace_parameter(layer, "w2_weight_scale", w2_scales)
417
417
- if not self.symmetric:
418
+ if w13_qzeros is not None:
419
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
420
+ if w2_qzeros is not None:
421
replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
422
423
# Marlin-specific parameters (not needed for Flashinfer)
@@ -437,9 +439,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
439
torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
440
)
441
440
- layer.workspace = marlin_make_workspace_new(
441
- layer.w13_weight_g_idx.device, 4
442
- )
442
+ if self.experts_cls is not None and issubclass(
443
+ self.experts_cls, FusedMoEExpertsModular
444
+ ):
445
+ layer.workspace = marlin_make_workspace_new(
446
+ layer.w13_weight_g_idx.device, 4
447
+ )
448
449
# Alias packed weights to w13_weight/w2_weight for the modular kernel interface
450
layer.w13_weight = layer.w13_weight_packed