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, 06:15 UTC
[CPU] Add FP8 W8A16 linear support (#41186)
Signed-off-by: yuwenzho <yuwen.zhou@intel.com>
Yuwen Zhou committed
May 6, 2026 at 15:05 UTC
809b98e5b71482e52ee0ecbf140084aae49f3652
9 files changed
+331
-3
.buildkite/hardware_tests/cpu.yaml
+4
-2
@@ -14,13 +14,15 @@ steps:
14
- tests/kernels/moe/test_cpu_fused_moe.py
15
- tests/kernels/test_onednn.py
16
- tests/kernels/test_awq_int4_to_int8.py
17
+ - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
18
commands:
19
- |
19
- bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
20
+ bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
21
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
22
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
23
pytest -x -v -s tests/kernels/test_onednn.py
23
- pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py"
24
+ pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
25
+ pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
26
27
- label: CPU-Compatibility Tests
28
depends_on: []
csrc/cpu/sgl-kernels/gemm_fp8.cpp
+1
-1
@@ -447,7 +447,7 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
447
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
448
449
at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
450
- std::vector<int64_t> block_size, std::optional<at::Tensor>& bias,
450
+ std::vector<int64_t> block_size, const std::optional<at::Tensor>& bias,
451
at::ScalarType out_dtype, bool is_vnni) {
452
RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, block_size, bias}));
453
csrc/cpu/torch_bindings.cpp
+14
@@ -77,6 +77,13 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
77
const std::optional<at::Tensor>& bias,
78
at::ScalarType out_dtype, bool is_vnni);
79
80
+// Adapted from sglang: FP8 W8A16 kernel
81
+at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
82
+ at::Tensor& scales2,
83
+ std::vector<int64_t> block_size,
84
+ const std::optional<at::Tensor>& bias,
85
+ at::ScalarType out_dtype, bool is_vnni);
86
+
87
// Adapted from sglang: INT4 W4A8 kernels
88
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
89
at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
@@ -376,6 +383,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
383
"int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
384
"Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
385
ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
386
+
387
+ // Adapted from sglang: FP8 W8A16 kernel
388
+ ops.def(
389
+ "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
390
+ "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
391
+ "bool is_vnni) -> Tensor");
392
+ ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
393
#endif
394
395
// CPU attention kernels
tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
new
+162
@@ -0,0 +1,162 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu).
4
+
5
+Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`.
6
+"""
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm import _custom_ops as ops
12
+from vllm.platforms import current_platform
13
+
14
+if not current_platform.is_cpu():
15
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
16
+
17
+if not ops._supports_cpu_fp8_w8a16:
18
+ pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
19
+
20
+BLOCK_SIZE = [128, 128]
21
+
22
+
23
+def cdiv(a: int, b: int) -> int:
24
+ return -(a // -b)
25
+
26
+
27
+def quantize_weight_block_fp8(
28
+ weight: torch.Tensor,
29
+ block_size: list[int],
30
+) -> tuple[torch.Tensor, torch.Tensor]:
31
+ """Quantize weight [N, K] to FP8 with block scales.
32
+
33
+ Returns:
34
+ fp8_weight: [N, K] float8_e4m3fn
35
+ scales: [n_tiles, k_tiles] float32
36
+ """
37
+ N, K = weight.shape
38
+ block_n, block_k = block_size
39
+ fp8_max = torch.finfo(torch.float8_e4m3fn).max
40
+
41
+ n_tiles = cdiv(N, block_n)
42
+ k_tiles = cdiv(K, block_k)
43
+
44
+ # Pad for even blocking
45
+ pad_N = (block_n - (N % block_n)) % block_n
46
+ pad_K = (block_k - (K % block_k)) % block_k
47
+ if pad_N > 0 or pad_K > 0:
48
+ weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
49
+
50
+ # Reshape into blocks
51
+ w_blocks = weight.view(n_tiles, block_n, k_tiles, block_k)
52
+ w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous()
53
+
54
+ # Per-block scale
55
+ abs_max = w_blocks.abs().amax(dim=(-2, -1), keepdim=True)
56
+ scales = abs_max / fp8_max
57
+ scales = torch.where(scales == 0, torch.ones_like(scales), scales)
58
+
59
+ # Quantize
60
+ q_fp8 = (w_blocks / scales).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
61
+
62
+ # Reshape back
63
+ fp8_weight = (
64
+ q_fp8.permute(0, 2, 1, 3)
65
+ .contiguous()
66
+ .view(N + pad_N, K + pad_K)[:N, :K]
67
+ .contiguous()
68
+ )
69
+
70
+ scales = scales.view(n_tiles, k_tiles)
71
+ return fp8_weight, scales
72
+
73
+
74
+def dequant_weight_block_fp8(
75
+ fp8_weight: torch.Tensor,
76
+ scales: torch.Tensor,
77
+ block_size: list[int],
78
+ out_dtype: torch.dtype,
79
+) -> torch.Tensor:
80
+ """Dequantize FP8 weight back to float for reference computation."""
81
+ N, K = fp8_weight.shape
82
+ block_n, block_k = block_size
83
+ n_tiles, k_tiles = scales.shape
84
+
85
+ pad_N = (block_n - (N % block_n)) % block_n
86
+ pad_K = (block_k - (K % block_k)) % block_k
87
+ if pad_N > 0 or pad_K > 0:
88
+ fp8_padded = torch.nn.functional.pad(fp8_weight.float(), (0, pad_K, 0, pad_N))
89
+ else:
90
+ fp8_padded = fp8_weight.float()
91
+
92
+ w_blocks = fp8_padded.view(n_tiles, block_n, k_tiles, block_k)
93
+ w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous()
94
+ dq = w_blocks * scales.view(n_tiles, k_tiles, 1, 1)
95
+ dq = dq.permute(0, 2, 1, 3).contiguous().view(N + pad_N, K + pad_K)
96
+ return dq[:N, :K].to(out_dtype)
97
+
98
+
99
+def ref_fp8_block_scaled_mm(
100
+ x: torch.Tensor,
101
+ fp8_weight: torch.Tensor,
102
+ scales: torch.Tensor,
103
+ block_size: list[int],
104
+ bias: torch.Tensor | None,
105
+ out_dtype: torch.dtype,
106
+) -> torch.Tensor:
107
+ """Reference: dequant FP8→float32, matmul in float32, cast to out_dtype."""
108
+ w_dq = dequant_weight_block_fp8(fp8_weight, scales, block_size, torch.float32)
109
+ out = torch.mm(x.float(), w_dq.t())
110
+ if bias is not None:
111
+ out = out + bias.float()
112
+ return out.to(out_dtype)
113
+
114
+
115
+# ---------------------------------------------------------------------------
116
+# Test parameters
117
+# ---------------------------------------------------------------------------
118
+M_SIZES = [1, 4, 16, 64, 128]
119
+# (N, K) — weight shape is [N, K], output has N columns.
120
+NK_SIZES = [
121
+ (128, 256),
122
+ (256, 512),
123
+ (512, 1024),
124
+ (1024, 2048),
125
+ (5120, 5120),
126
+ (17408, 5120),
127
+ (5120, 17408),
128
+]
129
+
130
+
131
+@pytest.mark.parametrize("M", M_SIZES)
132
+@pytest.mark.parametrize("N,K", NK_SIZES)
133
+@pytest.mark.parametrize("use_bias", [False, True])
134
+def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool):
135
+ """fp8_scaled_mm_cpu correctness against float reference."""
136
+ torch.manual_seed(42)
137
+ out_dtype = torch.bfloat16
138
+ block_size = BLOCK_SIZE
139
+
140
+ x = torch.randn(M, K, dtype=out_dtype) / (K**0.5)
141
+ w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5)
142
+ fp8_weight, scales = quantize_weight_block_fp8(w_f32, block_size)
143
+
144
+ bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None
145
+
146
+ ref_out = ref_fp8_block_scaled_mm(
147
+ x, fp8_weight, scales, block_size, bias, out_dtype
148
+ )
149
+
150
+ packed_weight = torch.ops._C.convert_weight_packed(fp8_weight)
151
+ kernel_out = ops.fp8_scaled_mm_cpu(
152
+ x,
153
+ packed_weight,
154
+ scales,
155
+ block_size,
156
+ bias,
157
+ out_dtype,
158
+ True,
159
+ )
160
+
161
+ assert kernel_out.dtype == out_dtype
162
+ torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
tests/quantization/test_cpu_wna16.py
+1
@@ -13,6 +13,7 @@ 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
17
]
18
DTYPE = ["bfloat16"]
19
vllm/_custom_ops.py
+33
@@ -3219,6 +3219,39 @@ if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
3219
3220
_supports_cpu_w4a8_int8 = bool(hasattr(torch.ops._C, "convert_weight_packed_scale_zp"))
3221
3222
+if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"):
3223
+
3224
+ @register_fake("_C::fp8_scaled_mm_cpu")
3225
+ def fp8_scaled_mm_cpu_fake(
3226
+ mat1: torch.Tensor,
3227
+ mat2: torch.Tensor,
3228
+ scales2: torch.Tensor,
3229
+ block_size: list[int],
3230
+ bias: torch.Tensor | None,
3231
+ out_dtype: torch.dtype,
3232
+ is_vnni: bool,
3233
+ ) -> torch.Tensor:
3234
+ M = mat1.size(0)
3235
+ N = mat2.size(0)
3236
+ return torch.empty((M, N), dtype=out_dtype, device=mat1.device)
3237
+
3238
+
3239
+_supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu"))
3240
+
3241
+
3242
+def fp8_scaled_mm_cpu(
3243
+ mat1: torch.Tensor,
3244
+ mat2: torch.Tensor,
3245
+ scales2: torch.Tensor,
3246
+ block_size: list[int],
3247
+ bias: torch.Tensor | None,
3248
+ out_dtype: torch.dtype,
3249
+ is_vnni: bool,
3250
+) -> torch.Tensor:
3251
+ return torch.ops._C.fp8_scaled_mm_cpu(
3252
+ mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni
3253
+ )
3254
+
3255
3256
class CPUDNNLGEMMHandler:
3257
def __init__(self) -> None:
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -110,6 +110,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
110
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
111
)
112
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
113
+ CPUFp8BlockScaledMMKernel,
114
CPUInt8ScaledMMLinearKernel,
115
)
116
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -199,6 +200,9 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
200
AiterFp8BlockScaledMMKernel,
201
TritonFp8BlockScaledMMKernel,
202
],
203
+ PlatformEnum.CPU: [
204
+ CPUFp8BlockScaledMMKernel,
205
+ ],
206
}
207
208
_POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+2
@@ -8,6 +8,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo
8
Fp8BlockScaledMMLinearKernel,
9
)
10
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
11
+ CPUFp8BlockScaledMMKernel,
12
CPUInt8ScaledMMLinearKernel,
13
)
14
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -58,4 +59,5 @@ __all__ = [
59
"ROCmFP8ScaledMMLinearKernel",
60
"TritonInt8ScaledMMLinearKernel",
61
"Fp8BlockScaledMMLinearKernel",
62
+ "CPUFp8BlockScaledMMKernel",
63
]
vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+110
@@ -14,6 +14,10 @@ from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
14
from vllm.platforms import current_platform
15
from vllm.platforms.interface import CpuArchEnum
16
17
+from .BlockScaledMMLinearKernel import (
18
+ Fp8BlockScaledMMLinearKernel,
19
+ FP8ScaledMMLinearLayerConfig,
20
+)
21
from .ScaledMMLinearKernel import (
22
Int8ScaledMMLinearKernel,
23
Int8ScaledMMLinearLayerConfig,
@@ -215,3 +219,109 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
219
x.dtype,
220
True,
221
)
222
+
223
+
224
+class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
225
+ """FP8 W8A16 block-quantized GEMM via AMX BRGEMM on CPU."""
226
+
227
+ # Input stays BF16 — no FP8 activation quantization.
228
+ apply_input_quant = False
229
+
230
+ @classmethod
231
+ def is_supported(
232
+ cls, compute_capability: int | None = None
233
+ ) -> tuple[bool, str | None]:
234
+ if not current_platform.is_cpu():
235
+ return False, "requires CPU platform."
236
+ if not torch.cpu._is_amx_tile_supported():
237
+ return False, "requires AMX tile support (Sapphire Rapids or newer)."
238
+ if not ops._supports_cpu_fp8_w8a16:
239
+ return False, "fp8_scaled_mm_cpu op not available."
240
+ return True, None
241
+
242
+ @classmethod
243
+ def can_implement(
244
+ cls, config: FP8ScaledMMLinearLayerConfig
245
+ ) -> tuple[bool, str | None]:
246
+ # Validate weight block shape
247
+ weight_gs = config.weight_quant_key.scale.group_shape
248
+ if weight_gs.col <= 0 or weight_gs.col != 128:
249
+ return False, (
250
+ "CPU FP8 kernel requires K-dimension block size of 128, "
251
+ f"got {weight_gs.col}."
252
+ )
253
+ if weight_gs.row <= 0 or weight_gs.row % 32 != 0:
254
+ return False, (
255
+ "CPU FP8 kernel requires N-dimension block size to be "
256
+ f"a positive multiple of 32, got {weight_gs.row}."
257
+ )
258
+ if config.out_dtype not in (torch.bfloat16, torch.float32):
259
+ return False, "Only bfloat16/float32 output dtype supported."
260
+ return True, None
261
+
262
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
263
+ # Skip the base class process (FP8 padding / fnuz normalization)
264
+ # which is GPU-oriented. Instead, VNNI-prepack weights for AMX.
265
+ params = self._get_layer_params(layer)
266
+ packed_weight = torch.ops._C.convert_weight_packed(params.weight)
267
+ replace_parameter(
268
+ layer,
269
+ params.WEIGHT,
270
+ torch.nn.Parameter(packed_weight, requires_grad=False),
271
+ )
272
+
273
+ # Re-wrap scale as a plain Parameter so the kernel can read it
274
+ # without weight-loader metadata interfering.
275
+ scale_attr = (
276
+ params.WEIGHT_SCALE_INV
277
+ if params.weight_scale_inv is not None
278
+ else params.WEIGHT_SCALE
279
+ )
280
+ weight_scale = (
281
+ params.weight_scale_inv
282
+ if params.weight_scale_inv is not None
283
+ else params.weight_scale
284
+ )
285
+ assert weight_scale is not None
286
+ replace_parameter(
287
+ layer,
288
+ scale_attr,
289
+ torch.nn.Parameter(weight_scale.data, requires_grad=False),
290
+ )
291
+
292
+ def apply_weights(
293
+ self,
294
+ layer: torch.nn.Module,
295
+ x: torch.Tensor,
296
+ bias: torch.Tensor | None = None,
297
+ **kwargs,
298
+ ) -> torch.Tensor:
299
+ params = self._get_layer_params(layer)
300
+ weight_scale = (
301
+ params.weight_scale_inv
302
+ if params.weight_scale_inv is not None
303
+ else params.weight_scale
304
+ )
305
+
306
+ x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
307
+ out = torch.ops._C.fp8_scaled_mm_cpu(
308
+ x_2d,
309
+ params.weight,
310
+ weight_scale,
311
+ list(self.weight_group_shape),
312
+ bias,
313
+ x.dtype,
314
+ True, # is_vnni (weight already prepacked)
315
+ )
316
+ return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
317
+
318
+ def apply_block_scaled_mm(
319
+ self,
320
+ A: torch.Tensor,
321
+ B: torch.Tensor,
322
+ As: torch.Tensor,
323
+ Bs: torch.Tensor,
324
+ ) -> torch.Tensor:
325
+ raise NotImplementedError(
326
+ "CPUFp8BlockScaledMMKernel overrides apply_weights directly."
327
+ )