# 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_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 _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(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.
    """
    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


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

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