92 lines
3.5 KiB
Python
92 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
|