tzhouamandClaude Opus 5
c4e61133fb
Two corrections and the paperwork they imply.
The fused-RMSNorm registration defaulted to *on*, which contradicts the
measurement it was committed with: those kernels are 3.6x and 4.5x faster in
isolation but measured 72.7 tok/s against 74.0 end to end, because inductor
already fuses the native norm into the residual add and an opaque call breaks
that fusion. The implementations stay registered, so the provider remains
selectable through `kernel_config.ir_op_priority` on a machine where the trade
lands differently; only the default ranking is left alone now, and the test
asserts both states rather than just the opt-in one.
The four switches this branch added (VLLM_OMNI_AOT_CACHE_SHIM,
_CPU_FUSED_NORMS, _CPU_RELEASE_MEMORY, _SPARK_GATE_REDUCTION) were not in the
reviewed environment-variable inventory, so
`test_statically_resolvable_environment_accesses_are_classified` failed. They
are classified now, the reviewed count moves 28 -> 32 with a note saying why,
and each has a row in docs/configuration/environment_variables.md carrying the
measurement that justifies its default.
`VLLM_OMNI_CPU_FAST_SAMPLER` and `VLLM_OMNI_STEP_STATS_SYNC` remain
unclassified and that test still fails on them. They arrived in 6ce8de1b, in
files this branch does not touch, so they are left to whoever owns that work
rather than documented from a guess.
Under `-m "core_model and cpu"`: 2593 passed before these fixes with 9
failures, of which two were mine. The rest are CUDA-only tests in a CPU venv
("Torch not compiled with CUDA enabled") plus test_qnn_export.py, which
another session is actively editing.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
107 lines
4.3 KiB
Python
107 lines
4.3 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_opt_in_and_idempotent(monkeypatch):
|
|
"""Off by default, because that is what the measurement says.
|
|
|
|
In isolation these kernels are 3.6x and 4.5x faster than the native path;
|
|
in the compiled model they measured 72.7 tok/s against 74.0, because
|
|
inductor already fuses the native norm into the residual add. So the
|
|
registration always happens -- the provider stays selectable through
|
|
`kernel_config.ir_op_priority` -- but the default ranking does not change
|
|
unless it is asked for.
|
|
"""
|
|
if not cpu_ir_ops.CPU_FUSED_NORMS:
|
|
pytest.skip(_REASON)
|
|
from vllm.platforms.cpu import CpuPlatform
|
|
|
|
monkeypatch.delenv("VLLM_OMNI_CPU_FUSED_NORMS", raising=False)
|
|
cpu_ir_ops.prefer_cpu_fused_norms()
|
|
assert CpuPlatform.get_default_ir_op_priority(None).rms_norm == ["native"]
|
|
|
|
monkeypatch.setenv("VLLM_OMNI_CPU_FUSED_NORMS", "1")
|
|
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
|