diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 010ef4ba7..8d0090ec0 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -65,7 +65,17 @@ vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with `QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts. In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into the image — so the first cold start does not recompile. The engine's startup dummy -run already exercises the kernels, so with a warm cache the first real request is fast. +run exercises the kernels, but new shapes, layouts, dtypes, or bias settings may +still need compilation or tuning. The warmup helper uses inference mode and +transposed weights without bias. Daemon workers compile candidates in-process +while retaining autotuning and caching. + +Scale validation and Quack/FlashInfer dispatch run inside a PyTorch custom op. +This keeps layer-specific scale addresses and validation-cache updates out of +Dynamo tracing without introducing a graph break. Unpopulated scales and Quack +failures still fall back to FlashInfer at runtime. CUDA graph capture additionally +requires warming the dispatch with populated scales; Python validation and +dispatch do not rerun during CUDA graph replay. To pre-warm specific shapes (e.g. at image build time): diff --git a/tests/quantization/test_quack_fp8_runtime.py b/tests/quantization/test_quack_fp8_runtime.py new file mode 100644 index 000000000..91605865a --- /dev/null +++ b/tests/quantization/test_quack_fp8_runtime.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU regressions for Quack compilation and warmup.""" + +import sys +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.quantization import quack_fp8 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_compile_pool_checks_daemon_at_tuning_time(monkeypatch): + process = SimpleNamespace(daemon=False) + compiler_pool = object() + pool = SimpleNamespace( + pool_scope=Mock(spec=[], side_effect=lambda: nullcontext(compiler_pool)), + suppress_pool=Mock(side_effect=lambda: nullcontext(object())), + ) + original_scope = pool.pool_scope + monkeypatch.setattr(quack_fp8, "current_process", lambda: process) + monkeypatch.setattr(quack_fp8, "import_module", lambda name: pool) + quack_fp8._configure_quack_compilation() + with pool.pool_scope() as active_pool: + assert active_pool is compiler_pool + pool.suppress_pool.assert_not_called() + + # spawn installs the child's daemon flag after importing the worker module. + process.daemon = True + with pool.pool_scope() as active_pool: + assert active_pool is None + original_scope.assert_called_once_with() + pool.suppress_pool.assert_called_once_with() + + +@pytest.fixture +def quack_dispatch(monkeypatch): + """Exercise the installed patch with CPU substitutes for the GPU kernels.""" + + def gemm(a, b, *, out, bias, alpha, tuned): + assert isinstance(alpha, torch.Tensor) + assert alpha.shape == (1,) + assert alpha.dtype == torch.float32 + assert tuned is True + result = (a.float() @ b.float()) * alpha + if bias is not None: + result += bias + out.copy_(result) + + def flashinfer(a, b, *, out_dtype, scale_a, scale_b, bias): + out = torch.full((a.shape[0], b.shape[1]), -7, dtype=out_dtype, device=a.device) + return out + bias if bias is not None else out + + class FlashInferKernel: + def apply_scaled_mm(self, **kwargs): + raise AssertionError("The installed custom op should handle dispatch") + + gemm_mock = Mock(side_effect=gemm) + fallback_mock = Mock(side_effect=flashinfer) + monkeypatch.setattr(quack_fp8, "quack_enabled", lambda: True) + monkeypatch.setattr(quack_fp8, "_gemm_interface", SimpleNamespace(gemm=gemm_mock)) + monkeypatch.setattr(quack_fp8, "_valid_scale_ptrs", set()) + monkeypatch.setattr(quack_fp8, "logger", Mock()) + monkeypatch.setitem( + sys.modules, + "vllm.model_executor.kernels.linear.scaled_mm.flashinfer", + SimpleNamespace(FlashInferFP8ScaledMMLinearKernel=FlashInferKernel), + ) + monkeypatch.setitem(sys.modules, "vllm.utils.flashinfer", SimpleNamespace(flashinfer_scaled_fp8_mm=fallback_mock)) + quack_fp8.install_quack_fp8_patch() + kernel = FlashInferKernel() + + def forward(a, b, scale_a, scale_b, bias=None, out_dtype=torch.bfloat16): + return kernel.apply_scaled_mm( + A=a, + B=b, + As=scale_a, + Bs=scale_b, + out_dtype=out_dtype, + bias=bias, + output_shape=(2, a.shape[0] // 2, b.shape[1]), + ) + + return SimpleNamespace(forward=forward, gemm=gemm_mock, fallback=fallback_mock) + + +@pytest.fixture +def compile_fullgraph(): + torch._dynamo.reset() + graphs = [] + + def backend(graph, example_inputs): + graphs.append(graph) + return graph.forward + + def compile_fn(fn, **kwargs): + return torch.compile(fn, backend=backend, fullgraph=True, **kwargs) + + try: + yield compile_fn, graphs + finally: + torch._dynamo.reset() + + +def _fp8_inputs(m=4): + a = torch.arange(m * 8, dtype=torch.float32).reshape(m, 8).to(torch.float8_e4m3fn) + b = torch.ones(6, 8, dtype=torch.float8_e4m3fn).t() + return a, b + + +@pytest.mark.parametrize("out_dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_scale_validation_does_not_recompile_for_each_layer(quack_dispatch, compile_fullgraph, out_dtype, with_bias): + compile_fn, graphs = compile_fullgraph + compiled = compile_fn(quack_dispatch.forward) + with torch.inference_mode(): + a, b = _fp8_inputs() + bias = torch.arange(6, dtype=out_dtype) if with_bias else None + # Keep all pairs alive, as separate model layers do. Both pointer + # addresses and scale values differ; tensor metadata stays identical. + pairs = [(torch.tensor([float(i + 1)]), torch.tensor([0.5])) for i in range(5)] + for _ in range(2): + for scale_a, scale_b in pairs: + expected = (a.float() @ b.float()) * scale_a * scale_b + if bias is not None: + expected += bias + out = compiled(a, b, scale_a, scale_b, bias, out_dtype) + torch.testing.assert_close(out, expected.to(out_dtype).view(2, 2, 6)) + assert len(graphs) == 1 + assert any(node.target == torch.ops.vllm_omni.quack_fp8_scaled_mm.default for node in graphs[0].graph.nodes) + assert quack_dispatch.gemm.call_count == 10 + quack_dispatch.fallback.assert_not_called() + + +def test_compiled_dispatch_rechecks_invalid_scales_after_loading(quack_dispatch, compile_fullgraph): + compile_fn, graphs = compile_fullgraph + compiled = compile_fn(quack_dispatch.forward) + with torch.inference_mode(): + a, b = _fp8_inputs() + pairs = [] + for bad_scale in (torch.finfo(torch.float32).min, 0.0, -1.0, float("nan"), float("inf")): + for scale_index in (0, 1): + scales = [torch.ones(1), torch.ones(1)] + scales[scale_index].fill_(bad_scale) + pairs.append(scales) + for scale_a, scale_b in pairs: + out = compiled(a, b, scale_a, scale_b) + torch.testing.assert_close(out, torch.full((2, 2, 6), -7, dtype=torch.bfloat16)) + quack_dispatch.gemm.assert_not_called() + assert quack_dispatch.fallback.call_count == len(pairs) + + # Populate the same buffers and switch branches without recompiling. + for scale_a, scale_b in pairs: + scale_a.fill_(0.25) + scale_b.fill_(0.5) + out = compiled(a, b, scale_a, scale_b) + expected = ((a.float() @ b.float()) * 0.125).to(torch.bfloat16) + torch.testing.assert_close(out, expected.view(2, 2, 6)) + assert quack_dispatch.gemm.call_count == len(pairs) + assert quack_dispatch.fallback.call_count == len(pairs) + assert len(graphs) == 1 + + +def test_compiled_dispatch_falls_back_on_quack_failure(monkeypatch, quack_dispatch, compile_fullgraph): + compile_fn, graphs = compile_fullgraph + compiled = compile_fn(quack_dispatch.forward) + with torch.inference_mode(): + a, b = _fp8_inputs() + scale_a, scale_b = torch.ones(1), torch.ones(1) + bias = torch.arange(6, dtype=torch.float32) + expected = (torch.full((2, 2, 6), -7, dtype=torch.bfloat16) + bias).to(torch.bfloat16) + compiled(a, b, scale_a, scale_b, bias) + quack_dispatch.fallback.assert_not_called() + + quack_dispatch.gemm.side_effect = RuntimeError("kernel failed") + out = compiled(a, b, scale_a, scale_b, bias) + torch.testing.assert_close(out, expected) + quack_fp8.logger.warning_once.assert_called_once() + + # The helper can also return None when Quack is unavailable. + monkeypatch.setattr(quack_fp8, "_gemm_interface", False) + out = compiled(a, b, scale_a, scale_b, bias) + torch.testing.assert_close(out, expected) + assert quack_dispatch.fallback.call_count == 2 + assert len(graphs) == 1 + + +@pytest.mark.parametrize("scale_index", [0, 1]) +def test_compiled_dispatch_falls_back_for_non_scalar_scales(quack_dispatch, compile_fullgraph, scale_index): + compile_fn, graphs = compile_fullgraph + compiled = compile_fn(quack_dispatch.forward) + with torch.inference_mode(): + a, b = _fp8_inputs() + scales = [torch.ones(1), torch.ones(1)] + scales[scale_index] = torch.ones(2) + out = compiled(a, b, *scales) + torch.testing.assert_close(out, torch.full((2, 2, 6), -7, dtype=torch.bfloat16)) + quack_dispatch.gemm.assert_not_called() + quack_dispatch.fallback.assert_called_once() + assert len(graphs) == 1 + + +@pytest.mark.parametrize("valid_scales", [False, True]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_dispatch_custom_op_contract(quack_dispatch, valid_scales, with_bias): + a, b = _fp8_inputs() + # opcheck's CPU mutation checker uses allclose, which does not support FP8. + # The compile tests above exercise real FP8 inputs with these same mocks. + a, b = a.float(), b.float() + scale_a = torch.tensor([0.25 if valid_scales else torch.finfo(torch.float32).min]) + scale_b = torch.tensor([0.5]) + bias = torch.arange(6, dtype=torch.float32) if with_bias else None + torch.library.opcheck(quack_fp8._quack_fp8_scaled_mm, (a, b, scale_a, scale_b, torch.bfloat16, bias)) + + +def test_warmup_matches_inference_weight_layout(monkeypatch): + calls = [] + + def gemm(a, b, *, out, bias, alpha, tuned): + assert torch.is_inference_mode_enabled() + assert a.dtype == b.dtype == torch.float8_e4m3fn + assert a.is_contiguous() + assert b.stride() == (1, b.shape[0]) + assert bias is None + assert tuned is True + calls.append((a.shape[0], a.shape[1], b.shape[1])) + out.zero_() + + monkeypatch.setattr(quack_fp8, "_gemm_interface", SimpleNamespace(gemm=gemm)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + shapes = [(4, 8, 6), (2, 16, 8)] + quack_fp8.warmup_quack_fp8(shapes, device="cpu") + assert calls == shapes diff --git a/vllm_omni/quantization/quack_fp8.py b/vllm_omni/quantization/quack_fp8.py index e92515432..3ea5ab4d0 100644 --- a/vllm_omni/quantization/quack_fp8.py +++ b/vllm_omni/quantization/quack_fp8.py @@ -5,6 +5,9 @@ from __future__ import annotations import os +from contextlib import contextmanager +from importlib import import_module +from multiprocessing import current_process import torch from vllm.logger import init_logger @@ -52,6 +55,41 @@ def _set_persistent_cache_dir() -> None: os.environ["QUACK_CACHE_DIR"] = os.path.join(root, "vllm_omni", "quack") +def _configure_quack_compilation() -> None: + """Keep autotuning in daemon workers without starting compiler children.""" + try: + async_compile = import_module("quack.cache.async_compile") + except ModuleNotFoundError as exc: + # Older Quack releases compile synchronously and have no async pool. + if exc.name not in {"quack.cache", "quack.cache.async_compile"}: + raise + return + original_pool_scope = getattr(async_compile, "pool_scope", None) + suppress_pool = getattr(async_compile, "suppress_pool", None) + if not callable(original_pool_scope) or not callable(suppress_pool): + logger.warning( + "Quack async compilation API is unsupported: pool_scope and suppress_pool must be callable. " + "Skipping the compilation patch; autotuning in daemon workers may fail." + ) + return + if getattr(original_pool_scope, "_omni_daemon_safe", False): + return + + @contextmanager + def daemon_safe_pool_scope(): + # spawn can import this module before installing the child's daemon flag. + # Check at tuning time; suppress_pool keeps compilation in-process. + if current_process().daemon: + with suppress_pool(): + yield None + else: + with original_pool_scope() as pool: + yield pool + + daemon_safe_pool_scope._omni_daemon_safe = True + async_compile.pool_scope = daemon_safe_pool_scope + + def _load_quack(): global _gemm_interface if _gemm_interface is not None: @@ -72,6 +110,7 @@ def _load_quack(): torch2cute_dtype_map.setdefault(torch.float8_e4m3fn, cutlass.Float8E4M3FN) torch2cute_dtype_map.setdefault(torch.float8_e5m2, cutlass.Float8E5M2) + _configure_quack_compilation() _gemm_interface = gemm_interface logger.info("Quack FP8 fused-bias GEMM enabled (CuteDSL).") return gemm_interface @@ -109,6 +148,9 @@ def _scales_valid(scale_a: torch.Tensor, scale_b: torch.Tensor) -> bool: make ``alpha = scale_a * scale_b`` overflow to ``+inf`` and return an all-inf tile. Cache only positive results: a buffer still holding the sentinel is re-checked and picks up the fast path once the real scale is written. + + Called inside the dispatch custom op, so pointer checks and cache updates + neither specialize Dynamo graphs on each layer's addresses nor break them. """ key = (scale_a.data_ptr(), scale_b.data_ptr()) if key in _valid_scale_ptrs: @@ -121,6 +163,40 @@ def _scales_valid(scale_a: torch.Tensor, scale_b: torch.Tensor) -> bool: return ok +@torch.library.custom_op("vllm_omni::quack_fp8_scaled_mm", mutates_args=()) +def _quack_fp8_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + bias: torch.Tensor | None = None, +) -> torch.Tensor: + """Keep scale validation and runtime fallback inside one opaque graph node.""" + if scale_a.numel() == 1 and scale_b.numel() == 1 and _scales_valid(scale_a, scale_b): + try: + out = quack_scaled_fp8_mm(a, b, scale_a, scale_b, out_dtype, bias) + if out is not None: + return out + except Exception as exc: # noqa: BLE001 + logger.warning_once("Quack FP8 GEMM failed (%s); using FlashInfer.", exc) + + # Use the same fallback as FlashInferFP8ScaledMMLinearKernel, without + # passing its Python instance through the custom-op schema. Unpopulated + # scales also take this path: FlashInfer tolerates the profiling sentinel. + from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm + + out = flashinfer_scaled_fp8_mm(a, b, out_dtype=out_dtype, scale_a=scale_a, scale_b=scale_b, bias=bias) + # FlashInfer's separate bias add may promote the dtype. Both branches must + # match the fake kernel's output metadata, including with an FP32 bias. + return out.to(out_dtype) + + +@_quack_fp8_scaled_mm.register_fake +def _quack_fp8_scaled_mm_fake(a, b, scale_a, scale_b, out_dtype, bias=None): + return torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype) + + def install_quack_fp8_patch() -> None: if not quack_enabled(): return @@ -138,34 +214,26 @@ def install_quack_fp8_patch() -> None: return def apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape): # noqa: N803 - # An unpopulated scale makes alpha overflow to +inf; hand those calls to - # FlashInfer, whose bmm_fp8 tolerates them. - if As.numel() != 1 or Bs.numel() != 1 or not _scales_valid(As, Bs): - return original(self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape) - try: - out = quack_scaled_fp8_mm(A, B, As, Bs, out_dtype, bias) - if out is not None: - return out.view(*output_shape) - except Exception as exc: # noqa: BLE001 - logger.warning_once("Quack FP8 GEMM failed (%s); using FlashInfer.", exc) - return original(self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape) + return _quack_fp8_scaled_mm(A, B, As, Bs, out_dtype, bias).view(*output_shape) apply_scaled_mm._omni_quack_patched = True FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = apply_scaled_mm logger.info("Patched FlashInfer FP8 ScaledMM to use quack fused-bias GEMM.") +@torch.inference_mode() def warmup_quack_fp8( shapes: list[tuple[int, int, int]], device: str = "cuda", out_dtype: torch.dtype = torch.bfloat16, ) -> None: + """Warm no-bias GEMMs with the transposed weight layout used by vLLM.""" if _load_quack() is None: return scale = torch.ones(1, device=device, dtype=torch.float32) for m, k, n in shapes: a = torch.zeros(m, k, device=device, dtype=torch.float8_e4m3fn) - b = torch.zeros(k, n, device=device, dtype=torch.float8_e4m3fn) + b = torch.zeros(n, k, device=device, dtype=torch.float8_e4m3fn).t() quack_scaled_fp8_mm(a, b, scale, scale, out_dtype) if torch.cuda.is_available(): torch.accelerator.synchronize()