wehub-resource-sync
76 lines
3.1 KiB
Python
76 lines
3.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for the quack FP8 auto-enable gate.
|
|
|
|
quack's CuteDSL FP8 GEMM uses the 5th-gen ``tcgen05`` tensor-core MMA, which
|
|
exists only on datacenter Blackwell (``sm_100a`` / ``sm_101a`` / ``sm_103a``,
|
|
compute capability ``10.x``). It must NOT auto-enable on workstation/consumer
|
|
Blackwell (``sm_120`` / ``sm_121``, cc ``12.x``, e.g. RTX PRO 6000 / RTX
|
|
50-series), where ``tcgen05`` is absent and every FP8 GEMM would fall back to
|
|
FlashInfer one call at a time (catastrophically slow). See
|
|
``vllm_omni/quantization/quack_fp8.py``.
|
|
|
|
Regression guard for the ``>= 10`` gate that matched cc ``12.x`` too. These
|
|
tests are hardware-free: the ``torch.cuda`` probes are monkeypatched, so no GPU
|
|
is required.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.quantization import quack_fp8
|
|
|
|
_ENV = "VLLM_OMNI_USE_QUACK_FP8"
|
|
|
|
|
|
def _fake_cuda(monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]) -> None:
|
|
"""Pretend a CUDA device with the given (major, minor) capability is present."""
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
|
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"capability, expected",
|
|
[
|
|
((10, 0), True), # sm_100 datacenter Blackwell — tcgen05 present
|
|
((10, 3), True), # sm_103 datacenter Blackwell
|
|
((12, 0), False), # sm_120 workstation Blackwell — the regression
|
|
((12, 1), False), # sm_121 consumer Blackwell
|
|
((9, 0), False), # Hopper — CUTLASS already fuses bias, quack unused
|
|
((8, 9), False), # Ada
|
|
],
|
|
)
|
|
def test_auto_enable_only_on_datacenter_blackwell(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
capability: tuple[int, int],
|
|
expected: bool,
|
|
) -> None:
|
|
monkeypatch.delenv(_ENV, raising=False)
|
|
_fake_cuda(monkeypatch, capability)
|
|
assert quack_fp8._is_quack_capable() is expected
|
|
assert quack_fp8.quack_enabled() is expected
|
|
|
|
|
|
def test_no_cuda_disables_quack(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv(_ENV, raising=False)
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
|
assert quack_fp8._is_quack_capable() is False
|
|
assert quack_fp8.quack_enabled() is False
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", "On"])
|
|
def test_env_override_forces_on_even_on_sm120(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
|
|
# sm_120 auto-disables, but an explicit truthy override forces quack on
|
|
# (e.g. once CuteDSL ships sm_120a support).
|
|
_fake_cuda(monkeypatch, (12, 0))
|
|
monkeypatch.setenv(_ENV, value)
|
|
assert quack_fp8.quack_enabled() is True
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
|
|
def test_env_override_forces_off_even_on_datacenter_blackwell(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
|
|
# sm_100 auto-enables, but an explicit non-truthy override forces quack off.
|
|
_fake_cuda(monkeypatch, (10, 0))
|
|
monkeypatch.setenv(_ENV, value)
|
|
assert quack_fp8.quack_enabled() is False
|