[Bugfix][CPU] Add per-tensor FP8 W8A16 kernel to fix Ministral crash on CPU (#56985)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Sep 16, 2026 at 15:49 UTC
b68408043d80faf9c99fc16c6aec1174336e9626
6 files changed
+233
-7
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+15
-6
@@ -83,14 +83,23 @@ rm -f "$build_log"
83
84
# Run the image, setting --shm-size=4g for tensor parallel. Default to
85
# HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network;
86
-# retry once online if the cache is missing something. vllm's get_config()
87
-# wraps the raw huggingface_hub offline-mode errors in a generic ValueError
88
-# (see transformers_utils/config.py), so match that message too or the
89
-# fallback never triggers for a config-lookup cache miss.
90
-OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified'
86
+# retry once online if the cache is missing something. vllm wraps offline
87
+# cache-miss errors in ways that don't preserve the raw huggingface_hub
88
+# exception name: get_config() re-raises as a generic ValueError
89
+# (transformers_utils/config.py), and default_loader.py raises its own
90
+# RuntimeError when a snapshot dir has a config but no weight files. Match
91
+# those messages too or the fallback never triggers for those cache misses.
92
+#
93
+# Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via
94
+# vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets
95
+# by default), a directory distinct from the HF hub cache above. Without a
96
+# persistent mount there, those assets never survive past one container's
97
+# lifetime, so every run re-fetches them from the network on the "offline"
98
+# attempt and unconditionally falls back to the online retry.
99
+OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with'
100
run_test() {
101
local hf_offline=$1
93
- docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \
102
+ docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface -v ~/.cache/vllm:/root/.cache/vllm --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 -e VLLM_CPU_ATTN_SPLIT_KV=0 -e HF_HUB_OFFLINE="$hf_offline" -e HF_DATASETS_OFFLINE="$hf_offline" --shm-size=4g "$IMAGE_NAME" \
103
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
104
}
105
tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
+113
@@ -9,6 +9,15 @@ import pytest
9
import torch
10
11
from vllm import _custom_ops as ops
12
+from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
13
+ CPUFp8PerTensorScaledMMLinearKernel,
14
+)
15
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
16
+ FP8ScaledMMLinearLayerConfig,
17
+)
18
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
+ kFp8StaticTensorSym,
20
+)
21
from vllm.platforms import current_platform
22
23
if not current_platform.is_cpu():
@@ -17,6 +26,9 @@ if not current_platform.is_cpu():
26
if not ops._supports_cpu_fp8_w8a16:
27
pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
28
29
+if not torch.cpu._is_amx_tile_supported():
30
+ pytest.skip("requires AMX tile support", allow_module_level=True)
31
+
32
BLOCK_SIZE = [128, 128]
33
34
@@ -160,3 +172,104 @@ def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool):
172
173
assert kernel_out.dtype == out_dtype
174
torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
175
+
176
+
177
+def quantize_weight_per_tensor_fp8(
178
+ weight: torch.Tensor,
179
+) -> tuple[torch.Tensor, torch.Tensor]:
180
+ """Quantize weight [N, K] to FP8 with a single per-tensor scale."""
181
+ fp8_max = torch.finfo(torch.float8_e4m3fn).max
182
+ scale = weight.abs().amax() / fp8_max
183
+ scale = torch.where(scale == 0, torch.ones_like(scale), scale)
184
+ q = (weight / scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
185
+ return q, scale
186
+
187
+
188
+def ref_fp8_per_tensor_scaled_mm(
189
+ x: torch.Tensor,
190
+ fp8_weight: torch.Tensor,
191
+ scale: torch.Tensor,
192
+ bias: torch.Tensor | None,
193
+ out_dtype: torch.dtype,
194
+) -> torch.Tensor:
195
+ w_dq = fp8_weight.float() * scale
196
+ out = torch.mm(x.float(), w_dq.t())
197
+ if bias is not None:
198
+ out = out + bias.float()
199
+ return out.to(out_dtype)
200
+
201
+
202
+NK_SIZES_PER_TENSOR = [
203
+ (32, 64),
204
+ (5120, 5120),
205
+]
206
+
207
+
208
+@pytest.mark.parametrize("M", [1, 64])
209
+@pytest.mark.parametrize("N,K", NK_SIZES_PER_TENSOR)
210
+@pytest.mark.parametrize("use_bias", [False, True])
211
+def test_cpu_fp8_per_tensor_scaled_mm_kernel(
212
+ M: int, N: int, K: int, use_bias: bool, default_vllm_config
213
+):
214
+ """CPUFp8PerTensorScaledMMLinearKernel correctness against float reference.
215
+
216
+ Exercises the full kernel class (not just the raw op), including the
217
+ weight-orientation fixup: Fp8LinearMethod stores `layer.weight` as
218
+ [K, N] (torch._scaled_mm convention) before calling
219
+ process_weights_after_loading, so the kernel must transpose back to
220
+ [N, K] before VNNI-packing.
221
+ """
222
+ torch.manual_seed(0)
223
+ out_dtype = torch.bfloat16
224
+
225
+ x = torch.randn(M, K, dtype=out_dtype) / (K**0.5)
226
+ w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5)
227
+ fp8_weight, scale = quantize_weight_per_tensor_fp8(w_f32)
228
+ bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None
229
+
230
+ ref_out = ref_fp8_per_tensor_scaled_mm(x, fp8_weight, scale, bias, out_dtype)
231
+
232
+ config = FP8ScaledMMLinearLayerConfig(
233
+ weight_quant_key=kFp8StaticTensorSym,
234
+ activation_quant_key=kFp8StaticTensorSym,
235
+ weight_shape=(N, K),
236
+ input_dtype=out_dtype,
237
+ out_dtype=out_dtype,
238
+ )
239
+ kernel = CPUFp8PerTensorScaledMMLinearKernel(
240
+ config,
241
+ layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
242
+ )
243
+
244
+ layer = torch.nn.Module()
245
+ # Fp8LinearMethod stores the weight transposed to [K, N] before calling
246
+ # process_weights_after_loading.
247
+ layer.register_parameter(
248
+ "weight", torch.nn.Parameter(fp8_weight.t().contiguous(), requires_grad=False)
249
+ )
250
+ layer.register_parameter(
251
+ "weight_scale", torch.nn.Parameter(scale.clone(), requires_grad=False)
252
+ )
253
+ kernel.process_weights_after_loading(layer)
254
+
255
+ kernel_out = kernel.apply_weights(layer, x, bias)
256
+
257
+ assert kernel_out.dtype == out_dtype
258
+ torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
259
+
260
+
261
+@pytest.mark.parametrize("n", [16, 48])
262
+def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int):
263
+ """The AMX tinygemm kernel tiles N in chunks of 32 and cannot handle a
264
+ remainder tile smaller than that, so can_implement must reject any N
265
+ that isn't a multiple of 32 rather than let it crash the process.
266
+ """
267
+ config = FP8ScaledMMLinearLayerConfig(
268
+ weight_quant_key=kFp8StaticTensorSym,
269
+ activation_quant_key=kFp8StaticTensorSym,
270
+ weight_shape=(n, 64),
271
+ input_dtype=torch.bfloat16,
272
+ out_dtype=torch.bfloat16,
273
+ )
274
+ supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config)
275
+ assert not supported
tests/v1/e2e/test_cpu_spec_decode.py
+1
-1
@@ -131,7 +131,7 @@ SMOKE_CONFIGS = [
131
id="dflash",
132
),
133
pytest.param(
134
- "Qwen/Qwen3.5-0.8B-Base",
134
+ "Qwen/Qwen3.5-0.8B",
135
{
136
"method": "mtp",
137
"num_speculative_tokens": 3,
vllm/model_executor/kernels/linear/__init__.py
+2
@@ -187,6 +187,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
187
)
188
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
189
CPUFp8BlockScaledMMKernel,
190
+ CPUFp8PerTensorScaledMMLinearKernel,
191
CPUInt8ScaledMMLinearKernel,
192
)
193
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -438,6 +439,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
439
ChannelWiseTorchFP8ScaledMMLinearKernel,
440
],
441
PlatformEnum.CPU: [
442
+ CPUFp8PerTensorScaledMMLinearKernel,
443
PerTensorTorchFP8ScaledMMLinearKernel,
444
ChannelWiseTorchFP8ScaledMMLinearKernel,
445
],
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+2
@@ -9,6 +9,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo
9
)
10
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
11
CPUFp8BlockScaledMMKernel,
12
+ CPUFp8PerTensorScaledMMLinearKernel,
13
CPUInt8ScaledMMLinearKernel,
14
)
15
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -67,5 +68,6 @@ __all__ = [
68
"ZentorchInt8ScaledMMLinearKernel",
69
"Fp8BlockScaledMMLinearKernel",
70
"CPUFp8BlockScaledMMKernel",
71
+ "CPUFp8PerTensorScaledMMLinearKernel",
72
"XPUFp8BlockScaledMMKernel",
73
]
vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+100
@@ -21,10 +21,14 @@ from .BlockScaledMMLinearKernel import (
21
FP8ScaledMMLinearLayerConfig,
22
)
23
from .ScaledMMLinearKernel import (
24
+ FP8ScaledMMLinearKernel,
25
Int8ScaledMMLinearKernel,
26
Int8ScaledMMLinearLayerConfig,
27
)
28
29
+# BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size.
30
+_BLOCK_SIZE_K = 128
31
+
32
33
class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
34
@classmethod
@@ -330,3 +334,99 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
334
raise NotImplementedError(
335
"CPUFp8BlockScaledMMKernel overrides apply_weights directly."
336
)
337
+
338
+
339
+class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel):
340
+ """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU.
341
+
342
+ Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single
343
+ synthetic block spanning the whole weight, so activations stay BF16/FP32
344
+ — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel.
345
+ """
346
+
347
+ @classmethod
348
+ def is_supported(
349
+ cls, compute_capability: int | None = None
350
+ ) -> tuple[bool, str | None]:
351
+ if not current_platform.is_cpu():
352
+ return False, "requires CPU platform."
353
+ if not torch.cpu._is_amx_tile_supported():
354
+ return False, "requires AMX tile support (Sapphire Rapids or newer)."
355
+ if not ops._supports_cpu_fp8_w8a16:
356
+ return False, "fp8_scaled_mm_cpu op not available."
357
+ return True, None
358
+
359
+ @classmethod
360
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
361
+ if not c.weight_quant_key.scale.group_shape.is_per_tensor():
362
+ return False, "requires a per-tensor weight scale."
363
+ if not c.activation_quant_key.scale.group_shape.is_per_tensor():
364
+ return False, "requires a per-tensor activation scale."
365
+ if c.out_dtype not in (torch.bfloat16, torch.float32):
366
+ return False, "Only bfloat16/float32 output dtype supported."
367
+ n = c.weight_shape[0]
368
+ if n % 32 != 0:
369
+ # The AMX tinygemm kernel tiles N in chunks of 32 and cannot
370
+ # handle a remainder tile smaller than that.
371
+ return False, f"requires weight output dim (N={n}) to be a multiple of 32."
372
+ return True, None
373
+
374
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
375
+ w_name, w_s_name, _, _ = self.layer_param_names
376
+ # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style
377
+ # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as
378
+ # produced by CUDA nn.Linear-style checkpoints.
379
+ weight = getattr(layer, w_name).t().contiguous()
380
+ n, k = weight.shape
381
+
382
+ packed_weight = torch.ops._C.convert_weight_packed(weight)
383
+ replace_parameter(
384
+ layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False)
385
+ )
386
+
387
+ # Synthesize a single-block "block scale" tensor from the scalar
388
+ # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is.
389
+ weight_scale = getattr(layer, w_s_name)
390
+ self._block_size_n = -(-n // 32) * 32 # round up to a multiple of 32
391
+ num_k_blocks = -(-k // _BLOCK_SIZE_K)
392
+ block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous()
393
+ replace_parameter(
394
+ layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False)
395
+ )
396
+
397
+ def apply_weights(
398
+ self,
399
+ layer: torch.nn.Module,
400
+ x: torch.Tensor,
401
+ bias: torch.Tensor | None = None,
402
+ ) -> torch.Tensor:
403
+ w_name, w_s_name, _, _ = self.layer_param_names
404
+ weight = getattr(layer, w_name)
405
+ weight_scale = getattr(layer, w_s_name)
406
+
407
+ x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
408
+ out = torch.ops._C.fp8_scaled_mm_cpu(
409
+ x_2d,
410
+ weight,
411
+ weight_scale,
412
+ [self._block_size_n, _BLOCK_SIZE_K],
413
+ bias,
414
+ x.dtype,
415
+ True, # is_vnni (weight already prepacked)
416
+ )
417
+ return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
418
+
419
+ def apply_scaled_mm(
420
+ self,
421
+ *,
422
+ A: torch.Tensor,
423
+ B: torch.Tensor,
424
+ out_dtype: torch.dtype,
425
+ As: torch.Tensor,
426
+ Bs: torch.Tensor,
427
+ bias: torch.Tensor | None,
428
+ output_shape: list,
429
+ ) -> torch.Tensor:
430
+ raise NotImplementedError(
431
+ "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly."
432
+ )