[Bugfix][CPU] Add per-tensor FP8 W8A16 kernel to fix Ministral crash on CPU (#56985)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

.buildkite/scripts/hardware_ci/run-cpu-test.sh

@@ -83,14 +83,23 @@ rm -f "$build_log"
 
 # Run the image, setting --shm-size=4g for tensor parallel. Default to
 # HF_HUB_OFFLINE so a warm ~/.cache/huggingface doesn't hit the network;
-# retry once online if the cache is missing something. vllm's get_config()
-# wraps the raw huggingface_hub offline-mode errors in a generic ValueError
-# (see transformers_utils/config.py), so match that message too or the
-# fallback never triggers for a config-lookup cache miss.
-OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified'
+# retry once online if the cache is missing something. vllm wraps offline
+# cache-miss errors in ways that don't preserve the raw huggingface_hub
+# exception name: get_config() re-raises as a generic ValueError
+# (transformers_utils/config.py), and default_loader.py raises its own
+# RuntimeError when a snapshot dir has a config but no weight files. Match
+# those messages too or the fallback never triggers for those cache misses.
+#
+# Also mount ~/.cache/vllm: multimodal test fixtures (e.g. VideoAsset, via
+# vllm/assets/video.py) download into VLLM_ASSETS_CACHE (~/.cache/vllm/assets
+# by default), a directory distinct from the HF hub cache above. Without a
+# persistent mount there, those assets never survive past one container's
+# lifetime, so every run re-fetches them from the network on the "offline"
+# attempt and unconditionally falls back to the online retry.
+OFFLINE_RETRY_PATTERN='huggingface_hub\.errors\.(LocalEntryNotFoundError|OfflineModeIsEnabled)|Invalid repository ID or local directory specified|Cannot find any model weights with'
 run_test() {
     local hf_offline=$1
-    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" \
+    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" \
         timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
 }
 

tests/kernels/quantization/test_cpu_fp8_scaled_mm.py

@@ -9,6 +9,15 @@ import pytest
 import torch
 
 from vllm import _custom_ops as ops
+from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
+    CPUFp8PerTensorScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
+    FP8ScaledMMLinearLayerConfig,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8StaticTensorSym,
+)
 from vllm.platforms import current_platform
 
 if not current_platform.is_cpu():
@@ -17,6 +26,9 @@ if not current_platform.is_cpu():
 if not ops._supports_cpu_fp8_w8a16:
     pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
 
+if not torch.cpu._is_amx_tile_supported():
+    pytest.skip("requires AMX tile support", allow_module_level=True)
+
 BLOCK_SIZE = [128, 128]
 
 
@@ -160,3 +172,104 @@ def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool):
 
     assert kernel_out.dtype == out_dtype
     torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
