From 809b98e5b71482e52ee0ecbf140084aae49f3652 Mon Sep 17 00:00:00 2001 From: Yuwen Zhou Date: Wed, 6 May 2026 15:05:27 +0800 Subject: [PATCH] [CPU] Add FP8 W8A16 linear support (#41186) Signed-off-by: yuwenzho --- .buildkite/hardware_tests/cpu.yaml | 6 +- csrc/cpu/sgl-kernels/gemm_fp8.cpp | 2 +- csrc/cpu/torch_bindings.cpp | 14 ++ .../quantization/test_cpu_fp8_scaled_mm.py | 162 ++++++++++++++++++ tests/quantization/test_cpu_wna16.py | 1 + vllm/_custom_ops.py | 33 ++++ .../model_executor/kernels/linear/__init__.py | 4 + .../kernels/linear/scaled_mm/__init__.py | 2 + .../kernels/linear/scaled_mm/cpu.py | 110 ++++++++++++ 9 files changed, 331 insertions(+), 3 deletions(-) create mode 100644 tests/kernels/quantization/test_cpu_fp8_scaled_mm.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index ba8fd497c39..77abb5532db 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -14,13 +14,15 @@ steps: - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/test_onednn.py - tests/kernels/test_awq_int4_to_int8.py + - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py commands: - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m " + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/test_onednn.py - pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py" + pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py + pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py" - label: CPU-Compatibility Tests depends_on: [] diff --git a/csrc/cpu/sgl-kernels/gemm_fp8.cpp b/csrc/cpu/sgl-kernels/gemm_fp8.cpp index ef29181cee5..487a73d0949 100644 --- a/csrc/cpu/sgl-kernels/gemm_fp8.cpp +++ b/csrc/cpu/sgl-kernels/gemm_fp8.cpp @@ -447,7 +447,7 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16); INSTANTIATE_TINYGEMM_TEMPLATE(at::Half); at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, - std::vector block_size, std::optional& bias, + std::vector block_size, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni) { RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector({mat1, mat2, scales2, block_size, bias})); diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 89c5765ee64..c5522521e78 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -77,6 +77,13 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni); +// Adapted from sglang: FP8 W8A16 kernel +at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, + at::Tensor& scales2, + std::vector block_size, + const std::optional& bias, + at::ScalarType out_dtype, bool is_vnni); + // Adapted from sglang: INT4 W4A8 kernels std::tuple convert_weight_packed_scale_zp( at::Tensor qweight, at::Tensor qzeros, at::Tensor scales); @@ -376,6 +383,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, " "Tensor(a3!) w_scales, Tensor? bias) -> Tensor"); ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu); + + // Adapted from sglang: FP8 W8A16 kernel + ops.def( + "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) " + "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, " + "bool is_vnni) -> Tensor"); + ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu); #endif // CPU attention kernels diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py new file mode 100644 index 00000000000..3154e2cb98b --- /dev/null +++ b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu). + +Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`. +""" + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +if not ops._supports_cpu_fp8_w8a16: + pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True) + +BLOCK_SIZE = [128, 128] + + +def cdiv(a: int, b: int) -> int: + return -(a // -b) + + +def quantize_weight_block_fp8( + weight: torch.Tensor, + block_size: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize weight [N, K] to FP8 with block scales. + + Returns: + fp8_weight: [N, K] float8_e4m3fn + scales: [n_tiles, k_tiles] float32 + """ + N, K = weight.shape + block_n, block_k = block_size + fp8_max = torch.finfo(torch.float8_e4m3fn).max + + n_tiles = cdiv(N, block_n) + k_tiles = cdiv(K, block_k) + + # Pad for even blocking + pad_N = (block_n - (N % block_n)) % block_n + pad_K = (block_k - (K % block_k)) % block_k + if pad_N > 0 or pad_K > 0: + weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N)) + + # Reshape into blocks + w_blocks = weight.view(n_tiles, block_n, k_tiles, block_k) + w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous() + + # Per-block scale + abs_max = w_blocks.abs().amax(dim=(-2, -1), keepdim=True) + scales = abs_max / fp8_max + scales = torch.where(scales == 0, torch.ones_like(scales), scales) + + # Quantize + q_fp8 = (w_blocks / scales).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn) + + # Reshape back + fp8_weight = ( + q_fp8.permute(0, 2, 1, 3) + .contiguous() + .view(N + pad_N, K + pad_K)[:N, :K] + .contiguous() + ) + + scales = scales.view(n_tiles, k_tiles) + return fp8_weight, scales + + +def dequant_weight_block_fp8( + fp8_weight: torch.Tensor, + scales: torch.Tensor, + block_size: list[int], + out_dtype: torch.dtype, +) -> torch.Tensor: + """Dequantize FP8 weight back to float for reference computation.""" + N, K = fp8_weight.shape + block_n, block_k = block_size + n_tiles, k_tiles = scales.shape + + pad_N = (block_n - (N % block_n)) % block_n + pad_K = (block_k - (K % block_k)) % block_k + if pad_N > 0 or pad_K > 0: + fp8_padded = torch.nn.functional.pad(fp8_weight.float(), (0, pad_K, 0, pad_N)) + else: + fp8_padded = fp8_weight.float() + + w_blocks = fp8_padded.view(n_tiles, block_n, k_tiles, block_k) + w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous() + dq = w_blocks * scales.view(n_tiles, k_tiles, 1, 1) + dq = dq.permute(0, 2, 1, 3).contiguous().view(N + pad_N, K + pad_K) + return dq[:N, :K].to(out_dtype) + + +def ref_fp8_block_scaled_mm( + x: torch.Tensor, + fp8_weight: torch.Tensor, + scales: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Reference: dequant FP8→float32, matmul in float32, cast to out_dtype.""" + w_dq = dequant_weight_block_fp8(fp8_weight, scales, block_size, torch.float32) + out = torch.mm(x.float(), w_dq.t()) + if bias is not None: + out = out + bias.float() + return out.to(out_dtype) + + +# --------------------------------------------------------------------------- +# Test parameters +# --------------------------------------------------------------------------- +M_SIZES = [1, 4, 16, 64, 128] +# (N, K) — weight shape is [N, K], output has N columns. +NK_SIZES = [ + (128, 256), + (256, 512), + (512, 1024), + (1024, 2048), + (5120, 5120), + (17408, 5120), + (5120, 17408), +] + + +@pytest.mark.parametrize("M", M_SIZES) +@pytest.mark.parametrize("N,K", NK_SIZES) +@pytest.mark.parametrize("use_bias", [False, True]) +def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool): + """fp8_scaled_mm_cpu correctness against float reference.""" + torch.manual_seed(42) + out_dtype = torch.bfloat16 + block_size = BLOCK_SIZE + + 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, scales = quantize_weight_block_fp8(w_f32, block_size) + + bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None + + ref_out = ref_fp8_block_scaled_mm( + x, fp8_weight, scales, block_size, bias, out_dtype + ) + + packed_weight = torch.ops._C.convert_weight_packed(fp8_weight) + kernel_out = ops.fp8_scaled_mm_cpu( + x, + packed_weight, + scales, + block_size, + bias, + out_dtype, + True, + ) + + assert kernel_out.dtype == out_dtype + torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01) diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index 5520dc1747a..f6e54fc1473 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -13,6 +13,7 @@ MODELS = [ "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", + "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized ] DTYPE = ["bfloat16"] diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index ffea35de881..430536ecca4 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3219,6 +3219,39 @@ if hasattr(torch.ops._C, "int4_scaled_mm_cpu"): _supports_cpu_w4a8_int8 = bool(hasattr(torch.ops._C, "convert_weight_packed_scale_zp")) +if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"): + + @register_fake("_C::fp8_scaled_mm_cpu") + def fp8_scaled_mm_cpu_fake( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, + ) -> torch.Tensor: + M = mat1.size(0) + N = mat2.size(0) + return torch.empty((M, N), dtype=out_dtype, device=mat1.device) + + +_supports_cpu_fp8_w8a16 = bool(hasattr(torch.ops._C, "fp8_scaled_mm_cpu")) + + +def fp8_scaled_mm_cpu( + mat1: torch.Tensor, + mat2: torch.Tensor, + scales2: torch.Tensor, + block_size: list[int], + bias: torch.Tensor | None, + out_dtype: torch.dtype, + is_vnni: bool, +) -> torch.Tensor: + return torch.ops._C.fp8_scaled_mm_cpu( + mat1, mat2, scales2, block_size, bias, out_dtype, is_vnni + ) + class CPUDNNLGEMMHandler: def __init__(self) -> None: diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 5d513f767f0..b9a96a03530 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -110,6 +110,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( + CPUFp8BlockScaledMMKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -199,6 +200,9 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[ AiterFp8BlockScaledMMKernel, TritonFp8BlockScaledMMKernel, ], + PlatformEnum.CPU: [ + CPUFp8BlockScaledMMKernel, + ], } _POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = { diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index e86684b2f8a..f8f12f7b0cb 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py @@ -8,6 +8,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel impo Fp8BlockScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( + CPUFp8BlockScaledMMKernel, CPUInt8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cutlass import ( @@ -58,4 +59,5 @@ __all__ = [ "ROCmFP8ScaledMMLinearKernel", "TritonInt8ScaledMMLinearKernel", "Fp8BlockScaledMMLinearKernel", + "CPUFp8BlockScaledMMKernel", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 3d67a73af43..083cb473aac 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -14,6 +14,10 @@ from vllm.model_executor.layers.utils import check_cpu_sgl_kernel from vllm.platforms import current_platform from vllm.platforms.interface import CpuArchEnum +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) from .ScaledMMLinearKernel import ( Int8ScaledMMLinearKernel, Int8ScaledMMLinearLayerConfig, @@ -215,3 +219,109 @@ class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): x.dtype, True, ) + + +class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """FP8 W8A16 block-quantized GEMM via AMX BRGEMM on CPU.""" + + # Input stays BF16 — no FP8 activation quantization. + apply_input_quant = False + + @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, config: FP8ScaledMMLinearLayerConfig + ) -> tuple[bool, str | None]: + # Validate weight block shape + weight_gs = config.weight_quant_key.scale.group_shape + if weight_gs.col <= 0 or weight_gs.col != 128: + return False, ( + "CPU FP8 kernel requires K-dimension block size of 128, " + f"got {weight_gs.col}." + ) + if weight_gs.row <= 0 or weight_gs.row % 32 != 0: + return False, ( + "CPU FP8 kernel requires N-dimension block size to be " + f"a positive multiple of 32, got {weight_gs.row}." + ) + if config.out_dtype not in (torch.bfloat16, torch.float32): + return False, "Only bfloat16/float32 output dtype supported." + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Skip the base class process (FP8 padding / fnuz normalization) + # which is GPU-oriented. Instead, VNNI-prepack weights for AMX. + params = self._get_layer_params(layer) + packed_weight = torch.ops._C.convert_weight_packed(params.weight) + replace_parameter( + layer, + params.WEIGHT, + torch.nn.Parameter(packed_weight, requires_grad=False), + ) + + # Re-wrap scale as a plain Parameter so the kernel can read it + # without weight-loader metadata interfering. + scale_attr = ( + params.WEIGHT_SCALE_INV + if params.weight_scale_inv is not None + else params.WEIGHT_SCALE + ) + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) + assert weight_scale is not None + replace_parameter( + layer, + scale_attr, + torch.nn.Parameter(weight_scale.data, requires_grad=False), + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + params = self._get_layer_params(layer) + weight_scale = ( + params.weight_scale_inv + if params.weight_scale_inv is not None + else params.weight_scale + ) + + x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x + out = torch.ops._C.fp8_scaled_mm_cpu( + x_2d, + params.weight, + weight_scale, + list(self.weight_group_shape), + 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_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "CPUFp8BlockScaledMMKernel overrides apply_weights directly." + )