diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index e8d0a24bc7..43555a19ab 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.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}" } diff --git a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py b/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py index 3154e2cb98..b70a1a8543 100644 --- a/tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +++ b/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 diff --git a/tests/v1/e2e/test_cpu_spec_decode.py b/tests/v1/e2e/test_cpu_spec_decode.py index 2d34c11656..a7ba5580c3 100644 --- a/tests/v1/e2e/test_cpu_spec_decode.py +++ b/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, diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index d3c41f97d8..c610e9b9c8 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/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, ], diff --git a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py b/vllm/model_executor/kernels/linear/scaled_mm/__init__.py index 39f9abd460..f4a062b7f2 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/__init__.py +++ b/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", ] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index 92ea073a1a..e72dffcbbc 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/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." + )