+
+
+def quantize_weight_per_tensor_fp8(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize weight [N, K] to FP8 with a single per-tensor scale."""
+    fp8_max = torch.finfo(torch.float8_e4m3fn).max
+    scale = weight.abs().amax() / fp8_max
+    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
+    q = (weight / scale).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
+    return q, scale
+
+
+def ref_fp8_per_tensor_scaled_mm(
+    x: torch.Tensor,
+    fp8_weight: torch.Tensor,
+    scale: torch.Tensor,
+    bias: torch.Tensor | None,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    w_dq = fp8_weight.float() * scale
+    out = torch.mm(x.float(), w_dq.t())
+    if bias is not None:
+        out = out + bias.float()
+    return out.to(out_dtype)
+
+
+NK_SIZES_PER_TENSOR = [
+    (32, 64),
+    (5120, 5120),
+]
+
+
+@pytest.mark.parametrize("M", [1, 64])
+@pytest.mark.parametrize("N,K", NK_SIZES_PER_TENSOR)
+@pytest.mark.parametrize("use_bias", [False, True])
+def test_cpu_fp8_per_tensor_scaled_mm_kernel(
+    M: int, N: int, K: int, use_bias: bool, default_vllm_config
+):
+    """CPUFp8PerTensorScaledMMLinearKernel correctness against float reference.
+
+    Exercises the full kernel class (not just the raw op), including the
+    weight-orientation fixup: Fp8LinearMethod stores `layer.weight` as
+    [K, N] (torch._scaled_mm convention) before calling
+    process_weights_after_loading, so the kernel must transpose back to
+    [N, K] before VNNI-packing.
+    """
+    torch.manual_seed(0)
+    out_dtype = torch.bfloat16
+
+    x = torch.randn(M, K, dtype=out_dtype) / (K**0.5)
+    w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5)
+    fp8_weight, scale = quantize_weight_per_tensor_fp8(w_f32)
+    bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None
+
+    ref_out = ref_fp8_per_tensor_scaled_mm(x, fp8_weight, scale, bias, out_dtype)
+
+    config = FP8ScaledMMLinearLayerConfig(
+        weight_quant_key=kFp8StaticTensorSym,
+        activation_quant_key=kFp8StaticTensorSym,
+        weight_shape=(N, K),
+        input_dtype=out_dtype,
+        out_dtype=out_dtype,
+    )
+    kernel = CPUFp8PerTensorScaledMMLinearKernel(
+        config,
+        layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
+    )
+
+    layer = torch.nn.Module()
+    # Fp8LinearMethod stores the weight transposed to [K, N] before calling
+    # process_weights_after_loading.
+    layer.register_parameter(
+        "weight", torch.nn.Parameter(fp8_weight.t().contiguous(), requires_grad=False)
+    )
+    layer.register_parameter(
+        "weight_scale", torch.nn.Parameter(scale.clone(), requires_grad=False)
+    )
+    kernel.process_weights_after_loading(layer)
+
+    kernel_out = kernel.apply_weights(layer, x, bias)
+
+    assert kernel_out.dtype == out_dtype
+    torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
+
+
+@pytest.mark.parametrize("n", [16, 48])
+def test_cpu_fp8_per_tensor_kernel_rejects_non_multiple_of_32_n(n: int):
+    """The AMX tinygemm kernel tiles N in chunks of 32 and cannot handle a
+    remainder tile smaller than that, so can_implement must reject any N
+    that isn't a multiple of 32 rather than let it crash the process.
+    """
+    config = FP8ScaledMMLinearLayerConfig(
+        weight_quant_key=kFp8StaticTensorSym,
+        activation_quant_key=kFp8StaticTensorSym,
+        weight_shape=(n, 64),
+        input_dtype=torch.bfloat16,
+        out_dtype=torch.bfloat16,
+    )
+    supported, _ = CPUFp8PerTensorScaledMMLinearKernel.can_implement(config)
+    assert not supported

tests/v1/e2e/test_cpu_spec_decode.py

@@ -131,7 +131,7 @@ SMOKE_CONFIGS = [
         id="dflash",
     ),
     pytest.param(
-        "Qwen/Qwen3.5-0.8B-Base",
+        "Qwen/Qwen3.5-0.8B",
         {
             "method": "mtp",
             "num_speculative_tokens": 3,

vllm/model_executor/kernels/linear/__init__.py

@@ -187,6 +187,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
     CPUFp8BlockScaledMMKernel,
+    CPUFp8PerTensorScaledMMLinearKernel,
     CPUInt8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -438,6 +439,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
         ChannelWiseTorchFP8ScaledMMLinearKernel,
     ],
     PlatformEnum.CPU: [
+        CPUFp8PerTensorScaledMMLinearKernel,
         PerTensorTorchFP8ScaledMMLinearKernel,
         ChannelWiseTorchFP8ScaledMMLinearKernel,
     ],

vllm/model_executor/kernels/linear/scaled_mm/__init__.py

@@ -9,6 +9,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
     CPUFp8BlockScaledMMKernel,
+    CPUFp8PerTensorScaledMMLinearKernel,
     CPUInt8ScaledMMLinearKernel,
 )
 from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
@@ -67,5 +68,6 @@ __all__ = [
     "ZentorchInt8ScaledMMLinearKernel",
     "Fp8BlockScaledMMLinearKernel",
     "CPUFp8BlockScaledMMKernel",
+    "CPUFp8PerTensorScaledMMLinearKernel",
     "XPUFp8BlockScaledMMKernel",
 ]

vllm/model_executor/kernels/linear/scaled_mm/cpu.py

@@ -21,10 +21,14 @@ from .BlockScaledMMLinearKernel import (
     FP8ScaledMMLinearLayerConfig,
 )
 from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
     Int8ScaledMMLinearKernel,
     Int8ScaledMMLinearLayerConfig,
 )
 
+# BLOCK_K in csrc/cpu/sgl-kernels/gemm.h — the AMX kernel's fixed K-block size.
+_BLOCK_SIZE_K = 128
+
 
 class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
     @classmethod
@@ -330,3 +334,99 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
         raise NotImplementedError(
             "CPUFp8BlockScaledMMKernel overrides apply_weights directly."
         )
+
+
+class CPUFp8PerTensorScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """FP8 W8A16 per-tensor-scaled GEMM via AMX BRGEMM on CPU.
+
+    Reuses the block-scaled AMX kernel (fp8_scaled_mm_cpu) with a single
+    synthetic block spanning the whole weight, so activations stay BF16/FP32
+    — no FP8 activation quantization, unlike PerTensorTorchFP8ScaledMMLinearKernel.
+    """
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU platform."
+        if not torch.cpu._is_amx_tile_supported():
+            return False, "requires AMX tile support (Sapphire Rapids or newer)."
+        if not ops._supports_cpu_fp8_w8a16:
+            return False, "fp8_scaled_mm_cpu op not available."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if not c.weight_quant_key.scale.group_shape.is_per_tensor():
+            return False, "requires a per-tensor weight scale."
+        if not c.activation_quant_key.scale.group_shape.is_per_tensor():
+            return False, "requires a per-tensor activation scale."
+        if c.out_dtype not in (torch.bfloat16, torch.float32):
+            return False, "Only bfloat16/float32 output dtype supported."
+        n = c.weight_shape[0]
+        if n % 32 != 0:
+            # The AMX tinygemm kernel tiles N in chunks of 32 and cannot
+            # handle a remainder tile smaller than that.
+            return False, f"requires weight output dim (N={n}) to be a multiple of 32."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_name, w_s_name, _, _ = self.layer_param_names
+        # Fp8LinearMethod transposes weight to (K, N) for torch._scaled_mm-style
+        # kernels; convert_weight_packed/fp8_scaled_mm_cpu expect (N, K), as
+        # produced by CUDA nn.Linear-style checkpoints.
+        weight = getattr(layer, w_name).t().contiguous()
+        n, k = weight.shape
+
+        packed_weight = torch.ops._C.convert_weight_packed(weight)
+        replace_parameter(
+            layer, w_name, torch.nn.Parameter(packed_weight, requires_grad=False)
+        )
+
+        # Synthesize a single-block "block scale" tensor from the scalar
+        # per-tensor weight scale so fp8_scaled_mm_cpu can be reused as-is.
+        weight_scale = getattr(layer, w_s_name)
+        self._block_size_n = -(-n // 32) * 32  # round up to a multiple of 32
+        num_k_blocks = -(-k // _BLOCK_SIZE_K)
+        block_scale = weight_scale.reshape(()).expand(1, num_k_blocks).contiguous()
+        replace_parameter(
+            layer, w_s_name, torch.nn.Parameter(block_scale, requires_grad=False)
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_name, w_s_name, _, _ = self.layer_param_names
+        weight = getattr(layer, w_name)
+        weight_scale = getattr(layer, w_s_name)
+
+        x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
+        out = torch.ops._C.fp8_scaled_mm_cpu(
+            x_2d,
+            weight,
+            weight_scale,
+            [self._block_size_n, _BLOCK_SIZE_K],
+            bias,
+            x.dtype,
+            True,  # is_vnni (weight already prepacked)
+        )
+        return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        raise NotImplementedError(
+            "CPUFp8PerTensorScaledMMLinearKernel overrides apply_weights directly."
+        )