tzhouamandClaude Opus 5
3dac97e0bd
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>
93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
# 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
|