240 lines
8.7 KiB
Python
240 lines
8.7 KiB
Python
# 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
|
|
from contextlib import contextmanager
|
|
from importlib import import_module
|
|
from multiprocessing import current_process
|
|
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_TRUTHY = {"1", "true", "yes", "on"}
|
|
_gemm_interface = None
|
|
|
|
|
|
def _is_quack_capable() -> bool:
|
|
"""quack's CuteDSL FP8 / block-scaled MMA is built on the 5th-gen tensor-core
|
|
``tcgen05`` instruction family, which is datacenter-Blackwell only
|
|
(``sm_100a`` / ``sm_101a`` / ``sm_103a``, compute capability ``10.x``).
|
|
|
|
Workstation/consumer Blackwell (``sm_120`` / ``sm_121``, compute capability
|
|
``12.x``, e.g. RTX PRO 6000 / RTX 50-series) lacks ``tcgen05``, so quack can
|
|
never run there — CuteDSL rejects the arch and every GEMM falls back to
|
|
FlashInfer one call at a time, which is catastrophically slow. Those GPUs
|
|
have working native FlashInfer FP8 kernels, so default quack off for them.
|
|
"""
|
|
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_quack_capable()
|
|
|
|
|
|
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 _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:
|
|
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)
|
|
|
|
_configure_quack_compilation()
|
|
_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(a, b, out=out, bias=bias, alpha=alpha, tuned=True)
|
|
return out
|
|
|
|
|
|
_valid_scale_ptrs: set[tuple[int, int]] = set()
|
|
|
|
|
|
def _scales_valid(scale_a: torch.Tensor, scale_b: torch.Tensor) -> bool:
|
|
"""True when both per-tensor scales are finite and positive.
|
|
|
|
vLLM initializes a per-tensor FP8 scale to ``finfo(float32).min`` and fills it at
|
|
weight-load time, so a call made before that (the dummy profiling forward) would
|
|
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:
|
|
return True
|
|
ok = bool(
|
|
torch.isfinite(scale_a).all() and torch.isfinite(scale_b).all() and (scale_a > 0).all() and (scale_b > 0).all()
|
|
)
|
|
if ok:
|
|
_valid_scale_ptrs.add(key)
|
|
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
|
|
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
|
|
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(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()
|