[Edge][CPU] W4A8 via vLLM's own AMX kernel, a per-device weight path, and stop recompiling on every start

Four changes, each measured on a Xeon 8480C with Spark-X2.5-1.7B at 24
threads. Numbers are best-of-interleaved-passes; this host is shared and a
single sequential grid swung the llama.cpp control 30%, so anything under
~8% is not treated as a result.

**vLLM's CPU backend does have a 4-bit kernel.** An earlier note in this
tree said it did not, on the strength of `woq_int4_linear`,
`int4_scaled_mm_with_quant` and `da8w4_linear` all being absent. Those are
simply not the names it uses: `CPUWNA16LinearKernel` dispatches GPTQ/AWQ
checkpoints to `int4_scaled_mm_cpu`, an AMX W4A8 GEMM, whenever activations
are bf16 and the checkpoint carries no act-order. Four-bit weights times
eight-bit activations is the arithmetic llama.cpp's Q4_K_M already does.
`gptq_cpu_export` builds checkpoints for it; the kernel pins the zero point
at 8, so the grid must be symmetric, and a per-group scale search recovers
most of what that costs.

**A per-device weight path** (`edge/weight_path.py`). The fast kernel on each
device needs a different, mutually exclusive weight layout, so the choice has
to come from a probe rather than a flag: AMX x86 -> W4A8/GPTQ; x86 without
AMX, or priority="fidelity" -> W4A16/tinygemm (bf16 activations, the more
faithful of the two); ARM -> W4A16 but flagged unverified, because PyTorch's
4-bit CPU GEMM dispatches a different kernel there than the AVX-512 one
measured here; phone -> runs_in_vllm=False, the deliverable is an exported
graph with calibrated w8a16.

**The CPU compile cache never loaded.** `_try_load_aot_compiled_fn` finishes
with `compiled_fn.finalize_loading(cfg)`, which the CUDA wrapper object has
and the inductor CPU backend's plain function does not. The AttributeError
was swallowed as a cache miss, so every engine start recompiled the model:
warm start 49.8 s against 14.9 s with the missing no-op attached. Self-
disabling if upstream gives the CPU path a real one.

**The head gate stops being a GEMM.** A 2048x8 projection cannot amortize a
thread fan-out: dispatched as `aten::mm` it costs 1.155 ms of a decode step,
nearly all barrier. As an eager reduction behind an opaque op it costs 0.444
ms. It has to be opaque -- written inline the reduction is traced, and
inductor's loop made the whole step 1.8x slower (78.8 -> 44.5 tok/s).

Also included and *off by default*: `cpu_ir_ops` registers the fused
`_C::rms_norm` / `_C::fused_add_rms_norm` kernels that vLLM's CPU build ships
but gates on `GPGPU_DEVICE`. In isolation they are 3.6x and 4.5x faster than
the native path; end to end they are slightly slower (72.7 vs 74.0 tok/s),
because inductor was already fusing the native norm into its neighbours and
an opaque call breaks that fusion. Kept with the measurement recorded so the
next person does not repeat it. `gelu_and_mul` is deliberately not registered:
its fused CPU kernel is 2.7x slower at Spark's GEGLU shape.

Decode, W4A16 g64 -> W4A8 g128 with the gate fix: 69.1 -> 84.5 tok/s (+22%),
against llama.cpp Q4_K_M's 90.0; prefill 3347 against 1121 (3.0x). The two
4-bit paths are not equally faithful: against each build's own bf16 over 12
greedy prompts, Q4_K_M reproduces 4/12, W4A16 2/12 and W4A8 1/12.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

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"

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

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():

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",
+    )

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):

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,
+    )

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

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,

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()