[Quant][Perf] Default Blackwell FP8 GEMM to quack CuteDSL fused-bias kernel (#4241)

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>

docs/user_guide/quantization/fp8.md

@@ -28,6 +28,50 @@ in deep DiT blocks.
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
 
+### Faster FP8 GEMM on Blackwell (quack)
+
+On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
+applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
+DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
+lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
+that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
+
+```bash
+# CUDA 12.9
+pip install vllm-omni[quack]
+
+# CUDA 13.x
+pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
+```
+
+It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
+on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
+Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
+not installed, FP8 still works — it just keeps the unfused FlashInfer path.
+
+#### Compile cache and warmup
+
+quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
+the first time across all autotuned configs). The compiled `.o` files are cached on
+disk and reused on later runs, so this is a one-time cost — **not** per request.
+
+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.
+
+To pre-warm specific shapes (e.g. at image build time):
+
+```python
+from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
+# (M, K, N) per linear; M = number of tokens for your resolution/frame count
+warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
+```
+
+> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
+> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
+
 ## Model Type Support
 
 ### Diffusion Model (Qwen-Image, Wan2.2)

pyproject.toml

@@ -90,6 +90,12 @@ minicpmo = [
     "stepaudio2-minicpmo",
 ]
 
+# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
+# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
+quack = [
+    "quack-kernels>=0.3.11",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

vllm_omni/patch.py

@@ -257,3 +257,15 @@ def _patch_flashinfer_fp8_scaled_mm_output_shape():
 
 
 _patch_flashinfer_fp8_scaled_mm_output_shape()
+
+
+def _patch_fp8_use_quack_fused_bias():
+    try:
+        from vllm_omni.quantization.quack_fp8 import install_quack_fp8_patch
+
+        install_quack_fp8_patch()
+    except Exception:  # noqa: BLE001
+        pass
+
+
+_patch_fp8_use_quack_fused_bias()

vllm_omni/quantization/quack_fp8.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FP8 scaled-MM with a fused bias epilogue, backed by quack's CuteDSL GEMM."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_TRUTHY = {"1", "true", "yes", "on"}
+_gemm_interface = None
+
+
+def _is_blackwell() -> bool:
+    try:
+        if not torch.cuda.is_available():
+            return False
+        return torch.cuda.get_device_capability()[0] >= 10
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def quack_enabled() -> bool:
+    override = os.environ.get("VLLM_OMNI_USE_QUACK_FP8")
+    if override is not None:
+        return override.lower() in _TRUTHY
+    return _is_blackwell()
+
+
+def _set_persistent_cache_dir() -> None:
+    if os.environ.get("QUACK_CACHE_DIR"):
+        return
+    root = (
+        os.environ.get("VLLM_CACHE_ROOT")
+        or os.environ.get("XDG_CACHE_HOME")
+        or os.path.join(os.path.expanduser("~"), ".cache")
+    )
+    os.environ["QUACK_CACHE_DIR"] = os.path.join(root, "vllm_omni", "quack")
+
+
+def _load_quack():
+    global _gemm_interface
+    if _gemm_interface is not None:
+        return _gemm_interface or None
+    try:
+        _set_persistent_cache_dir()
+
+        import cutlass
+        import cutlass.base_dsl
+        import cutlass.base_dsl.arch as arch
+
+        if not hasattr(cutlass.base_dsl, "Arch"):
+            cutlass.base_dsl.Arch = arch.Arch
+
+        import quack.gemm_interface as gemm_interface
+        from quack.cute_dsl_utils import torch2cute_dtype_map
+
+        torch2cute_dtype_map.setdefault(torch.float8_e4m3fn, cutlass.Float8E4M3FN)
+        torch2cute_dtype_map.setdefault(torch.float8_e5m2, cutlass.Float8E5M2)
+
+        _gemm_interface = gemm_interface
+        logger.info("Quack FP8 fused-bias GEMM enabled (CuteDSL).")
+        return gemm_interface
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("Quack FP8 unavailable, using FlashInfer: %s", exc)
+        _gemm_interface = False
+        return None
+
+
+def quack_scaled_fp8_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 | None:
+    gemm = _load_quack()
+    if gemm is None:
+        return None
+    out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
+    alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
+    gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
+    return out
+
+
+def install_quack_fp8_patch() -> None:
+    if not quack_enabled():
+        return
+    if _load_quack() is None:
+        return
+    try:
+        from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
+            FlashInferFP8ScaledMMLinearKernel,
+        )
+    except ImportError:
+        return
+
+    original = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
+    if getattr(original, "_omni_quack_patched", False):
+        return
+
+    def apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape):  # noqa: N803
+        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)
+
+    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.")
+
+
+def warmup_quack_fp8(
+    shapes: list[tuple[int, int, int]],
+    device: str = "cuda",
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> None:
+    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)
+        quack_scaled_fp8_mm(a, b, scale, scale, out_dtype)
+    if torch.cuda.is_available():
+        torch.accelerator.synchronize()