diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index c8f14548c..1abcd4389 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 19cdc0a26..09b8e23ca 100644 --- a/pyproject.toml +++ b/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", diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index c343276b1..59d697c16 100644 --- a/vllm_omni/patch.py +++ b/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() diff --git a/vllm_omni/quantization/quack_fp8.py b/vllm_omni/quantization/quack_fp8.py new file mode 100644 index 000000000..ba68d078d --- /dev/null +++ b/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()