[bugfix] Fix FP8 Quack kernels when using Context Parallelism and reduce number of recompilations (#7458)

Signed-off-by: Maciej Bala <mbala@nvidia.com>

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):
 

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

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()