diff --git a/tests/edge/test_weight_path.py b/tests/edge/test_weight_path.py new file mode 100644 index 000000000..9afca0a24 --- /dev/null +++ b/tests/edge/test_weight_path.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The per-device weight/kernel selector.""" + +import pytest + +from vllm_omni.edge.hardware_probe import HardwareProfile +from vllm_omni.edge.weight_path import has_amx, select_weight_path + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _x86(*, amx: bool = True, **kw) -> HardwareProfile: + flags = ["avx2", "avx512f"] + (["amx_tile", "amx_bf16"] if amx else []) + return HardwareProfile( + arch="x86_64", cpu_flags=flags, n_physical=32, n_logical=64, + ram_total_bytes=64 << 30, ram_available_bytes=48 << 30, + supports_bf16=True, **kw, + ) + + +def test_amx_x86_takes_the_vllm_int4_kernel(): + path = select_weight_path(_x86(amx=True)) + assert path.name == "w4a8-amx" + assert path.quantization == "gptq" + assert path.runs_in_vllm and path.verified + assert path.as_engine_kwargs() == {"quantization": "gptq"} + + +def test_x86_without_amx_falls_back_to_bf16_activations(): + """No AMX means no int4_scaled_mm_cpu, so tinygemm is the only 4-bit path.""" + path = select_weight_path(_x86(amx=False)) + assert path.name == "w4a16-tinygemm" + assert path.quantization == "cpu_int4" + + +def test_fidelity_priority_gives_up_the_amx_kernel(): + """W4A8 quantizes activations to int8; asking for fidelity must not.""" + fast = select_weight_path(_x86(amx=True), priority="speed") + faithful = select_weight_path(_x86(amx=True), priority="fidelity") + assert fast.name == "w4a8-amx" + assert faithful.name == "w4a16-tinygemm" + + +def test_phone_does_not_run_in_vllm_at_all(): + profile = HardwareProfile(arch="aarch64", accelerator="npu:qualcomm", n_physical=8) + path = select_weight_path(profile) + assert path.runs_in_vllm is False + assert path.as_engine_kwargs() == {} + assert path.build_with and "export" in path.build_with + # The calibration lesson has to travel with the recommendation. + assert any("calibrat" in n.lower() for n in path.notes) + + +def test_gpu_keeps_the_model_dtype(): + profile = HardwareProfile( + arch="x86_64", accelerator="cuda_discrete", gpu_mem_bytes=80 << 30, + supports_bf16=True, n_physical=32, + ) + path = select_weight_path(profile) + assert path.name == "bf16" + assert path.quantization is None + assert path.fused_cpu_norms is False + + +def test_arm_cpu_is_offered_but_flagged_unverified(): + """The 4-bit CPU GEMM has an ARM path, but a different kernel than the one + measured here, so the selector must not claim it was.""" + profile = HardwareProfile(arch="aarch64", cpu_flags=["neon"], n_physical=8) + path = select_weight_path(profile) + assert path.quantization == "cpu_int4" + assert path.verified is False + + +def test_has_amx_needs_x86_and_the_tile_flag(): + assert has_amx(_x86(amx=True)) + assert not has_amx(_x86(amx=False)) + assert not has_amx(HardwareProfile(arch="aarch64", cpu_flags=["amx_tile"])) + + +def test_every_path_states_why(): + for profile in ( + _x86(amx=True), _x86(amx=False), + HardwareProfile(arch="aarch64", accelerator="npu:qualcomm"), + HardwareProfile(arch="x86_64", accelerator="cuda_discrete"), + ): + path = select_weight_path(profile) + assert len(path.rationale) > 40, path.name diff --git a/tests/model_executor/layers/quantization/test_gptq_cpu_export.py b/tests/model_executor/layers/quantization/test_gptq_cpu_export.py new file mode 100644 index 000000000..72371e886 --- /dev/null +++ b/tests/model_executor/layers/quantization/test_gptq_cpu_export.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The GPTQ W4A8 export convention, checked against the kernel that reads it.""" + +import pytest +import torch + +from vllm_omni.model_executor.layers.quantization.gptq_cpu_export import ( + PACK_FACTOR, + dequantize_symmetric, + gptq_tensors, + pack_int32, + quantize_symmetric, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_pack_round_trips_through_vllms_unpacker(): + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + unpack_quantized_values_into_int32, + ) + from vllm.scalar_type import scalar_types + + codes = torch.randint(0, 16, (64, 32), dtype=torch.int32) + packed = pack_int32(codes, dim=0) + assert packed.shape == (64 // PACK_FACTOR, 32) + back = unpack_quantized_values_into_int32(packed, scalar_types.uint4, 0) + assert torch.equal(back, codes) + + packed1 = pack_int32(codes, dim=1) + assert packed1.shape == (64, 32 // PACK_FACTOR) + back1 = unpack_quantized_values_into_int32(packed1, scalar_types.uint4, 1) + assert torch.equal(back1, codes) + + +def test_codes_stay_in_range_and_shapes_match_the_loader(): + torch.manual_seed(0) + w = torch.randn(64, 256) + qweight, qzeros, scales, g_idx, snr = gptq_tensors(w, 64) + assert qweight.shape == (256 // 8, 64) and qweight.dtype == torch.int32 + assert qzeros.shape == (256 // 64, 64 // 8) + assert scales.shape == (256 // 64, 64) and scales.dtype == torch.bfloat16 + assert g_idx.shape == (256,) and g_idx.max().item() == 3 + assert snr > 18.0 + + +def test_scale_search_beats_the_fixed_grid(): + torch.manual_seed(1) + w = torch.randn(128, 256) + w[:, 0] *= 40.0 # one outlier per row is what a fixed scale handles worst + err = {} + for search in (False, True): + codes, scale = quantize_symmetric(w, 64, search=search) + deq = dequantize_symmetric(codes, scale, 64) + err[search] = (w - deq).norm().item() + assert err[True] < err[False] + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), reason="W4A8 kernel needs AMX tiles" +) +def test_kernel_reads_back_what_we_wrote(): + """End-to-end against `int4_scaled_mm_cpu`, through the exact repack + `CPUWNA16LinearKernel._process_gptq_weights_w4a8` performs. + + This is the guard that matters: a checkpoint written to a different + zero-point convention still loads and still generates text, just wrong + text, so nothing short of running the kernel catches it. + """ + from vllm import _custom_ops as ops + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + pack_quantized_values_into_int32, + unpack_quantized_values_into_int32, + ) + from vllm.scalar_type import scalar_types + + wt = scalar_types.uint4 + torch.manual_seed(0) + n, k, group = 64, 256, 64 + w = torch.randn(n, k) + qweight, qzeros, scales, _, _ = gptq_tensors(w, group) + codes, scale = quantize_symmetric(w, group) + deq = dequantize_symmetric(codes, scale, group) + + weight = unpack_quantized_values_into_int32(qweight, wt, 0) + weight = weight.view(k, n // 8, 8)[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(k, n) + weight = pack_quantized_values_into_int32(weight, wt, 1).contiguous() + fake_zp = torch.ones(k // group, n // 8, dtype=torch.int32) * -2004318072 + zp = unpack_quantized_values_into_int32(fake_zp, wt, 1) + zp = zp.view(k // group, n // 8, 8)[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape( + k // group, n + ) + zp = pack_quantized_values_into_int32(zp, wt, 1).contiguous() + bw, bzp, bs = ops.convert_weight_packed_scale_zp( + weight, zp, scales, ops.CPUQuantAlgo.AWQ + ) + + x = torch.randn(4, k, dtype=torch.bfloat16) + got = ops.int4_scaled_mm_cpu(x=x, w=bw, w_zeros=bzp, w_scales=bs, bias=None).float() + exp = x.float() @ deq.t() + snr = 20 * torch.log10(exp.norm() / (got - exp).norm()) + # The kernel quantizes activations to int8, so ~40 dB is the ceiling here; + # a convention mismatch lands far below zero. + assert snr > 30.0, f"kernel disagrees with our grid: {snr:.1f} dB" diff --git a/tests/model_executor/layers/test_cpu_ir_ops.py b/tests/model_executor/layers/test_cpu_ir_ops.py new file mode 100644 index 000000000..d89436a65 --- /dev/null +++ b/tests/model_executor/layers/test_cpu_ir_ops.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU fused-RMSNorm IR registration.""" + +import pytest +import torch + +from vllm import ir +from vllm_omni.model_executor.layers import cpu_ir_ops + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +pytest.importorskip("vllm") + +_SKIP = not cpu_ir_ops.CPU_FUSED_NORMS +_REASON = "vLLM CPU build without _C rms_norm kernels" + + +def _snr(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.float(), b.float() + return float(20 * torch.log10(a.norm() / (a - b).norm().clamp(min=1e-9))) + + +def test_registered_under_its_own_provider(): + """Registering a new provider, rather than editing vllm_c.py, is what keeps + this working across wheel upgrades.""" + assert cpu_ir_ops.PROVIDER in ir.ops.rms_norm.impls + assert cpu_ir_ops.PROVIDER in ir.ops.fused_add_rms_norm.impls + + +@pytest.mark.skipif(_SKIP, reason=_REASON) +@pytest.mark.parametrize("tokens", [1, 5, 64]) +def test_rms_norm_matches_the_native_impl(tokens): + torch.manual_seed(0) + x = torch.randn(tokens, 2048, dtype=torch.bfloat16) + w = torch.randn(2048, dtype=torch.bfloat16) + native = ir.ops.rms_norm.impls["native"].impl_fn(x, w, 1e-6) + ours = ir.ops.rms_norm.impls[cpu_ir_ops.PROVIDER].impl_fn(x, w, 1e-6) + assert ours.shape == native.shape and ours.dtype == native.dtype + # Both accumulate in their own order at bf16 output precision, so ~48 dB is + # the ceiling; a wrong kernel or layout lands far below. + assert _snr(native, ours) > 40.0, _snr(native, ours) + + +@pytest.mark.skipif(_SKIP, reason=_REASON) +def test_fused_add_rms_norm_writes_both_outputs(): + torch.manual_seed(0) + x = torch.randn(4, 2048, dtype=torch.bfloat16) + res = (x * 0.5).contiguous() + w = torch.randn(2048, dtype=torch.bfloat16) + n_out, n_res = ir.ops.fused_add_rms_norm.impls["native"].impl_fn( + x.clone(), res.clone(), w, 1e-6 + ) + o_out, o_res = ir.ops.fused_add_rms_norm.impls[cpu_ir_ops.PROVIDER].impl_fn( + x.clone(), res.clone(), w, 1e-6 + ) + assert _snr(n_out, o_out) > 40.0, _snr(n_out, o_out) + # The residual is a plain sum, so it should agree far more tightly. + assert _snr(n_res, o_res) > 100.0, _snr(n_res, o_res) + + +@pytest.mark.skipif(_SKIP, reason=_REASON) +def test_three_dimensional_input_round_trips_its_shape(): + x = torch.randn(2, 3, 2048, dtype=torch.bfloat16) + w = torch.randn(2048, dtype=torch.bfloat16) + out = ir.ops.rms_norm.impls[cpu_ir_ops.PROVIDER].impl_fn(x, w, 1e-6) + assert out.shape == x.shape + + +def test_declines_arguments_the_kernel_cannot_take(): + """variance_size and a non-contiguous input must fall back to native.""" + impl = ir.ops.rms_norm.impls[cpu_ir_ops.PROVIDER] + x = torch.randn(4, 2048, dtype=torch.bfloat16) + w = torch.randn(2048, dtype=torch.bfloat16) + assert impl.supports_args(x, w, 1e-6) + assert not impl.supports_args(x, w, 1e-6, 1024) + assert not impl.supports_args(x.t(), w, 1e-6) + # Mismatched weight dtype is the other case the kernel rejects. + assert not impl.supports_args(x, w.float(), 1e-6) + + +def test_priority_patch_is_idempotent_and_puts_us_first(): + if not cpu_ir_ops.CPU_FUSED_NORMS: + pytest.skip(_REASON) + from vllm.platforms.cpu import CpuPlatform + + cpu_ir_ops.prefer_cpu_fused_norms() + cpu_ir_ops.prefer_cpu_fused_norms() # second call must not double-wrap + priority = CpuPlatform.get_default_ir_op_priority(None) + assert priority.rms_norm[0] == cpu_ir_ops.PROVIDER + assert priority.rms_norm.count(cpu_ir_ops.PROVIDER) == 1 + assert "native" in priority.rms_norm diff --git a/tests/model_executor/models/spark2_5/test_spark2_5.py b/tests/model_executor/models/spark2_5/test_spark2_5.py index 2620c7931..823677f2e 100644 --- a/tests/model_executor/models/spark2_5/test_spark2_5.py +++ b/tests/model_executor/models/spark2_5/test_spark2_5.py @@ -147,6 +147,49 @@ def test_headwise_gate_scales_each_head_independently(): ) +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("tokens", [1, 2, 3, 16]) +def test_head_gate_branches_compute_the_same_projection(tokens): + """The gate op picks a reduction at one or two tokens and a GEMM above. + + Both arms have to be the same linear map, or a sequence would be gated + one way while it is prefilled and another way while it decodes. + """ + from vllm_omni.model_executor.models.spark2_5.spark2_5 import _head_gate_score + + torch.manual_seed(0) + g = torch.randn(8, 2048, dtype=torch.bfloat16) + x = torch.randn(tokens, 2048, dtype=torch.bfloat16) + + got = _head_gate_score(x, g) + expected = torch.nn.functional.linear(x.float(), g.float()) + assert got.shape == (tokens, 8) + assert got.dtype == torch.float32 + # The reduction accumulates in fp32 and the GEMM in bf16, so the GEMM arm + # is the looser of the two against an fp32 reference. + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_head_gate_op_is_registered_and_traceable(): + """It must stay opaque. Written inline the reduction gets traced, and + inductor's generated loop made the whole decode step 1.8x slower -- + 78.8 tok/s to 44.5 -- than leaving the matmul alone.""" + from vllm_omni.model_executor.models.spark2_5 import spark2_5 # noqa: F401 + + g = torch.randn(8, 64, dtype=torch.bfloat16) + x = torch.randn(1, 64, dtype=torch.bfloat16) + out = torch.ops.vllm.spark_head_gate(x, g) + assert out.shape == (1, 8) + + compiled = torch.compile( + lambda a, b: torch.ops.vllm.spark_head_gate(a, b), dynamic=True + ) + torch.testing.assert_close(compiled(x, g), out) + + @pytest.mark.core_model @pytest.mark.cpu def test_registered_in_the_omni_registry(): diff --git a/vllm_omni/edge/weight_path.py b/vllm_omni/edge/weight_path.py new file mode 100644 index 000000000..3cf7e64b0 --- /dev/null +++ b/vllm_omni/edge/weight_path.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Which weight format and which kernel, per class of device. + +A laptop, an edge gateway and a phone do not want the same build of the same +model, and the reason is not preference -- it is that the fast kernel on each +one needs a different weight layout, and the layouts are mutually exclusive: + +* an x86 core with AMX tiles has a 4-bit GEMM (`_C::int4_scaled_mm_cpu`) that + wants GPTQ-packed weights and quantizes activations to int8 -- the same + arithmetic llama.cpp's Q4_K_M does, and the fastest thing measured here; +* an x86 core without AMX has no such kernel, so 4-bit has to go through + PyTorch's tinygemm path, which wants its own packing and keeps activations + in bf16 -- slower, but more faithful, and the only option; +* a phone NPU does not run vLLM at all. The deliverable there is an exported + graph, and its quantization is decided by the export toolchain, not by us. + +So the choice is a property of the device, and picking it at deploy time from +a probe is the difference between "4-bit is 22% faster" and "4-bit is broken +here". Measured on a Xeon 8480C (24 threads, Spark-X2.5-1.7B, best of three +interleaved passes, tok/s decode / prefill): + + bf16 41.0 / 4796 no 4-bit kernel used + w4a16 tinygemm g64 75.1 / 3207 + w4a16 tinygemm g128 78.7 / 3346 more faithful of the two + w4a8 AMX g128 84.5 / 3347 fastest; least faithful + llama.cpp Q4_K_M 90.0 / 1121 the bar + +Fidelity is the reason both 4-bit rows stay: against each build's own bf16 +output over 12 greedy prompts, llama.cpp's Q4_K_M reproduced 4/12 exactly, +w4a16 2/12 and w4a8 1/12. W4A8 pins the zero point at 8 (so the grid must be +symmetric) *and* quantizes activations to int8; W4A16 does neither. +""" + +from dataclasses import dataclass, field +from typing import Literal + +from vllm_omni.edge.hardware_probe import ( + HW_CLASS_ARM64_CPU, + HW_CLASS_CUDA_DISCRETE, + HW_CLASS_JETSON, + HW_CLASS_NPU_PHONE, + HW_CLASS_X86_CPU, + HardwareProfile, + hardware_class, +) + +Priority = Literal["speed", "fidelity"] + + +@dataclass(frozen=True) +class WeightPath: + """How to build and run the weights on one class of device.""" + + name: str + quantization: str | None + """vLLM ``--quantization`` name, or None to keep the model dtype.""" + group_size: int | None + bits_per_weight: float | None + fused_cpu_norms: bool + runs_in_vllm: bool + """False means vLLM cannot execute here and an exported artifact is the + deliverable -- the phone case.""" + verified: bool + """True only where this combination was actually measured on hardware.""" + rationale: str + build_with: str | None = None + """Script that produces a checkpoint in this format.""" + notes: list[str] = field(default_factory=list) + + def as_engine_kwargs(self) -> dict[str, object]: + """The subset that goes straight into ``LLM(...)`` / deploy overrides.""" + if not self.runs_in_vllm: + return {} + kwargs: dict[str, object] = {} + if self.quantization: + kwargs["quantization"] = self.quantization + return kwargs + + +_AMX_FLAGS = ("amx_tile", "amx_bf16") + + +def has_amx(profile: HardwareProfile) -> bool: + """AMX tiles, which is what gates vLLM's CPU 4-bit GEMM.""" + return profile.arch == "x86_64" and any(f in profile.cpu_flags for f in _AMX_FLAGS) + + +def select_weight_path( + profile: HardwareProfile, priority: Priority = "speed" +) -> WeightPath: + """Pick the weight format for this device. + + ``priority='fidelity'`` keeps bf16 activations wherever a choice exists, + which costs ~7% of decode on an AMX machine and buys back the int8-activation + error. + """ + cls = hardware_class(profile) + + if cls == HW_CLASS_NPU_PHONE: + return WeightPath( + name="export-qnn-w8a16", + quantization=None, + group_size=None, + bits_per_weight=8.0, + fused_cpu_norms=False, + runs_in_vllm=False, + verified=True, + rationale=( + "vLLM has no NPU backend; the artifact is an exported graph. " + "Calibrated w8a16 is the only scheme measured correct on device " + "(36.9 dB); int8 activations break it (5.1 dB) and 4-bit weights " + "fail the converter on the calibrated path." + ), + build_with="vllm_omni/edge/spark_export.py", + notes=[ + "Calibration must use real activations: uncalibrated w4a16 " + "scores -1.6 dB on device, i.e. uncorrelated with the reference.", + "4-bit saves memory, not time, on this NPU: 0.917 vs 0.912 ms " + "per layer against 8-bit.", + ], + ) + + if cls in (HW_CLASS_CUDA_DISCRETE, HW_CLASS_JETSON): + return WeightPath( + name="bf16" if profile.supports_bf16 else "fp16", + quantization=None, + group_size=None, + bits_per_weight=16.0, + fused_cpu_norms=False, + runs_in_vllm=True, + verified=cls == HW_CLASS_CUDA_DISCRETE, + rationale=( + "The CPU 4-bit kernels are x86 CPU kernels. On a GPU the weights " + "stay in the model dtype and the GPU quantization paths " + "(marlin/AWQ) are a separate question this selector does not make." + ), + ) + + if cls == HW_CLASS_ARM64_CPU: + return WeightPath( + name="w4a16-tinygemm", + quantization="cpu_int4", + group_size=128, + bits_per_weight=4.16, + fused_cpu_norms=True, + runs_in_vllm=True, + verified=False, + rationale=( + "No AMX, so no int4_scaled_mm_cpu. PyTorch's 4-bit CPU GEMM has " + "an ARM path, but it dispatches to a different kernel " + "(KleidiAI-style packing) than the AVX-512 one measured here." + ), + build_with="analysis/experiments/spark_edge/quantize_int4.py", + notes=["Unverified on this host: x86 only. Measure before trusting."], + ) + + # x86 CPU: the one case with a real choice to make. + if has_amx(profile) and priority == "speed": + return WeightPath( + name="w4a8-amx", + quantization="gptq", + group_size=128, + bits_per_weight=4.16, + fused_cpu_norms=True, + runs_in_vllm=True, + verified=True, + rationale=( + "AMX tiles present, so vLLM's own int4_scaled_mm_cpu applies: " + "84.5 tok/s decode against 78.7 for tinygemm and 41.0 for bf16, " + "and 3.0x llama.cpp on prefill." + ), + build_with="analysis/experiments/spark_edge/quantize_gptq_int4.py", + notes=[ + "Least faithful of the three: 1/12 greedy prompts reproduced " + "exactly against its own bf16, vs 2/12 for w4a16.", + "Activations are quantized to int8 by the kernel; that is what " + "llama.cpp's Q4_K_M also does.", + ], + ) + + return WeightPath( + name="w4a16-tinygemm", + quantization="cpu_int4", + group_size=128, + bits_per_weight=4.16, + fused_cpu_norms=True, + runs_in_vllm=True, + verified=True, + rationale=( + "No AMX tiles, or fidelity was asked for: keep activations in bf16 " + "and take the tinygemm 4-bit GEMM. 78.7 tok/s decode, and the more " + "faithful of the two 4-bit paths." + ), + build_with="analysis/experiments/spark_edge/quantize_int4.py", + ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 6b9ae78d4..46cf640ab 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -144,6 +144,12 @@ def register_omni_models_to_vllm(): # the CPU backend the 4-bit weight path vLLM 0.28 does not ship). import vllm_omni.model_executor.layers.quantization # noqa: F401 + # Let the CPU backend use the fused RMSNorm kernels it already ships but + # marks unsupported (see cpu_ir_ops): ~1.8 ms of a 12 ms decode step. + from vllm_omni.model_executor.layers import cpu_ir_ops + + cpu_ir_ops.prefer_cpu_fused_norms() + @dataclass class OmniEngineArgs(EngineArgs): diff --git a/vllm_omni/model_executor/layers/cpu_ir_ops.py b/vllm_omni/model_executor/layers/cpu_ir_ops.py new file mode 100644 index 000000000..8fe76107c --- /dev/null +++ b/vllm_omni/model_executor/layers/cpu_ir_ops.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Give vLLM's CPU backend the fused RMSNorm kernels it already ships. + +vLLM's CPU build registers `_C::rms_norm` and `_C::fused_add_rms_norm`, and +`vllm/kernels/vllm_c.py` already wraps both as IR-op implementations -- but it +gates them on ``GPGPU_DEVICE = is_cuda_alike() or is_xpu()``, so on CPU the IR +layer marks them unsupported and every norm falls through to the PyTorch +implementation. Nothing then fuses it: inductor emits an elementwise loop and +spreads a 4 KB reduction across every core, which at one token is almost +entirely thread barrier. + +Measured on a Xeon 8480C at Spark-X2.5's decode shape (1 token, hidden 2048), +24 threads: + + rms_norm native 20.99 us fused 5.76 us 3.6x + fused_add_rms_norm native 41.14 us fused 9.24 us 4.5x + +Two norms per layer over 28 layers is ~1.8 ms of an ~12 ms decode step. + +The same file's `gelu_and_mul` is deliberately *not* re-registered here: the +fused CPU kernel measured 64.11 us against the native path's 23.36 us at +Spark's GEGLU shape ([1, 13312] -> [1, 6656]), i.e. 2.7x slower, so the native +form is the right choice there and the asymmetry is the point -- "use the +fused kernel" is not a rule, it is a per-op measurement. + +Registering under our own provider name rather than editing `vllm_c.py` keeps +this working across vLLM wheel upgrades; it is the mechanism vLLM documents for +out-of-tree kernels (see `vllm/kernels/oink_ops.py`). The equivalent upstream +change is two lines: add `is_cpu()` to `GPGPU_DEVICE` and give `CpuPlatform` a +`get_default_ir_op_priority`. +""" + +import os + +import torch +from torch import Tensor + +from vllm import ir +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +PROVIDER = "vllm_c_cpu" +_FUSED_NORM_OPS = ("rms_norm", "fused_add_rms_norm") + + +def _cpu_kernels_available() -> bool: + if not current_platform.is_cpu(): + return False + try: + current_platform.import_kernels() + return all(hasattr(torch.ops._C, name) for name in _FUSED_NORM_OPS) + except Exception as exc: # pragma: no cover - build without the extension + logger.debug("CPU fused norm kernels unavailable: %s", exc) + return False + + +CPU_FUSED_NORMS = _cpu_kernels_available() + + +def _supports(x: Tensor, weight, epsilon, variance_size=None) -> bool: + # The kernel normalizes over the last dim of a contiguous 2-D buffer and + # has no variance_size override; anything else stays on the native path. + return ( + variance_size is None + and (weight is None or weight.dtype == x.dtype) + and x.is_contiguous() + ) + + +def _supports_add(x: Tensor, x_residual: Tensor, weight, epsilon, variance_size=None): + return _supports(x, weight, epsilon, variance_size) and x_residual.is_contiguous() + + +@ir.ops.rms_norm.register_impl( + PROVIDER, supports_args=_supports, supported=CPU_FUSED_NORMS +) +def rms_norm( + x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None +) -> Tensor: + assert variance_size is None + shape = x.shape + x2d = x if x.dim() == 2 else x.reshape(-1, shape[-1]) + out = torch.empty_like(x2d) + torch.ops._C.rms_norm(out, x2d, weight, epsilon) + return out if x.dim() == 2 else out.reshape(shape) + + +@ir.ops.fused_add_rms_norm.register_impl( + PROVIDER, supports_args=_supports_add, supported=CPU_FUSED_NORMS, inplace=True +) +def fused_add_rms_norm( + x: Tensor, + x_residual: Tensor, + weight: Tensor | None, + epsilon: float, + variance_size: int | None = None, +) -> tuple[Tensor, Tensor]: + assert variance_size is None + if x.dim() == 2: + torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon) + return x, x_residual + shape = x.shape + x2d, r2d = x.reshape(-1, shape[-1]), x_residual.reshape(-1, shape[-1]) + torch.ops._C.fused_add_rms_norm(x2d, r2d, weight, epsilon) + return x2d.reshape(shape), r2d.reshape(shape) + + +def prefer_cpu_fused_norms() -> None: + """Make the CPU platform rank these kernels ahead of the native path. + + The platform hook only fills in ops the user left unset, so an explicit + ``kernel_config.ir_op_priority`` still wins. + """ + if not CPU_FUSED_NORMS: + return + if os.environ.get("VLLM_OMNI_CPU_FUSED_NORMS", "1") != "1": + logger.info("CPU fused RMSNorm kernels disabled by environment.") + return + from vllm.platforms.cpu import CpuPlatform + + if getattr(CpuPlatform, "_omni_fused_norm_priority", False): + return + + base = CpuPlatform.get_default_ir_op_priority.__func__ + + def get_default_ir_op_priority(cls, vllm_config): + priority = base(cls, vllm_config) + for op_name in _FUSED_NORM_OPS: + current = list(getattr(priority, op_name)) + if PROVIDER not in current: + setattr(priority, op_name, [PROVIDER] + current) + return priority + + CpuPlatform.get_default_ir_op_priority = classmethod(get_default_ir_op_priority) + CpuPlatform._omni_fused_norm_priority = True + logger.info( + "CPU fused RMSNorm kernels enabled (provider %s) ahead of the native path.", + PROVIDER, + ) diff --git a/vllm_omni/model_executor/layers/quantization/gptq_cpu_export.py b/vllm_omni/model_executor/layers/quantization/gptq_cpu_export.py new file mode 100644 index 000000000..2cdd7cd68 --- /dev/null +++ b/vllm_omni/model_executor/layers/quantization/gptq_cpu_export.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Build GPTQ-format 4-bit tensors for vLLM's CPU W4A8 kernel. + +`CPUWNA16LinearKernel` routes a GPTQ checkpoint to `int4_scaled_mm_cpu`, an +AMX W4A8 GEMM, when the activations are bf16 and the checkpoint carries no +activation reordering. That kernel builds its own zero point -- a constant 8 +-- and ignores the checkpoint's `qzeros`, so the grid written here has to be +symmetric about 8: + + w = (code - 8) * scale, code in [0, 15] + +Keeping that assumption in the package, with a test that checks it against the +kernel itself, is the point of this module: it is the one thing that can go +silently wrong -- a checkpoint quantized to a different convention still loads, +still generates, and is simply wrong. +""" + +import torch + +PACK_FACTOR = 8 +# Scale candidates as a fraction of amax/8. 0.875 puts +amax exactly on code +# 15; larger values trade clipping the extreme for a finer step, and which +# wins depends on the group, so it is searched rather than assumed. +DEFAULT_SHRINKS = (1.06, 1.0, 0.95, 0.9, 0.875, 0.85, 0.8, 0.75) + + +def pack_int32(codes: torch.Tensor, dim: int) -> torch.Tensor: + """Pack 4-bit codes into int32 words, eight per word, along ``dim``.""" + codes = codes.to(torch.int32) + if dim == 0: + out = torch.zeros( + codes.shape[0] // PACK_FACTOR, codes.shape[1], dtype=torch.int32 + ) + for i in range(PACK_FACTOR): + out |= (codes[i::PACK_FACTOR] & 0xF) << (4 * i) + else: + out = torch.zeros( + codes.shape[0], codes.shape[1] // PACK_FACTOR, dtype=torch.int32 + ) + for i in range(PACK_FACTOR): + out |= (codes[:, i::PACK_FACTOR] & 0xF) << (4 * i) + return out + + +def quantize_symmetric(w: torch.Tensor, group: int, search: bool = True): + """[N, K] float -> (codes [N, K] in [0, 15], scales [N, K // group]).""" + n, k = w.shape + if k % group: + raise ValueError(f"input dim {k} is not a multiple of group size {group}") + wf = w.float().reshape(n, k // group, group) + amax = wf.abs().amax(-1).clamp(min=1e-9) + + best_err = best = None + for f in (DEFAULT_SHRINKS if search else (0.875,)): + scale = amax * f / 8.0 + codes = (wf / scale.unsqueeze(-1)).round().add_(8).clamp_(0, 15) + err = ((codes - 8.0) * scale.unsqueeze(-1) - wf).pow_(2).sum(-1) + if best_err is None: + best_err, best = err, (codes, scale) + else: + take = err < best_err + best_err = torch.where(take, err, best_err) + best = ( + torch.where(take.unsqueeze(-1), codes, best[0]), + torch.where(take, scale, best[1]), + ) + codes, scale = best + return codes.reshape(n, k).to(torch.int32), scale + + +def dequantize_symmetric(codes: torch.Tensor, scale: torch.Tensor, group: int): + n, k = codes.shape + c = codes.float().reshape(n, k // group, group) + return ((c - 8.0) * scale.unsqueeze(-1)).reshape(n, k) + + +def gptq_tensors(w: torch.Tensor, group: int, search: bool = True): + """[N, K] float -> the four tensors AutoGPTQLinearMethod expects. + + Returns ``(qweight [K//8, N], qzeros [K//g, N//8], scales [K//g, N], + g_idx [K], snr_db)``. + """ + n, k = w.shape + codes, scale = quantize_symmetric(w, group, search) + deq = dequantize_symmetric(codes, scale, group) + wf = w.float() + snr = float(20 * torch.log10(wf.norm() / (wf - deq).norm().clamp(min=1e-12))) + + qweight = pack_int32(codes.t().contiguous(), dim=0) + scales = scale.t().contiguous().to(torch.bfloat16) + qzeros = pack_int32(torch.full((k // group, n), 8, dtype=torch.int32), dim=1) + g_idx = torch.arange(k, dtype=torch.int32) // group + return qweight, qzeros, scales, g_idx, snr diff --git a/vllm_omni/model_executor/models/spark2_5/spark2_5.py b/vllm_omni/model_executor/models/spark2_5/spark2_5.py index e477b4e02..22b473d1e 100644 --- a/vllm_omni/model_executor/models/spark2_5/spark2_5.py +++ b/vllm_omni/model_executor/models/spark2_5/spark2_5.py @@ -21,6 +21,7 @@ The checkpoint stores QKV pre-fused as a single ``q_k_v_proj``; load_weights splits it so vLLM's sharded ``QKVParallelLinear`` can take it. """ +import os from collections.abc import Iterable import torch @@ -49,7 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP -from vllm.model_executor.utils import set_weight_attrs from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -58,7 +58,10 @@ from vllm.model_executor.models.utils import ( make_layers, maybe_prefix, ) +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.torch_utils import direct_register_custom_op from .configuration_spark2_5 import Spark2_5Config @@ -173,6 +176,15 @@ class Spark2_5Attention(nn.Module): else: self.g_weight = None + # Only CPU pays a thread barrier per op that a 2048x8 projection + # cannot amortize; on GPU the plain matmul is already free. The + # override exists so the choice can be ablated on a new machine + # rather than assumed. + override = os.environ.get("VLLM_OMNI_SPARK_GATE_REDUCTION", "auto") + self.gate_via_reduction = ( + current_platform.is_cpu() if override == "auto" else override == "1" + ) + layer_idx = extract_layer_index(prefix) layer_type = config.layer_types[layer_idx] self.is_sliding = layer_type == "sliding_attention" @@ -217,11 +229,18 @@ class Spark2_5Attention(nn.Module): attn_output = self.attn(q, k, v) if self.g_weight is not None: - gate_score = torch.nn.functional.linear(hidden_states, self.g_weight) - if self.gate_attn_act_mode == "sigmoid": - gate = torch.sigmoid(gate_score.float()) + if self.gate_via_reduction: + gate_score = torch.ops.vllm.spark_head_gate( + hidden_states, self.g_weight + ) else: - gate = nn.functional.silu(gate_score.float()) + gate_score = torch.nn.functional.linear( + hidden_states, self.g_weight + ).float() + if self.gate_attn_act_mode == "sigmoid": + gate = torch.sigmoid(gate_score) + else: + gate = nn.functional.silu(gate_score) gate = gate.to(attn_output.dtype) # Reshape only the feature dim: a view(-1, ...) here would make # torch.compile re-derive the symbolic token count and fail. @@ -352,6 +371,46 @@ class Spark2_5Model(nn.Module): return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) +def _head_gate_score( + hidden_states: torch.Tensor, g_weight: torch.Tensor +) -> torch.Tensor: + """Head-gate projection, routed by token count. + + Eight outputs is far too narrow for a GEMM to pay for its thread fan-out: + at one token `aten::mm` costs 34 us on 24 cores, 0.9 ms of a decode step, + nearly all of it barrier. The same arithmetic as a broadcast-and-reduce + costs 8 us, because ATen's grain size keeps a 16 K-element reduction on + one thread. Above a couple of tokens the GEMM wins again. + + It has to be an opaque op. Written inline the reduction is traced by + `@support_torch_compile`, and inductor's generated loop is far slower than + either eager form -- measured 44.5 tok/s against 78.8. + """ + if hidden_states.shape[0] <= _GATE_REDUCE_MAX_TOKENS: + return (hidden_states.unsqueeze(-2).float() * g_weight.float()).sum(-1) + return torch.nn.functional.linear(hidden_states, g_weight).float() + + +def _head_gate_score_fake( + hidden_states: torch.Tensor, g_weight: torch.Tensor +) -> torch.Tensor: + return torch.empty( + (hidden_states.shape[0], g_weight.shape[0]), + dtype=torch.float32, + device=hidden_states.device, + ) + + +_GATE_REDUCE_MAX_TOKENS = 2 + +direct_register_custom_op( + op_name="spark_head_gate", + op_func=_head_gate_score, + mutates_args=[], + fake_impl=_head_gate_score_fake, +) + + def _split_fused_qkv( weights: Iterable[tuple[str, torch.Tensor]], config: Spark2_5Config, diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 37086986a..e2cd479ce 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -579,3 +579,61 @@ def _patch_cumem_free_callback_cuda() -> None: _patch_cumem_free_callback_cuda() + + +def _patch_cpu_aot_compile_cache_load() -> None: + """Let the CPU backend actually reuse its compiled model across starts. + + ``vllm.compilation.decorators._try_load_aot_compiled_fn`` finishes a cache + load with ``loaded_fn._artifacts.compiled_fn.finalize_loading(cfg)``. On + CUDA that artifact is a wrapper object carrying the method; with the + inductor CPU backend it is a plain function, so the call raises + ``AttributeError``, the exception handler treats the whole load as a cache + miss ("Compiling model again due to a load failure"), and **every engine + start recompiles the model from scratch** -- most of a 50-100 s init on a + 28-layer model, paid again on every restart. + + Rather than copy vLLM's loader, attach the missing no-op: a plain Python + function accepts attribute assignment, and the CPU artifact loads its + generated shared object lazily on first call anyway (that call is the + warm-up run, so a genuine failure still surfaces during init, not in + production traffic). + + Self-disabling: if the artifact already has ``finalize_loading`` -- as it + does on GPU, and as it will if upstream gives the CPU path one -- this + changes nothing. + """ + import torch + + from vllm.platforms import current_platform + + if not current_platform.is_cpu(): + return + if os.environ.get("VLLM_OMNI_AOT_CACHE_SHIM", "1") != "1": + return + if getattr(torch.compiler, "_omni_aot_finalize_shim", False): + return + + original = torch.compiler.load_compiled_function + + def load_compiled_function(*args, **kwargs): + loaded = original(*args, **kwargs) + artifacts = getattr(loaded, "_artifacts", None) + compiled_fn = getattr(artifacts, "compiled_fn", None) + if compiled_fn is not None and not hasattr(compiled_fn, "finalize_loading"): + try: + compiled_fn.finalize_loading = lambda *_a, **_k: None + except (AttributeError, TypeError): + # Slot-based artifact: leave it alone and let vLLM recompile. + pass + return loaded + + torch.compiler.load_compiled_function = load_compiled_function + torch.compiler._omni_aot_finalize_shim = True + _PATCH_LOGGER.info( + "CPU AOT compile-cache load shim installed: compiled models are reused " + "across engine starts instead of being recompiled." + ) + + +_patch_cpu_aot_compile_cache_load()