88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
|
# 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
|