[CPU] Support cpu compressed-tensor w8a8 int8 moe (#42920)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com> Signed-off-by: Yuwen Zhou <yuwen.zhou@intel.com>
Yuwen Zhou committed
Jun 29, 2026 at 11:04 UTC
58d6a6e60ae6bd94a20ea6da27eb224188b24dca
6 files changed
+410
-9
.buildkite/hardware_tests/cpu.yaml
+3
-1
@@ -91,11 +91,13 @@ steps:
91
- vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
92
- tests/quantization/test_compressed_tensors.py
93
- tests/quantization/test_cpu_wna16.py
94
+ - tests/quantization/test_cpu_w8a8.py
95
commands:
96
- |
97
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
98
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
98
- pytest -x -v -s tests/quantization/test_cpu_wna16.py"
99
+ pytest -x -v -s tests/quantization/test_cpu_wna16.py
100
+ pytest -x -v -s tests/quantization/test_cpu_w8a8.py"
101
102
- label: CPU-Distributed Tests (PP+TP)
103
depends_on: []
tests/kernels/moe/test_cpu_quant_fused_moe.py
+139
-4
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for CPU quantized fused MoE kernels (FP8 W8A16 and MXFP4 W4A16)."""
3
+"""Tests for CPU quantized fused MoE kernels."""
4
5
import math
6
import sys
@@ -31,7 +31,10 @@ def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
31
return torch.ops._C.convert_weight_packed(w)
32
33
34
-# FP8 W8A16 block-scaled fused MoE
34
+# ===========================================================================
35
+# FP8 W8A16 MoE
36
+# ===========================================================================
37
+
38
39
BLOCK_SIZE = [128, 128] # [block_n, block_k]
40
@@ -216,7 +219,9 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
219
torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
220
221
219
-# MXFP4 W4A16 fused MoE
222
+# ===========================================================================
223
+# MXFP4 W4A16 MoE
224
+# ===========================================================================
225
226
227
class MXFP4QuantizeUtil:
@@ -496,7 +501,9 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
501
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
502
503
499
-# INT4 W4A16 group-quantized MoE
504
+# ===========================================================================
505
+# INT4 W4A16 MoE
506
+# ===========================================================================
507
508
509
def _pack_int4_gptq(w_int4: torch.Tensor) -> torch.Tensor:
@@ -749,5 +756,133 @@ def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed
756
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
757
758
759
+# ===========================================================================
760
+# INT8 W8A8 MoE
761
+# ===========================================================================
762
+
763
+
764
+def _quantize_per_channel(w):
765
+ """Symmetric per-channel INT8 quantisation. w: [N, K] -> (int8, scale)."""
766
+ amax = w.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
767
+ scale = amax / 127.0
768
+ w_q = (w / scale).round().clamp(-128, 127).to(torch.int8)
769
+ return w_q, scale.float()
770
+
771
+
772
+def _quantize_per_token(x):
773
+ """Symmetric per-token INT8 quantisation. x: [M, K] -> (int8, scale)."""
774
+ amax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-12)
775
+ scale = amax / 127.0
776
+ x_q = (x / scale).round().clamp(-128, 127).to(torch.int8)
777
+ return x_q, scale.float()
778
+
779
+
780
+def _ref_int8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids):
781
+ """Reference INT8 W8A8 per-channel fused MoE in pure torch."""
782
+ B, D = a.shape
783
+ topk = topk_ids.size(1)
784
+
785
+ out = torch.zeros(B, topk, w2.shape[1], dtype=torch.float32)
786
+ for b in range(B):
787
+ for t in range(topk):
788
+ eid = topk_ids[b, t].item()
789
+
790
+ x = a[b : b + 1].float()
791
+ x_q, x_s = _quantize_per_token(x)
792
+ ic = torch.matmul(x_q.float(), w1[eid].float().t())
793
+ ic = ic * x_s * w1_s[eid].view(1, -1)
794
+ ic = _silu_and_mul(ic)
795
+
796
+ ic_q, ic_s = _quantize_per_token(ic)
797
+ oc = torch.matmul(ic_q.float(), w2[eid].float().t())
798
+ oc = oc * ic_s * w2_s[eid].view(1, -1)
799
+ out[b, t] = oc.squeeze(0)
800
+
801
+ result = (out * topk_weight.unsqueeze(-1)).sum(dim=1)
802
+ return result.to(a.dtype)
803
+
804
+
805
+def _make_int8_moe_weights(E, N, K):
806
+ factor = 1e-2
807
+ w1_f = (torch.randn(E, 2 * N, K) - 0.5) * 2
808
+ w2_f = (torch.randn(E, K, N) - 0.5) * 2
809
+
810
+ w1_q_list, w1_s_list = [], []
811
+ w2_q_list, w2_s_list = [], []
812
+ for e in range(E):
813
+ q, s = _quantize_per_channel(w1_f[e])
814
+ w1_q_list.append(q)
815
+ w1_s_list.append(s)
816
+ q, s = _quantize_per_channel(w2_f[e])
817
+ w2_q_list.append(q)
818
+ w2_s_list.append(s)
819
+
820
+ return (
821
+ torch.stack(w1_q_list),
822
+ torch.stack(w2_q_list),
823
+ torch.stack(w1_s_list) * factor,
824
+ torch.stack(w2_s_list) * factor,
825
+ )
826
+
827
+
828
+INT8_NUM_TOKENS = [1, 2, 64, 121]
829
+INT8_MOE_CONFIGS = [
830
+ # (N, K, E, topk)
831
+ (256, 512, 8, 2),
832
+ (512, 256, 8, 2),
833
+ (512, 512, 8, 4),
834
+ (768, 2048, 8, 2),
835
+]
836
+
837
+
838
+@pytest.mark.parametrize("M", INT8_NUM_TOKENS)
839
+@pytest.mark.parametrize("N,K,E,topk", INT8_MOE_CONFIGS)
840
+@pytest.mark.parametrize("seed", [0])
841
+@pytest.mark.parametrize("is_vnni", [False, True])
842
+@pytest.mark.parametrize("inplace", [False, True])
843
+def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
844
+ """Test fused_experts_cpu INT8 W8A8 against torch reference."""
845
+ set_random_seed(seed)
846
+
847
+ a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
848
+ w1_q, w2_q, w1_s, w2_s = _make_int8_moe_weights(E, N, K)
849
+
850
+ score = torch.randn(M, E, dtype=torch.bfloat16)
851
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
852
+ topk_weight, topk_ids = torch.topk(score, topk)
853
+ topk_ids = topk_ids.to(torch.int32)
854
+
855
+ ref_out = _ref_int8_moe(a, w1_q, w2_q, w1_s, w2_s, topk_weight, topk_ids)
856
+
857
+ w1 = _prepack_experts(w1_q) if is_vnni else w1_q
858
+ w2 = _prepack_experts(w2_q) if is_vnni else w2_q
859
+
860
+ out = ops.fused_experts_cpu(
861
+ a.clone(),
862
+ w1,
863
+ w2,
864
+ topk_weight,
865
+ topk_ids,
866
+ inplace,
867
+ ops.CPUQuantMethod.INT8_W8A8,
868
+ w1_s,
869
+ w2_s,
870
+ None, # w1_zero
871
+ None, # w2_zero
872
+ None, # block_size
873
+ None, # w1_bias
874
+ None, # w2_bias
875
+ None, # alpha
876
+ None, # limit
877
+ is_vnni,
878
+ )
879
+ torch.testing.assert_close(
880
+ ref_out.bfloat16(),
881
+ out,
882
+ atol=2e-1,
883
+ rtol=2e-1,
884
+ )
885
+
886
+
887
if __name__ == "__main__":
888
sys.exit(pytest.main([__file__, "-v"]))
tests/quantization/test_cpu_w8a8.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from vllm.platforms import current_platform
6
+
7
+if not current_platform.is_cpu():
8
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
9
+
10
+MODELS = [
11
+ "RedHatAI/Qwen3-30B-A3B-Instruct-2507-quantized.w8a8", # INT8 W8A8 MoE
12
+]
13
+DTYPE = ["bfloat16"]
14
+
15
+
16
+@pytest.mark.parametrize("model", MODELS)
17
+@pytest.mark.parametrize("dtype", DTYPE)
18
+def test_cpu_w8a8(vllm_runner, model, dtype):
19
+ with vllm_runner(model, dtype=dtype) as llm:
20
+ output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
21
+ assert output
22
+ print(output)
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+173
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""CPU FP8 W8A16 and MXFP4 W4A16 fused MoE experts."""
3
+"""CPU quantized fused MoE experts."""
4
5
import torch
6
@@ -23,10 +23,16 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
kFp8Dynamic128Sym,
24
kFp8Static128BlockSym,
25
kInt4Static,
26
+ kInt8DynamicTokenSym,
27
+ kInt8StaticChannelSym,
28
kMxfp4Static,
29
)
30
from vllm.platforms import current_platform
31
32
+# ===========================================================================
33
+# FP8 W8A16 MoE
34
+# ===========================================================================
35
+
36
37
def prepare_fp8_moe_layer_for_cpu(
38
w13: torch.Tensor,
@@ -177,6 +183,11 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
183
)
184
185
186
+# ===========================================================================
187
+# MXFP4 W4A16 MoE
188
+# ===========================================================================
189
+
190
+
191
def prepare_mxfp4_moe_layer_for_cpu(
192
w13: torch.Tensor,
193
w2: torch.Tensor,
@@ -326,6 +337,11 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
337
)
338
339
340
+# ===========================================================================
341
+# INT4 W4A16 MoE
342
+# ===========================================================================
343
+
344
+
345
def prepare_int4_moe_layer_for_cpu(
346
w13_packed: torch.Tensor,
347
w2_packed: torch.Tensor,
@@ -529,3 +545,159 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
545
None, # limit
546
True, # is_vnni
547
)
548
+
549
+
550
+# ===========================================================================
551
+# INT8 W8A8 MoE
552
+# ===========================================================================
553
+
554
+
555
+def prepare_int8_moe_layer_for_cpu(
556
+ w13: torch.Tensor,
557
+ w2: torch.Tensor,
558
+) -> tuple[torch.Tensor, torch.Tensor]:
559
+ """VNNI-prepack INT8 MoE weights for CPU kernel."""
560
+ packed_w13 = torch.ops._C.convert_weight_packed(w13)
561
+ packed_w2 = torch.ops._C.convert_weight_packed(w2)
562
+ return packed_w13, packed_w2
563
+
564
+
565
+class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
566
+ """CPU INT8 W8A8 per-channel weight / dynamic per-token activation
567
+ monolithic MoE experts."""
568
+
569
+ def __init__(
570
+ self,
571
+ moe_config: FusedMoEConfig,
572
+ quant_config: FusedMoEQuantConfig,
573
+ ):
574
+ super().__init__(
575
+ moe_config,
576
+ quant_config,
577
+ )
578
+
579
+ @property
580
+ def expects_unquantized_inputs(self) -> bool:
581
+ return True
582
+
583
+ @staticmethod
584
+ def activation_format() -> mk.FusedMoEActivationFormat:
585
+ return mk.FusedMoEActivationFormat.Standard
586
+
587
+ @staticmethod
588
+ def _supports_current_device() -> bool:
589
+ return current_platform.is_cpu()
590
+
591
+ @staticmethod
592
+ def _supports_no_act_and_mul() -> bool:
593
+ return False
594
+
595
+ @staticmethod
596
+ def _supports_activation(activation: MoEActivation) -> bool:
597
+ return activation == MoEActivation.SILU
598
+
599
+ @staticmethod
600
+ def _supports_parallel_config(
601
+ moe_parallel_config: FusedMoEParallelConfig,
602
+ ) -> bool:
603
+ return True
604
+
605
+ @staticmethod
606
+ def _supports_quant_scheme(
607
+ weight_key: QuantKey | None,
608
+ activation_key: QuantKey | None,
609
+ ) -> bool:
610
+ SUPPORTED_W_A = [
611
+ (kInt8StaticChannelSym, kInt8DynamicTokenSym),
612
+ ]
613
+ return (weight_key, activation_key) in SUPPORTED_W_A
614
+
615
+ @staticmethod
616
+ def _supports_routing_method(
617
+ routing_method: RoutingMethodType,
618
+ weight_key: QuantKey | None,
619
+ activation_key: QuantKey | None,
620
+ ) -> bool:
621
+ return routing_method in [
622
+ RoutingMethodType.Default,
623
+ RoutingMethodType.Renormalize,
624
+ RoutingMethodType.RenormalizeNaive,
625
+ ]
626
+
627
+ @staticmethod
628
+ def _supports_router_logits_dtype(
629
+ router_logits_dtype: torch.dtype | None,
630
+ routing_method: RoutingMethodType,
631
+ ) -> bool:
632
+ return True
633
+
634
+ def supports_expert_map(self) -> bool:
635
+ return False
636
+
637
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
638
+ """VNNI-prepack INT8 MoE weights for CPU kernel."""
639
+ from vllm.model_executor.utils import replace_parameter
640
+
641
+ w13, w2 = prepare_int8_moe_layer_for_cpu(layer.w13_weight, layer.w2_weight)
642
+ replace_parameter(layer, "w13_weight", w13)
643
+ replace_parameter(layer, "w2_weight", w2)
644
+
645
+ def apply(
646
+ self,
647
+ hidden_states: torch.Tensor,
648
+ w1: torch.Tensor,
649
+ w2: torch.Tensor,
650
+ router_logits: torch.Tensor,
651
+ activation: MoEActivation,
652
+ global_num_experts: int,
653
+ expert_map: torch.Tensor | None,
654
+ a1q_scale: torch.Tensor | None,
655
+ apply_router_weight_on_input: bool,
656
+ # grouped topk + fused topk bias parameters
657
+ num_expert_group: int | None = None,
658
+ e_score_correction_bias: torch.Tensor | None = None,
659
+ routed_scaling_factor: float | None = None,
660
+ topk_group: int | None = None,
661
+ ) -> torch.Tensor:
662
+ from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
663
+ select_experts,
664
+ )
665
+
666
+ topk_weights, topk_ids = select_experts(
667
+ hidden_states=hidden_states,
668
+ router_logits=router_logits,
669
+ use_grouped_topk=num_expert_group is not None,
670
+ top_k=self.moe_config.experts_per_token,
671
+ renormalize=self.moe_config.routing_method
672
+ in (
673
+ RoutingMethodType.Renormalize,
674
+ RoutingMethodType.RenormalizeNaive,
675
+ ),
676
+ topk_group=topk_group,
677
+ num_expert_group=num_expert_group,
678
+ scoring_func="softmax",
679
+ routed_scaling_factor=(
680
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
681
+ ),
682
+ e_score_correction_bias=e_score_correction_bias,
683
+ )
684
+
685
+ return fused_experts_cpu(
686
+ hidden_states,
687
+ w1,
688
+ w2,
689
+ topk_weights,
690
+ topk_ids,
691
+ False, # inplace
692
+ CPUQuantMethod.INT8_W8A8,
693
+ self.w1_scale,
694
+ self.w2_scale,
695
+ None, # w1_zero
696
+ None, # w2_zero
697
+ None, # block_size (per-channel, no block)
698
+ None, # w1_bias
699
+ None, # w2_bias
700
+ None, # alpha
701
+ None, # limit
702
+ True, # is_vnni
703
+ )
vllm/model_executor/layers/fused_moe/oracle/int8.py
+39
-1
@@ -22,12 +22,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
kInt8DynamicTokenSym,
23
kInt8StaticChannelSym,
24
)
25
+from vllm.platforms import current_platform
26
27
logger = init_logger(__name__)
28
29
30
class Int8MoeBackend(Enum):
31
TRITON = "TRITON"
32
+ CPU = "CPU"
33
34
35
def _get_priority_backends(
@@ -36,7 +38,18 @@ def _get_priority_backends(
38
"""
39
Get available backends in priority order based on platform and config.
40
"""
39
- return [Int8MoeBackend.TRITON]
41
+ _AVAILABLE_BACKENDS = [
42
+ Int8MoeBackend.TRITON,
43
+ Int8MoeBackend.CPU,
44
+ ]
45
+
46
+ def _move_to_front(backends: list[Int8MoeBackend], backend: Int8MoeBackend) -> None:
47
+ backends.insert(0, backends.pop(backends.index(backend)))
48
+
49
+ if current_platform.is_cpu():
50
+ _move_to_front(_AVAILABLE_BACKENDS, Int8MoeBackend.CPU)
51
+
52
+ return _AVAILABLE_BACKENDS
53
54
55
def backend_to_kernel_cls(
@@ -49,6 +62,13 @@ def backend_to_kernel_cls(
62
63
return [TritonExperts]
64
65
+ elif backend == Int8MoeBackend.CPU:
66
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
67
+ CPUExpertsInt8,
68
+ )
69
+
70
+ return [CPUExpertsInt8]
71
+
72
else:
73
raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
74
@@ -176,6 +196,24 @@ def make_int8_moe_quant_config(
196
)
197
198
199
+def convert_to_int8_moe_kernel_format(
200
+ int8_backend: Int8MoeBackend,
201
+ w13: torch.Tensor,
202
+ w2: torch.Tensor,
203
+) -> tuple[torch.Tensor, torch.Tensor]:
204
+ """Convert INT8 MoE weights to backend-specific kernel format."""
205
+ if int8_backend == Int8MoeBackend.CPU:
206
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
207
+ prepare_int8_moe_layer_for_cpu,
208
+ )
209
+
210
+ w13, w2 = prepare_int8_moe_layer_for_cpu(w13, w2)
211
+ elif int8_backend != Int8MoeBackend.TRITON:
212
+ raise ValueError(f"Unsupported Int8 MoE backend: {int8_backend.value}")
213
+
214
+ return w13, w2
215
+
216
+
217
def make_int8_moe_kernel(
218
moe_quant_config: FusedMoEQuantConfig,
219
moe_config: FusedMoEConfig,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+34
-2
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-
4
import torch
5
from compressed_tensors.quantization import (
6
QuantizationArgs,
@@ -20,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import (
19
FusedMoEQuantConfig,
20
)
21
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
22
+ convert_to_int8_moe_kernel_format,
23
make_int8_moe_kernel,
24
make_int8_moe_quant_config,
25
select_int8_moe_backend,
@@ -31,7 +31,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
kInt8DynamicTokenSym,
32
kInt8StaticChannelSym,
33
)
34
-from vllm.model_executor.utils import set_weight_attrs
34
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
35
36
logger = init_logger(__name__)
37
@@ -142,6 +142,14 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
142
layer.w2_input_scale = None
143
144
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
145
+ w13, w2 = convert_to_int8_moe_kernel_format(
146
+ int8_backend=self.int8_backend,
147
+ w13=layer.w13_weight,
148
+ w2=layer.w2_weight,
149
+ )
150
+ replace_parameter(layer, "w13_weight", w13)
151
+ replace_parameter(layer, "w2_weight", w2)
152
+
153
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
154
assert self.experts_cls is not None
155
self.moe_kernel = make_int8_moe_kernel(
@@ -193,3 +201,27 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
201
shared_experts=shared_experts,
202
shared_experts_input=shared_experts_input,
203
)
204
+
205
+ def apply_monolithic(
206
+ self,
207
+ layer: RoutedExperts,
208
+ x: torch.Tensor,
209
+ router_logits: torch.Tensor,
210
+ input_ids: torch.Tensor | None = None,
211
+ ) -> torch.Tensor:
212
+ assert self.is_monolithic
213
+ assert self.moe_kernel is not None
214
+ return self.moe_kernel.apply_monolithic(
215
+ x,
216
+ layer.w13_weight,
217
+ layer.w2_weight,
218
+ router_logits,
219
+ activation=layer.activation,
220
+ global_num_experts=layer.global_num_experts,
221
+ expert_map=layer.expert_map,
222
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
223
+ num_expert_group=layer.num_expert_group,
224
+ topk_group=layer.topk_group,
225
+ e_score_correction_bias=layer.e_score_correction_bias,
226
+ routed_scaling_factor=layer.routed_scaling_factor,
227
+ )