A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 05:00 UTC
[CPU] Add FP8 W8A16 MoE support (#41314)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com>
Yuwen Zhou committed
May 7, 2026 at 14:17 UTC
713b28bd0bf4d7f8226ee20940e3ae171fabd09b
6 files changed
+498
-1
.buildkite/hardware_tests/cpu.yaml
+2
@@ -12,6 +12,7 @@ steps:
12
- vllm/_custom_ops.py
13
- tests/kernels/attention/test_cpu_attn.py
14
- tests/kernels/moe/test_cpu_fused_moe.py
15
+ - tests/kernels/moe/test_cpu_fp8_fused_moe.py
16
- tests/kernels/test_onednn.py
17
- tests/kernels/test_awq_int4_to_int8.py
18
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -20,6 +21,7 @@ steps:
21
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
22
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
23
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
24
+ pytest -x -v -s tests/kernels/moe/test_cpu_fp8_fused_moe.py
25
pytest -x -v -s tests/kernels/test_onednn.py
26
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
27
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
tests/kernels/moe/test_cpu_fp8_fused_moe.py
new
+258
@@ -0,0 +1,258 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for CPU FP8 W8A16 block-scaled fused MoE kernel."""
4
+
5
+import math
6
+import sys
7
+
8
+import pytest
9
+import torch
10
+import torch.nn.functional as F
11
+
12
+from vllm.platforms import current_platform
13
+from vllm.utils.torch_utils import set_random_seed
14
+
15
+if not current_platform.is_cpu():
16
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
17
+
18
+import vllm._custom_ops as ops # noqa: E402
19
+
20
+if not hasattr(torch.ops._C, "fused_experts_cpu"):
21
+ pytest.skip("fused_experts_cpu op not available", allow_module_level=True)
22
+
23
+
24
+BLOCK_SIZE = [128, 128] # [block_n, block_k]
25
+
26
+# FP8 weight generation parameters
27
+_FP8_INFO = torch.finfo(torch.float8_e4m3fn)
28
+FP8_SCALE = _FP8_INFO.max # 448.0
29
+FACTOR_FOR_SCALE = 1e-3
30
+
31
+# Tolerance for FP8 W8A16
32
+FP8_W8A16_ATOL = 1e-2
33
+FP8_W8A16_RTOL = 1e-2
34
+
35
+
36
+def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
37
+ d = x.shape[-1] // 2
38
+ return F.silu(x[..., :d]) * x[..., d:]
39
+
40
+
41
+def _block_dequant_weight(
42
+ weight: torch.Tensor,
43
+ scales: torch.Tensor,
44
+ block_size: list[int],
45
+) -> torch.Tensor:
46
+ """Block-dequantize FP8 weight [E, N, K] → float [E, N, K].
47
+
48
+ Each (block_n × block_k) tile is multiplied by its per-block scale.
49
+ """
50
+ E, N, K = weight.shape
51
+ block_n, block_k = block_size
52
+ pad_N = (block_n - N % block_n) % block_n
53
+ pad_K = (block_k - K % block_k) % block_k
54
+
55
+ if pad_N > 0 or pad_K > 0:
56
+ weight = F.pad(weight, (0, pad_K, 0, pad_N))
57
+
58
+ n_tiles = math.ceil(N / block_n)
59
+ k_tiles = math.ceil(K / block_k)
60
+
61
+ weight_block = (
62
+ weight.view(E, n_tiles, block_n, k_tiles, block_k)
63
+ .permute(0, 1, 3, 2, 4)
64
+ .float()
65
+ .contiguous()
66
+ )
67
+ weight_scaled = (
68
+ (weight_block * scales.view(E, n_tiles, k_tiles, 1, 1))
69
+ .permute(0, 1, 3, 2, 4)
70
+ .contiguous()
71
+ )
72
+ if pad_N > 0 or pad_K > 0:
73
+ weight_scaled = weight_scaled.view(E, N + pad_N, K + pad_K)
74
+ weight_scaled = weight_scaled[..., :N, :K].contiguous()
75
+ else:
76
+ weight_scaled = weight_scaled.view(E, N, K)
77
+ return weight_scaled
78
+
79
+
80
+def ref_w8a16_block_fp8_moe(
81
+ a: torch.Tensor,
82
+ w1: torch.Tensor,
83
+ w2: torch.Tensor,
84
+ w1_s: torch.Tensor,
85
+ w2_s: torch.Tensor,
86
+ topk_weight: torch.Tensor,
87
+ topk_ids: torch.Tensor,
88
+ block_size: list[int],
89
+) -> torch.Tensor:
90
+ """Reference FP8 W8A16 block-scaled fused MoE in pure torch.
91
+
92
+ Steps:
93
+ 1. Block-dequant FP8 weights → float
94
+ 2. For each expert: matmul → SiLU+Mul → matmul
95
+ 3. Weighted sum across top-k experts
96
+ """
97
+ B, D = a.shape
98
+ topk = topk_ids.size(1)
99
+
100
+ w1_dq = _block_dequant_weight(w1, w1_s, block_size)
101
+ w2_dq = _block_dequant_weight(w2, w2_s, block_size)
102
+
103
+ a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
104
+ out = torch.zeros(B * topk, w2_dq.shape[1], dtype=torch.float32)
105
+
106
+ topk_weight_flat = topk_weight.view(-1)
107
+ topk_ids_flat = topk_ids.view(-1)
108
+
109
+ for i in range(w1_dq.shape[0]):
110
+ mask = topk_ids_flat == i
111
+ if mask.sum():
112
+ ic0 = torch.matmul(a_exp[mask], w1_dq[i].transpose(0, 1))
113
+ ic1 = _silu_and_mul(ic0)
114
+ out[mask] = torch.matmul(ic1, w2_dq[i].transpose(0, 1))
115
+
116
+ return (
117
+ (out.view(B, -1, w2_dq.shape[1]) * topk_weight_flat.view(B, -1, 1))
118
+ .sum(dim=1)
119
+ .to(a.dtype)
120
+ )
121
+
122
+
123
+def _make_fp8_moe_weights(
124
+ E: int,
125
+ N: int,
126
+ K: int,
127
+ block_size: list[int],
128
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
129
+ """Generate random FP8 MoE weights with random block scales.
130
+
131
+ Weight generation follows SGLang: ``randn * FP8_SCALE`` → clamp → cast.
132
+ Scales are small random values (``FACTOR_FOR_SCALE``), independent of
133
+ the actual weight magnitudes — this is sufficient to test the kernel's
134
+ block-dequant + matmul correctness.
135
+
136
+ Returns: (w1, w2, w1_s, w2_s)
137
+ """
138
+ block_n, block_k = block_size
139
+
140
+ w1 = (
141
+ (torch.randn(E, 2 * N, K) * FP8_SCALE)
142
+ .clamp(min=-FP8_SCALE, max=FP8_SCALE)
143
+ .to(torch.float8_e4m3fn)
144
+ )
145
+ w2 = (
146
+ (torch.randn(E, K, N) * FP8_SCALE)
147
+ .clamp(min=-FP8_SCALE, max=FP8_SCALE)
148
+ .to(torch.float8_e4m3fn)
149
+ )
150
+
151
+ w1_s = (
152
+ torch.randn(E, math.ceil(2 * N / block_n), math.ceil(K / block_k))
153
+ * FACTOR_FOR_SCALE
154
+ )
155
+ w2_s = (
156
+ torch.randn(E, math.ceil(K / block_n), math.ceil(N / block_k))
157
+ * FACTOR_FOR_SCALE
158
+ )
159
+ return w1, w2, w1_s, w2_s
160
+
161
+
162
+def _prepack_experts(w: torch.Tensor) -> torch.Tensor:
163
+ """VNNI-prepack each expert's weight via ``convert_weight_packed``."""
164
+ return torch.stack(
165
+ [torch.ops._C.convert_weight_packed(w[e]) for e in range(w.shape[0])]
166
+ )
167
+
168
+
169
+NUM_TOKENS = [1, 2, 64, 121]
170
+# (M, intermediate_size N, hidden_size K, num_experts E, topk)
171
+MoE_CONFIGS = [
172
+ (256, 512, 8, 2),
173
+ (256, 512, 8, 4),
174
+ (512, 256, 8, 2),
175
+ (512, 256, 8, 4),
176
+ (512, 512, 8, 2),
177
+ (512, 512, 8, 4),
178
+ (768, 2048, 8, 2),
179
+ (768, 2048, 8, 4),
180
+ (768, 2048, 128, 8),
181
+]
182
+SEEDS = [0]
183
+
184
+
185
+@pytest.mark.parametrize("M", NUM_TOKENS)
186
+@pytest.mark.parametrize("N,K,E,topk", MoE_CONFIGS)
187
+@pytest.mark.parametrize("seed", SEEDS)
188
+def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
189
+ """Test fused_experts_cpu FP8 W8A16 against dequantised torch reference."""
190
+ set_random_seed(seed)
191
+
192
+ a = torch.randn(M, K, dtype=torch.bfloat16) / math.sqrt(K)
193
+ w1, w2, w1_s, w2_s = _make_fp8_moe_weights(E, N, K, BLOCK_SIZE)
194
+
195
+ score = torch.randn(M, E, dtype=torch.bfloat16)
196
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
197
+ topk_weight, topk_ids = torch.topk(score, topk)
198
+ topk_ids = topk_ids.to(torch.int32)
199
+
200
+ ref_out = ref_w8a16_block_fp8_moe(
201
+ a,
202
+ w1,
203
+ w2,
204
+ w1_s,
205
+ w2_s,
206
+ topk_weight,
207
+ topk_ids,
208
+ BLOCK_SIZE,
209
+ )
210
+
211
+ pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
212
+
213
+ # Test inplace=False against reference
214
+ out = ops.fused_experts_cpu(
215
+ a.clone(),
216
+ pw1,
217
+ pw2,
218
+ topk_weight,
219
+ topk_ids,
220
+ False,
221
+ False,
222
+ True,
223
+ w1_s,
224
+ w2_s,
225
+ BLOCK_SIZE,
226
+ None,
227
+ None,
228
+ True,
229
+ )
230
+ torch.testing.assert_close(
231
+ ref_out.bfloat16(),
232
+ out,
233
+ atol=FP8_W8A16_ATOL,
234
+ rtol=FP8_W8A16_RTOL,
235
+ )
236
+
237
+ # Test inplace=True produces identical output
238
+ out_inplace = ops.fused_experts_cpu(
239
+ a.clone(),
240
+ pw1,
241
+ pw2,
242
+ topk_weight,
243
+ topk_ids,
244
+ True,
245
+ False,
246
+ True,
247
+ w1_s,
248
+ w2_s,
249
+ BLOCK_SIZE,
250
+ None,
251
+ None,
252
+ True,
253
+ )
254
+ torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
255
+
256
+
257
+if __name__ == "__main__":
258
+ sys.exit(pytest.main([__file__, "-v"]))
tests/quantization/test_cpu_wna16.py
+2
-1
@@ -13,7 +13,8 @@ MODELS = [
13
"Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
14
"RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
15
"OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
16
- "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized
16
+ "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear
17
+ "Qwen/Qwen3-30B-A3B-FP8", # FP8 W8A16 block-quantized MoE
18
]
19
DTYPE = ["bfloat16"]
20
vllm/_custom_ops.py
+34
@@ -3172,6 +3172,40 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):
3172
return torch.empty_like(hidden_states)
3173
3174
3175
+def fused_experts_cpu(
3176
+ hidden_states: torch.Tensor,
3177
+ w1: torch.Tensor,
3178
+ w2: torch.Tensor,
3179
+ topk_weights: torch.Tensor,
3180
+ topk_ids: torch.Tensor,
3181
+ inplace: bool,
3182
+ use_int8_w8a8: bool,
3183
+ use_fp8_w8a16: bool,
3184
+ w1_scale: torch.Tensor | None,
3185
+ w2_scale: torch.Tensor | None,
3186
+ block_size: list[int] | None,
3187
+ a1_scale: torch.Tensor | None,
3188
+ a2_scale: torch.Tensor | None,
3189
+ is_vnni: bool,
3190
+) -> torch.Tensor:
3191
+ return torch.ops._C.fused_experts_cpu(
3192
+ hidden_states,
3193
+ w1,
3194
+ w2,
3195
+ topk_weights,
3196
+ topk_ids,
3197
+ inplace,
3198
+ use_int8_w8a8,
3199
+ use_fp8_w8a16,
3200
+ w1_scale,
3201
+ w2_scale,
3202
+ block_size,
3203
+ a1_scale,
3204
+ a2_scale,
3205
+ is_vnni,
3206
+ )
3207
+
3208
+
3209
if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
3210
3211
@register_fake("_C::int8_scaled_mm_with_quant")
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
new
+175
@@ -0,0 +1,175 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU FP8 W8A16 block-quantized fused MoE experts."""
4
+
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
+from vllm.model_executor.layers.fused_moe.config import (
10
+ FusedMoEConfig,
11
+ FusedMoEParallelConfig,
12
+ FusedMoEQuantConfig,
13
+ RoutingMethodType,
14
+)
15
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
+ QuantKey,
17
+ kFp8Dynamic128Sym,
18
+ kFp8Static128BlockSym,
19
+)
20
+from vllm.platforms import current_platform
21
+
22
+
23
+def prepare_fp8_moe_layer_for_cpu(
24
+ w13: torch.Tensor,
25
+ w2: torch.Tensor,
26
+) -> tuple[torch.Tensor, torch.Tensor]:
27
+ """VNNI-prepack FP8 MoE weights for CPU kernel."""
28
+ num_experts = w13.size(0)
29
+ packed_w13_list = []
30
+ packed_w2_list = []
31
+ for i in range(num_experts):
32
+ packed_w13_list.append(torch.ops._C.convert_weight_packed(w13[i]))
33
+ packed_w2_list.append(torch.ops._C.convert_weight_packed(w2[i]))
34
+ packed_w13 = torch.stack(packed_w13_list)
35
+ packed_w2 = torch.stack(packed_w2_list)
36
+ return packed_w13, packed_w2
37
+
38
+
39
+class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
40
+ """CPU FP8 W8A16 block-quantized monolithic MoE experts."""
41
+
42
+ def __init__(
43
+ self,
44
+ moe_config: FusedMoEConfig,
45
+ quant_config: FusedMoEQuantConfig,
46
+ ):
47
+ super().__init__(
48
+ moe_config,
49
+ quant_config,
50
+ )
51
+
52
+ @property
53
+ def expects_unquantized_inputs(self) -> bool:
54
+ return True
55
+
56
+ @staticmethod
57
+ def activation_format() -> mk.FusedMoEActivationFormat:
58
+ return mk.FusedMoEActivationFormat.Standard
59
+
60
+ @staticmethod
61
+ def _supports_current_device() -> bool:
62
+ return current_platform.is_cpu()
63
+
64
+ @staticmethod
65
+ def _supports_no_act_and_mul() -> bool:
66
+ return False
67
+
68
+ @staticmethod
69
+ def _supports_activation(activation: MoEActivation) -> bool:
70
+ return activation == MoEActivation.SILU
71
+
72
+ @staticmethod
73
+ def _supports_parallel_config(
74
+ moe_parallel_config: FusedMoEParallelConfig,
75
+ ) -> bool:
76
+ return True
77
+
78
+ @staticmethod
79
+ def _supports_quant_scheme(
80
+ weight_key: QuantKey | None,
81
+ activation_key: QuantKey | None,
82
+ ) -> bool:
83
+ SUPPORTED_W_A = [
84
+ (kFp8Static128BlockSym, kFp8Dynamic128Sym),
85
+ ]
86
+ return (weight_key, activation_key) in SUPPORTED_W_A
87
+
88
+ @staticmethod
89
+ def _supports_routing_method(
90
+ routing_method: RoutingMethodType,
91
+ weight_key: QuantKey | None,
92
+ activation_key: QuantKey | None,
93
+ ) -> bool:
94
+ return routing_method in [
95
+ RoutingMethodType.Default,
96
+ RoutingMethodType.Renormalize,
97
+ RoutingMethodType.RenormalizeNaive,
98
+ ]
99
+
100
+ @staticmethod
101
+ def _supports_router_logits_dtype(
102
+ router_logits_dtype: torch.dtype | None,
103
+ routing_method: RoutingMethodType,
104
+ ) -> bool:
105
+ return True
106
+
107
+ def supports_expert_map(self) -> bool:
108
+ return False
109
+
110
+ def apply(
111
+ self,
112
+ hidden_states: torch.Tensor,
113
+ w1: torch.Tensor,
114
+ w2: torch.Tensor,
115
+ router_logits: torch.Tensor,
116
+ activation: MoEActivation,
117
+ global_num_experts: int,
118
+ expert_map: torch.Tensor | None,
119
+ a1q_scale: torch.Tensor | None,
120
+ apply_router_weight_on_input: bool,
121
+ # grouped topk + fused topk bias parameters
122
+ num_expert_group: int | None = None,
123
+ e_score_correction_bias: torch.Tensor | None = None,
124
+ routed_scaling_factor: float | None = None,
125
+ topk_group: int | None = None,
126
+ ) -> torch.Tensor:
127
+ from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
128
+ select_experts,
129
+ )
130
+
131
+ topk_weights, topk_ids = select_experts(
132
+ hidden_states=hidden_states,
133
+ router_logits=router_logits,
134
+ use_grouped_topk=num_expert_group is not None,
135
+ top_k=self.moe_config.experts_per_token,
136
+ renormalize=self.moe_config.routing_method
137
+ in (
138
+ RoutingMethodType.Renormalize,
139
+ RoutingMethodType.RenormalizeNaive,
140
+ ),
141
+ topk_group=topk_group,
142
+ num_expert_group=num_expert_group,
143
+ scoring_func="softmax",
144
+ routed_scaling_factor=(
145
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
146
+ ),
147
+ e_score_correction_bias=e_score_correction_bias,
148
+ )
149
+
150
+ block_shape = (
151
+ list(self.quant_config.block_shape)
152
+ if self.quant_config.block_shape
153
+ else (
154
+ [self.quant_config._w1.shape.row, self.quant_config._w1.shape.col]
155
+ if self.quant_config._w1.shape is not None
156
+ else None
157
+ )
158
+ )
159
+
160
+ return torch.ops._C.fused_experts_cpu(
161
+ hidden_states,
162
+ w1,
163
+ w2,
164
+ topk_weights,
165
+ topk_ids,
166
+ False, # inplace
167
+ False, # use_int8_w8a8
168
+ True, # use_fp8_w8a16
169
+ self.w1_scale, # w1_scale
170
+ self.w2_scale, # w2_scale
171
+ block_shape, # block_size
172
+ None, # a1_scale (W8A16: no activation quantization)
173
+ None, # a2_scale
174
+ True, # is_vnni
175
+ )
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+27
@@ -55,6 +55,7 @@ class Fp8MoeBackend(Enum):
55
VLLM_CUTLASS = "VLLM_CUTLASS"
56
BATCHED_VLLM_CUTLASS = "BATCHED_VLLM_CUTLASS"
57
XPU = "XPU"
58
+ CPU = "CPU"
59
60
61
def _get_priority_backends(
@@ -80,6 +81,7 @@ def _get_priority_backends(
81
Fp8MoeBackend.BATCHED_VLLM_CUTLASS,
82
Fp8MoeBackend.BATCHED_TRITON,
83
Fp8MoeBackend.XPU,
84
+ Fp8MoeBackend.CPU,
85
]
86
87
def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None:
@@ -102,6 +104,10 @@ def _get_priority_backends(
104
# move XPU backend to the front.
105
_move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.XPU)
106
107
+ if current_platform.is_cpu():
108
+ # CPU platform uses FP8 W8A16 fused MoE kernel.
109
+ _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.CPU)
110
+
111
return _AVAILABLE_BACKENDS
112
113
@@ -186,6 +192,13 @@ def backend_to_kernel_cls(
192
193
return [XPUExpertsFp8]
194
195
+ elif backend == Fp8MoeBackend.CPU:
196
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
197
+ CPUExpertsFp8,
198
+ )
199
+
200
+ return [CPUExpertsFp8]
201
+
202
else:
203
raise ValueError(f"Unknown FP8 MoE backend: {backend.value}")
204
@@ -472,6 +485,12 @@ def convert_to_fp8_moe_kernel_format(
485
)
486
487
w13, w2 = prepare_fp8_moe_layer_for_xpu(w13, w2)
488
+ elif fp8_backend == Fp8MoeBackend.CPU:
489
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
490
+ prepare_fp8_moe_layer_for_cpu,
491
+ )
492
+
493
+ w13, w2 = prepare_fp8_moe_layer_for_cpu(w13, w2)
494
else:
495
if fp8_backend not in [
496
Fp8MoeBackend.TRITON,
@@ -516,6 +535,14 @@ def make_fp8_moe_quant_config(
535
block_shape=block_shape,
536
)
537
538
+ # CPU is mixed precision W8A16 config.
539
+ if fp8_backend == Fp8MoeBackend.CPU:
540
+ return fp8_w8a16_moe_quant_config(
541
+ w1_scale=w1_scale,
542
+ w2_scale=w2_scale,
543
+ block_shape=block_shape,
544
+ )
545
+
546
# Flashinfer CUTLASS per-tensor uses single dq scale
547
# (alpha = w_scale * a_scale) and inverse a2 scale.
548
if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None: