[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/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"