vLLM/vllm/utils/deep_gemm.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Compatibility wrapper for DeepGEMM API changes.

Users of vLLM should always import **only** these wrappers.
"""

import contextlib
import functools
import importlib
import os
from collections.abc import Callable
from enum import Enum
from typing import Any, NoReturn

import torch

import vllm.envs as envs
from vllm.logger import logger
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    get_fp8_min_max,
)
from vllm.platforms import current_platform
from vllm.utils.import_utils import has_deep_gemm
from vllm.utils.math_utils import cdiv

# Model type -> Blackwell device capability families to auto-disable DeepGemm on.
_DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES: dict[str, tuple[int, ...]] = {
    # Accuracy
    "qwen3_5_text": (100, 120),
    "qwen3_5_moe_text": (100, 120),
    # Latency
    "kolibri1": (120,),
}


def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
    """Check if DeepGemm should be auto-disabled for this model on Blackwell.

    Returns True if DeepGemm is known to hurt this model's accuracy or
    latency on the current Blackwell GPU family.
    """
    if model_type is None:
        return False
    return any(
        current_platform.is_device_capability_family(family)
        for family in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES.get(model_type, ())
    )


class DeepGemmQuantScaleFMT(Enum):
    # Float32 scales in Float32 tensor
    FLOAT32 = 0
    # Compute float32 scales and ceil the scales to UE8M0.
    # Keep the scales in Float32 tensor.
    FLOAT32_CEIL_UE8M0 = 1
    # Compute float32 scales and ceil the scales to UE8M0.
    # Pack the scales into a int32 tensor where each int32
    # element contains 4 scale values.
    UE8M0 = 2

    @classmethod
    def init_oracle_cache(cls) -> None:
        """Initialize the oracle decision and store it in the class cache."""
        cached = getattr(cls, "_oracle_cache", None)
        if cached is not None:
            return

        use_e8m0 = (
            envs.VLLM_USE_DEEP_GEMM_E8M0
            and is_deep_gemm_supported()
            and (_fp8_gemm_nt_impl is not None)
        )
        if not use_e8m0:
            cls._oracle_cache = cls.FLOAT32  # type: ignore
            return

        cls._oracle_cache = (  # type: ignore
            cls.UE8M0
            if (
                current_platform.is_device_capability_family(100)
                or current_platform.is_device_capability_family(120)
            )
            else cls.FLOAT32_CEIL_UE8M0
        )

    @classmethod
    def from_oracle(cls) -> "DeepGemmQuantScaleFMT":
        """Return the oracle decision, initializing it on first use.

        The cache is normally populated by ``_lazy_init()`` (e.g. during
        engine startup), but standalone consumers such as ``QuantFP8`` with an
        explicit ``use_ue8m0=True`` can reach this before any DeepGEMM kernel
        wrapper has run. Resolve the DeepGEMM symbols and initialize the
        decision here instead of asserting; without DeepGEMM this yields
        FLOAT32, matching ``is_deep_gemm_e8m0_used()``.
        """
        cached = getattr(cls, "_oracle_cache", None)
        if cached is None:
            _lazy_init()
            cls.init_oracle_cache()
            cached = cls._oracle_cache  # type: ignore[attr-defined]
        return cached


@functools.cache
def is_deep_gemm_supported() -> bool:
    """Return `True` if DeepGEMM is supported on the current platform.
    Currently, only Hopper and Blackwell GPUs are supported.
    """
    is_supported_arch = current_platform.support_deep_gemm()
    return is_supported_arch and envs.VLLM_USE_DEEP_GEMM and has_deep_gemm()


@functools.cache
def is_deep_gemm_e8m0_used() -> bool:
    """Return `True` if vLLM is configured to use DeepGEMM "
    "E8M0 scale on a Hopper or Blackwell-class GPU.
    """
    if not is_deep_gemm_supported():
        logger.debug_once(
            "DeepGEMM E8M0 disabled: DeepGEMM not supported on this system."
        )
        return False

    _lazy_init()

    if _fp8_gemm_nt_impl is None:
        logger.info_once("DeepGEMM E8M0 disabled: _fp8_gemm_nt_impl not found")
        return False

    if envs.VLLM_USE_DEEP_GEMM_E8M0:
        logger.info_once("DeepGEMM E8M0 enabled on current platform.")
        return True

    logger.info_once("DeepGEMM E8M0 disabled on current configuration.")
    return False


def _missing(*_: Any, **__: Any) -> NoReturn:
    """Placeholder for unavailable DeepGEMM backend."""
    raise RuntimeError(
        "DeepGEMM backend is unavailable in the current vLLM environment, "
        "or the available DeepGEMM package does not provide the required APIs "
        "for these kernels."
    )


_cublaslt_gemm_nt_impl: Callable[..., Any] | None = None
_fp8_gemm_nt_impl: Callable[..., Any] | None = None
_fp8_einsum_impl: Callable[..., Any] | None = None
_grouped_impl: Callable[..., Any] | None = None
_grouped_masked_impl: Callable[..., Any] | None = None
_grouped_fp4_impl: Callable[..., Any] | None = None
_fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
_fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None
_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None
_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None
_tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
_mega_mhc_impl: Callable[..., Any] | None = None
_bf16_mega_gate_impl: Callable[..., Any] | None = None
_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
_get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (
    None
)
_transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
_pack_ue8m0_to_int_impl: Callable[..., Any] | None = None
_get_mn_major_tma_aligned_packed_ue8m0_tensor_impl: Callable[..., Any] | None = None
_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl: (
    Callable[..., Any] | None
) = None


@functools.cache
def _import_deep_gemm():
    """Import the deep_gemm module.

    Prefers an externally installed ``deep_gemm`` package (so users can
    pin a specific version), then falls back to the vendored copy bundled
    in the vLLM wheel.

    Returns ``None`` when neither source is usable.
    """
    # 1. Try the external (pip-installed) package first.
    try:
        module = importlib.import_module("deep_gemm")
        logger.debug_once("Imported deep_gemm module from site-packages")
        return module
    except ImportError:
        logger.info_once(
            "deep_gemm not found in site-packages, "
            "trying vendored vllm.third_party.deep_gemm"
        )

    # 2. Fall back to the vendored copy bundled in the vLLM wheel.
    try:
        module = importlib.import_module("vllm.third_party.deep_gemm")
        logger.debug_once("Imported deep_gemm module from vllm.third_party.deep_gemm")
        return module
    except ImportError:
        logger.info_once("Vendored deep_gemm not found either")
    except Exception as e:
        # The vendored module may raise RuntimeError during _C.init()
        # if JIT include files are missing (e.g. incomplete wheel).
        logger.warning_once("Failed to import vendored deep_gemm: %s", e)

    return None


def _apply_pdl(mod, enable: bool = True) -> None:
    mod_name = getattr(mod, "__name__", str(mod))
    try:
        set_pdl_fn = getattr(mod, "set_pdl", None)
        if set_pdl_fn is None:
            return
        set_pdl_fn(enable)
        logger.info_once(
            "DeepGEMM PDL %s on %s.",
            "enabled" if enable else "disabled",
            mod_name,
        )
    except Exception as e:  # noqa: BLE001
        logger.warning_once("Failed to set DeepGEMM PDL on %s: %s", mod_name, e)


def _lazy_init() -> None:
    """Import deep_gemm and resolve symbols on first use."""
    global _cublaslt_gemm_nt_impl
    global _fp8_gemm_nt_impl, _fp8_einsum_impl
    global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
    global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
    global _get_paged_mqa_logits_metadata_impl
    global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl
    global _get_sparse_mqa_logits_metadata_impl
    global _get_paged_sparse_mqa_logits_metadata_impl
    global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl
    global _bf16_mega_gate_impl
    global _get_mn_major_tma_aligned_tensor_impl
    global _get_mk_alignment_for_contiguous_layout_impl
    global _get_theoretical_mk_alignment_for_contiguous_layout_impl
    global _transform_sf_into_required_layout_impl
    global _pack_ue8m0_to_int_impl
    global _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl
    global _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl
    # fast path
    if (
        _cublaslt_gemm_nt_impl is not None
        or _fp8_gemm_nt_impl is not None
        or _fp8_einsum_impl is not None
        or _grouped_impl is not None
        or _grouped_masked_impl is not None
        or _grouped_fp4_impl is not None
        or _fp8_fp4_mqa_logits_impl is not None
        or _fp8_fp4_paged_mqa_logits_impl is not None
        or _get_paged_mqa_logits_metadata_impl is not None
        or _tf32_hc_prenorm_gemm_impl is not None
        or _mega_mhc_impl is not None
        or _bf16_mega_gate_impl is not None
        or _get_mk_alignment_for_contiguous_layout_impl is not None
        or _transform_sf_into_required_layout_impl is not None
        or _pack_ue8m0_to_int_impl is not None
        or _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
        or _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
    ):
        return

    if not has_deep_gemm():
        return

    # Set up deep_gemm cache path
    DEEP_GEMM_JIT_CACHE_ENV_NAME = "DG_JIT_CACHE_DIR"
    if not os.environ.get(DEEP_GEMM_JIT_CACHE_ENV_NAME, None):
        os.environ[DEEP_GEMM_JIT_CACHE_ENV_NAME] = os.path.join(
            envs.VLLM_CACHE_ROOT, "deep_gemm"
        )

    _dg = _import_deep_gemm()
    if _dg is None:
        return

    # Enable PDL for DeepGEMM on architectures that support it (SM90+).
    if current_platform.is_arch_support_pdl():
        _apply_pdl(_dg, True)
    _cublaslt_gemm_nt_impl = getattr(_dg, "cublaslt_gemm_nt", None)
    _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None)
    _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None)
    _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None)
    # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so
    # an externally pinned older DeepGEMM still resolves.
    _grouped_masked_impl = getattr(
        _dg, "m_grouped_fp8_gemm_nt_masked", None
    ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None)
    _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None)
    # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that
    # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`.
    _fp8_fp4_mqa_logits_impl = getattr(_dg, "fp8_fp4_mqa_logits", None)
    _fp8_fp4_paged_mqa_logits_impl = getattr(_dg, "fp8_fp4_paged_mqa_logits", None)
    _get_paged_mqa_logits_metadata_impl = getattr(
        _dg, "get_paged_mqa_logits_metadata", None
    )
    # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only).
    _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None)
    _fp8_fp4_paged_sparse_mqa_logits_impl = getattr(
        _dg, "fp8_fp4_paged_sparse_mqa_logits", None
    )
    _get_sparse_mqa_logits_metadata_impl = getattr(
        _dg, "get_sparse_mqa_logits_metadata", None
    )
    _get_paged_sparse_mqa_logits_metadata_impl = getattr(
        _dg, "get_paged_sparse_mqa_logits_metadata", None
    )
    _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
    _mega_mhc_impl = getattr(_dg, "mega_mhc", None)
    _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None)
    _get_mn_major_tma_aligned_tensor_impl = getattr(
        _dg, "get_mn_major_tma_aligned_tensor", None
    )
    _get_mk_alignment_for_contiguous_layout_impl = getattr(
        _dg, "get_mk_alignment_for_contiguous_layout", None
    )
    _get_theoretical_mk_alignment_for_contiguous_layout_impl = getattr(
        _dg, "get_theoretical_mk_alignment_for_contiguous_layout", None
    )
    _transform_sf_into_required_layout_impl = getattr(
        _dg, "transform_sf_into_required_layout", None
    )
    _pack_ue8m0_to_int_impl = getattr(_dg, "pack_ue8m0_to_int", None)
    _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
        _dg, "get_mn_major_tma_aligned_packed_ue8m0_tensor", None
    )
    _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
        _dg, "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", None
    )
    DeepGemmQuantScaleFMT.init_oracle_cache()


def get_num_sms() -> int:
    _lazy_init()
    dg = _import_deep_gemm()
    if dg is None:
        raise RuntimeError("DeepGEMM is not available")
    return int(dg.get_num_sms())


def set_num_sms(num_sms: int) -> None:
    _lazy_init()
    dg = _import_deep_gemm()
    if dg is None:
        raise RuntimeError("DeepGEMM is not available")
    dg.set_num_sms(num_sms)


def get_mk_alignment_for_contiguous_layout() -> list[int]:
    _lazy_init()
    if _get_mk_alignment_for_contiguous_layout_impl is None:
        return _missing()
    mk_align_size = _get_mk_alignment_for_contiguous_layout_impl()
    return [mk_align_size, mk_align_size]


def get_theoretical_mk_alignment_for_contiguous_layout(
    expected_m: int | None = None,
    num_groups: int | None = None,
) -> int:
    """Per-call optimal M alignment for grouped contiguous GEMMs.

    `expected_m` is the TOTAL routed tokens (sum across experts, typically
    M × num_topk). `num_groups` is the number of experts on this rank.
    The helper divides to recover per-expert em and picks an alignment based
    on data-driven thresholds (see deep_gemm runtime.hpp comments).

    Older callers that omit `num_groups` are interpreted as passing already
    per-expert em (legacy behaviour preserved for backward compat).
    """
    _lazy_init()
    if _get_theoretical_mk_alignment_for_contiguous_layout_impl is None:
        return _missing()
    if num_groups is None:
        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(expected_m)
    if num_groups <= 0:
        raise ValueError(f"num_groups must be positive, got {num_groups}")
    try:
        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(
            expected_m, num_groups
        )
    except TypeError:
        per_group_m = None if expected_m is None else cdiv(expected_m, num_groups)
        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(per_group_m)


def set_mk_alignment_for_contiguous_layout(value: int) -> None:
    """Set DeepGEMM's BLOCK_M cap for grouped contiguous GEMMs.

    The DG heuristic constrains BLOCK_M ≤ this value when picking a kernel
    layout. Use this in concert with `compute_aligned_M_and_alignment`'s
    per-call alignment so the workspace's per-expert padding matches the
    kernel's BLOCK_M; a mismatch leads to the scheduler reading the wrong
    expert_id from `m_indices` at `m_block_idx * BLOCK_M` stride and
    OOB-indexing the B-weights tensor (manifests as IMA under CUDA-graph
    replay).
    """
    _lazy_init()
    dg = _import_deep_gemm()
    if dg is None:
        raise RuntimeError("DeepGEMM is not available")
    dg.set_mk_alignment_for_contiguous_layout(value)


@contextlib.contextmanager
def mk_alignment_scope(value: int):
    """Temporarily set DeepGEMM's BLOCK_M cap, restoring on exit.

    Use around a sequence of grouped-contiguous GEMM calls whose workspace
    is padded to `value` (typically the per_call_align returned by
    `compute_aligned_M_and_alignment`).
    """
    prev = get_mk_alignment_for_contiguous_layout()[0]
    set_mk_alignment_for_contiguous_layout(value)
    try:
        yield
    finally:
        set_mk_alignment_for_contiguous_layout(prev)


def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
    """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor."""
    _lazy_init()
    if _get_mn_major_tma_aligned_tensor_impl is None:
        return _missing()
    return _get_mn_major_tma_aligned_tensor_impl(x)


def pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor:
    """Pack 4 UE8M0 (uint8) scales into one int32.

    DeepGEMM's SM100/SM120 FP8/FP4 kernels accept either ``float32`` scales
    (legacy format, 4 B/scale) or ``int32`` packed UE8M0 scales (1 B/scale
    after 4:1 packing — 4× smaller than the legacy fp32 representation).
    """
    _lazy_init()
    if _pack_ue8m0_to_int_impl is None:
        return _missing()
    return _pack_ue8m0_to_int_impl(x)


def get_mn_major_tma_aligned_packed_ue8m0_tensor(x: torch.Tensor) -> torch.Tensor:
    """Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the
    DeepGEMM kernels consume directly. 16× smaller than the fp32 legacy SF
    format. Use for non-grouped 2D scale tensors.
    """
    _lazy_init()
    if _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
        return _missing()
    return _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl(x)


def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(
    sf: torch.Tensor,
    ks_tensor: torch.Tensor,
    ks: list[int],
    gran_k: int,
) -> torch.Tensor:
    """Grouped (3D, expert-batched) variant of
    ``get_mn_major_tma_aligned_packed_ue8m0_tensor``. Use for MoE weight
    scale tensors of shape ``(num_experts, mn, k_scale)``.
    """
    _lazy_init()
    if _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
        return _missing()
    return _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl(
        sf, ks_tensor, ks, gran_k
    )


def cublaslt_gemm_nt(*args, **kwargs):
    _lazy_init()
    if _cublaslt_gemm_nt_impl is None:
        return _missing(*args, **kwargs)
    return _cublaslt_gemm_nt_impl(*args, **kwargs)


def bf16_mega_gate(
    x: torch.Tensor,
    weight: torch.Tensor,
    num_topk: int,
    *,
    scoring_func: str,
    routed_scaling_factor: float,
    ep_rank: int,
    bias: torch.Tensor | None = None,
    image_bias: torch.Tensor | None = None,
    image_token_mask: torch.Tensor | None = None,
    fix_routing_mask: torch.Tensor | None = None,
    unmapped_topk_idx: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run DeepGEMM Mega Gate.

    The output order follows vLLM's router contract: weights, then indices.
    Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA
    graph capture.
    """
    _lazy_init()
    if _bf16_mega_gate_impl is None:
        raise RuntimeError(
            "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate."
        )

    topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device)
    topk_weights = torch.empty(
        (x.shape[0], num_topk), dtype=torch.float32, device=x.device
    )
    _bf16_mega_gate_impl(
        x,
        weight,
        num_topk,
        use_shared_as_routed=False,
        num_shared_experts=0,
        routed_scaling_factor=routed_scaling_factor,
        ep_rank=ep_rank,
        scoring_func=scoring_func,
        bias=bias,
        image_bias=image_bias,
        image_token_mask=image_token_mask,
        fix_routing_mask=fix_routing_mask,
        unmapped_topk_idx=unmapped_topk_idx,
        out=(topk_idx, topk_weights),
    )
    return topk_weights, topk_idx


def fp8_gemm_nt(*args, **kwargs):
    _lazy_init()
    if _fp8_gemm_nt_impl is None:
        return _missing(*args, **kwargs)
    if "is_deep_gemm_e8m0_used" in kwargs:
        use_ue8m0 = kwargs["is_deep_gemm_e8m0_used"]
        del kwargs["is_deep_gemm_e8m0_used"]
    else:
        use_ue8m0 = is_deep_gemm_e8m0_used()
    return _fp8_gemm_nt_impl(*args, disable_ue8m0_cast=not use_ue8m0, **kwargs)


def fp8_einsum(*args, **kwargs):
    _lazy_init()
    if _fp8_einsum_impl is None:
        return _missing(*args, **kwargs)
    return _fp8_einsum_impl(*args, **kwargs)


def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
    _lazy_init()
    if _grouped_impl is None:
        return _missing(*args, **kwargs)
    return _grouped_impl(
        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
    )


def m_grouped_fp8_fp4_gemm_nt_contiguous(*args, **kwargs):
    _lazy_init()
    if _grouped_fp4_impl is None:
        return _missing(*args, **kwargs)
    return _grouped_fp4_impl(
        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
    )


def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):
    _lazy_init()
    if _grouped_masked_impl is None:
        return _missing(*args, **kwargs)
    return _grouped_masked_impl(
        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
    )


def transform_sf_into_required_layout(*args, **kwargs):
    _lazy_init()
    if _transform_sf_into_required_layout_impl is None:
        return _missing(*args, **kwargs)
    return _transform_sf_into_required_layout_impl(
        *args, disable_ue8m0_cast=not is_deep_gemm_e8m0_used(), **kwargs
    )


@functools.cache
def mqa_logits_use_mx_scales() -> bool:
    """Whether the MQA logits kernels take only MX-scaled inputs.

    On SM100 they need MXFP8/MXFP4 Q and K with packed UE8M0 scales and BF16
    weights, and return BF16 logits. FP8 indexer K caches then store each
    token's power-of-two scale as four equal UE8M0 bytes instead of an fp32.
    """
    return current_platform.is_cuda() and current_platform.is_device_capability_family(
        100
    )


# Four packed UE8M0 bytes encoding 1.0.
_UNIT_UE8M0_X4 = 0x7F7F7F7F
# Grown buffers are kept alive: CUDA graphs captured earlier still read them.
_unit_mx_scales: dict[torch.device, list[torch.Tensor]] = {}


def mqa_logits_mx_inputs(
    q: torch.Tensor, q_scale: torch.Tensor | None, weights: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor]:
    """Adapt indexer Q scales and weights to the SM100 MQA logits kernels.

    FP8 Q folds its per-(token, head) power-of-two scale into ``weights``, so
    it is passed with unit MX scales, which is exact since relu commutes with
    a positive scale. Returns ``(q_scale, weights)`` unchanged off SM100.
    """
    if not mqa_logits_use_mx_scales():
        return q_scale, weights
    if q_scale is None:
        numel = q.numel() // q.shape[-1]
        bufs = _unit_mx_scales.setdefault(q.device, [])
        if not bufs or bufs[-1].numel() < numel:
            size = max(numel, 2 * bufs[-1].numel() if bufs else 0)
            bufs.append(
                torch.full((size,), _UNIT_UE8M0_X4, dtype=torch.int32, device=q.device)
            )
        q_scale = bufs[-1][:numel].view(q.shape[:-1])
    return q_scale, weights.to(torch.bfloat16)


def fp8_fp4_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv: tuple[torch.Tensor, torch.Tensor],
    weights: torch.Tensor,
    cu_seqlen_ks: torch.Tensor,
    cu_seqlen_ke: torch.Tensor,
    max_seqlen_k: int,
) -> torch.Tensor:
    """Compute MQA logits for a single sequence without KV paging.

    Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
    ``q = (values, scales_or_None)``. SM100 takes MXFP8 or MXFP4 with packed
    UE8M0 block scales; SM90/SM120 also take FP8 Q with ``q_scale=None``
    (per-token scale folded into ``weights``).

    Args:
        q: Tuple ``(q_values, q_scale)``, q_values of shape [M, H, D] (packed
            to D/2 bytes for MXFP4) and q_scale int32 [M, H] or None.
        kv: Tuple ``(k_values, k_scales)`` of shape [N, D] and [N]; k_scales
            is int32 packed UE8M0 with MX scales, fp32 otherwise.
        weights: Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere.
        cu_seqlen_ks: Start indices (inclusive) for valid K per query
            position, shape [M], dtype int32.
        cu_seqlen_ke: End indices (exclusive) for valid K per query
            position, shape [M], dtype int32.
        max_seqlen_k: Upper bound on ``cu_seqlen_ke - cu_seqlen_ks``.

    Returns:
        Compressed logits of shape [M, max_seqlen_k]: row ``i`` holds K
        positions ``[cu_seqlen_ks[i], cu_seqlen_ke[i])`` starting at column
        zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere.

    """
    _lazy_init()
    if _fp8_fp4_mqa_logits_impl is None:
        return _missing()
    return _fp8_fp4_mqa_logits_impl(
        q,
        kv,
        weights,
        cu_seqlen_ks,
        cu_seqlen_ke,
        max_seqlen_k,
    )


def native_next_n_supported(next_n: int) -> bool:
    """Whether the paged MQA logits kernel takes `next_n` Q rows per request.

    SM90 implements only {1, 2, 4}; SM100 and SM120 schedule any `next_n` via
    multi-atom tiles. Unsupported values must be flattened to one row per query.
    """
    if current_platform.is_device_capability_family(90):
        return next_n in (1, 2, 4)
    return True


def _paged_mqa_logits_schedule_slots(num_sms: int, next_n: int) -> int:
    """Scheduler tasks the paged MQA logits kernel launches.

    SM90 `next_n=4` runs one task per 2-CTA multicast cluster rather than per
    SM, and `fp8_fp4_paged_mqa_logits` asserts its metadata is sized to match.
    """
    num_kv_multicast = (
        2 if next_n == 4 and current_platform.is_device_capability_family(90) else 1
    )
    return num_sms // num_kv_multicast


def get_paged_mqa_logits_metadata(
    context_lens: torch.Tensor,
    block_size: int,
    num_sms: int,
    indices: torch.Tensor | None = None,
) -> torch.Tensor:
    """Build scheduling metadata for paged MQA logits.

    Args:
        context_lens: Tensor of shape [B, next_n], dtype int32; effective
            context length per Q row.
        block_size: KV-cache block size in tokens (e.g., 64).
        num_sms: Number of SMs available. 132 for Hopper
        indices: Optional request index for each varlen row.

    Returns:
        Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits`
        to schedule work across SMs.

    """
    _lazy_init()
    if _get_paged_mqa_logits_metadata_impl is None:
        return _missing()
    next_n = context_lens.shape[1] if context_lens.dim() == 2 else 1
    num_slots = _paged_mqa_logits_schedule_slots(num_sms, next_n)
    kwargs = {} if indices is None else {"indices": indices}
    return _get_paged_mqa_logits_metadata_impl(
        context_lens, block_size, num_slots, **kwargs
    )


def fp8_fp4_paged_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor | None],
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    context_lens: torch.Tensor,
    block_tables: torch.Tensor,
    schedule_metadata: torch.Tensor,
    max_model_len: int,
    indices: torch.Tensor | None = None,
) -> torch.Tensor:
    """Compute MQA logits using a paged KV-cache.

    Unified FP8/FP4 dispatch — the underlying DeepGEMM kernel takes
    ``q = (values, scales_or_None)``; see `fp8_fp4_mqa_logits` for the
    per-architecture formats. SM100 requires the varlen form (``indices``
    with next_n == 1).

    Args:
        q: Tuple ``(q_values, q_scale)``, q_values of shape
            [B, next_n, H, D] (packed to D/2 bytes for MXFP4) and q_scale
            int32 [B, next_n, H] or None.
        kv_cache: Paged KV-cache of shape [num_blocks, block_size, 1, D+4]
            (D/2+4 for MXFP4), dtype `torch.uint8`; the last 4 bytes per
            (block, pos) hold the int32 packed UE8M0 scales with MX scales,
            or the fp32 dequant scale otherwise.
        weights: Tensor of shape [B * next_n, H], bf16 on SM100 and fp32
            elsewhere.
        context_lens: Tensor of shape [B], dtype int32; effective context length
            for each batch element.
        block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
            block indices to physical blocks in the paged cache.
        schedule_metadata: Returned by `get_paged_mqa_logits_metadata`;
            used to distribute work across SMs.
        max_model_len: Maximum sequence length used to size the logits output.
        indices: Optional request index for each varlen row.

    Returns:
        Logits tensor of shape [B * next_n, max_model_len]; columns past each
        row's context length are not cleaned. bf16 on SM100, fp32 elsewhere.

    """
    _lazy_init()
    if _fp8_fp4_paged_mqa_logits_impl is None:
        return _missing()
    # DeepGEMM asserts block_tables.stride(-1)==1. A trailing size-1 dim
    # (e.g. block_table shape [B,1] for short seqs under a large block_size)
    # can be a transposed view where .contiguous() is a no-op (torch treats the
    # size-1 dim's stride as irrelevant) yet stride(-1)!=1, failing the kernel.
    # clone to contiguous format to force stride(-1)==1.
    if block_tables.dim() >= 2 and block_tables.stride(-1) != 1:
        block_tables = block_tables.clone(memory_format=torch.contiguous_format)
    kwargs = {} if indices is None else {"indices": indices}
    return _fp8_fp4_paged_mqa_logits_impl(
        q,
        kv_cache,
        weights,
        context_lens,
        block_tables,
        schedule_metadata,
        max_model_len,
        **kwargs,
    )


def get_sparse_mqa_logits_metadata(
    cu_seqlen_ks: torch.Tensor,
    cu_seqlen_ke: torch.Tensor,
    num_kv_tokens: int,
    sparse_kv_block_indices: torch.Tensor,
    qk_dtype: torch.dtype,
    sparse_block_kv: int,
) -> torch.Tensor:
    """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`.

    Always uses the unaligned-ks variant, which degenerates to the aligned
    variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never
    need a host-side sync to pick one.

    Args:
        cu_seqlen_ks: Per-row K range start bounds in the packed KV
            workspace, shape [num_q_tokens], dtype int32.
        cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace,
            shape [num_q_tokens], dtype int32.
        num_kv_tokens: Total KV tokens in the packed workspace.
        sparse_kv_block_indices: Per-row candidate block ids, shape
            [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's
            valid prefix must be sorted ascending (repeats are tolerated but
            wasteful); pad by repeating the last valid block. Block ``i``
            covers tokens
            ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``.
        qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4,
            ``torch.float8_e4m3fn`` for FP8).
        sparse_block_kv: Tokens per sparse block, 8 or 16.

    """
    _lazy_init()
    if _get_sparse_mqa_logits_metadata_impl is None:
        return _missing()
    return _get_sparse_mqa_logits_metadata_impl(
        cu_seqlen_ks,
        cu_seqlen_ke,
        num_kv_tokens,
        sparse_kv_block_indices,
        qk_dtype,
        sparse_block_kv,
        use_unaligned_ks=True,
    )


def get_paged_sparse_mqa_logits_metadata(
    context_lens: torch.Tensor,
    block_table: torch.Tensor,
    indices: torch.Tensor,
    page_kv: int,
    sparse_kv_block_indices: torch.Tensor,
    qk_dtype: torch.dtype,
    sparse_block_kv: int,
) -> torch.Tensor:
    """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`.

    Args:
        context_lens: Per-row context length, shape [num_q_tokens], int32.
        block_table: Per-row page table, shape [num_q_tokens, max_pages],
            int32; rows of queries paired with the same request must be
            identical (expand it per query row before calling).
        indices: Request index for each query row, shape [num_q_tokens].
        page_kv: Tokens per KV page; must be a multiple of sparse_block_kv.
        sparse_kv_block_indices: Per-row logical (context-relative) candidate
            block ids, [num_q_tokens, num_max_sparse_blocks], int32.
        qk_dtype: dtype of the packed Q values.
        sparse_block_kv: Tokens per sparse block, 8 or 16.

    """
    _lazy_init()
    if _get_paged_sparse_mqa_logits_metadata_impl is None:
        return _missing()
    return _get_paged_sparse_mqa_logits_metadata_impl(
        context_lens,
        block_table,
        indices,
        page_kv,
        sparse_kv_block_indices,
        qk_dtype,
        sparse_block_kv,
    )


def fp8_fp4_sparse_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor],
    kv: tuple[torch.Tensor, torch.Tensor],
    weights: torch.Tensor,
    metadata: torch.Tensor,
    num_max_sparse_blocks: int,
    sparse_block_kv: int,
) -> torch.Tensor:
    """Compute MQA logits only at the candidate blocks (prefill path).

    Args:
        q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed
            as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H]
            int32 (mandatory, unlike the dense kernels).
        kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0
            int32 scales.
        weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in).
        metadata: From `get_sparse_mqa_logits_metadata`.
        num_max_sparse_blocks: Width of each sparse-index row.
        sparse_block_kv: Tokens per sparse block, 8 or 16.

    Returns:
        bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv];
        column ``j * sparse_block_kv + o`` scores the token at
        ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``.

    """
    _lazy_init()
    if _fp8_fp4_sparse_mqa_logits_impl is None:
        return _missing()
    return _fp8_fp4_sparse_mqa_logits_impl(
        q,
        kv,
        weights,
        metadata,
        num_max_sparse_blocks,
        sparse_block_kv,
        use_unaligned_ks=True,
    )


def fp8_fp4_paged_sparse_mqa_logits(
    q: tuple[torch.Tensor, torch.Tensor],
    kv_cache: torch.Tensor,
    weights: torch.Tensor,
    metadata: torch.Tensor,
    num_max_sparse_blocks: int,
    sparse_block_kv: int,
) -> torch.Tensor:
    """Compute MQA logits only at the candidate blocks (paged decode path).

    Args:
        q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D].
        kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8,
            page stride 256B-aligned.
        weights: [num_q_tokens, H] ``torch.bfloat16``.
        metadata: From `get_paged_sparse_mqa_logits_metadata`.
        num_max_sparse_blocks: Candidate blocks per row.
        sparse_block_kv: Tokens per sparse block, 8 or 16.

    Returns:
        bf16 logits of shape
        [num_q_tokens, num_max_sparse_blocks * sparse_block_kv].

    """
    _lazy_init()
    if _fp8_fp4_paged_sparse_mqa_logits_impl is None:
        return _missing()
    return _fp8_fp4_paged_sparse_mqa_logits_impl(
        q,
        kv_cache,
        weights,
        metadata,
        num_max_sparse_blocks,
        sparse_block_kv,
    )


def tf32_hc_prenorm_gemm(
    x: torch.Tensor,
    fn: torch.Tensor,
    out: torch.Tensor,
    sqrsum: torch.Tensor,
    num_split: int,
) -> torch.Tensor:
    """Perform the following computation:
        out = x.float() @ fn.T
        sqrsum = x.float().square().sum(-1)

    See the caller function for shape requirement
    """
    _lazy_init()
    if _tf32_hc_prenorm_gemm_impl is None:
        return _missing()
    return _tf32_hc_prenorm_gemm_impl(
        x,
        fn,
        out,
        sqrsum,
        num_split,
    )


def mega_mhc(*args: Any, **kwargs: Any) -> None:
    """Run DeepGEMM Mega mHC with caller-owned output tensors."""
    _lazy_init()
    if _mega_mhc_impl is None:
        _missing()
        return
    _mega_mhc_impl(*args, **kwargs)


def _ceil_to_ue8m0(x: torch.Tensor):
    return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))


def _align(x: int, y: int) -> int:
    return cdiv(x, y) * y


# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19
def get_tma_aligned_size(x: int, element_size: int) -> int:
    return _align(x, 16 // element_size)


DEFAULT_BLOCK_SIZE = [128, 128]


# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/dd6ed14acbc7445dcef224248a77ab4d22b5f240/deep_gemm/utils/math.py#L38
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
def per_block_cast_to_fp8(
    x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False
) -> tuple[torch.Tensor, torch.Tensor]:
    fp8_dtype = current_platform.fp8_dtype()
    assert x.dim() == 2
    m, n = x.shape
    block_m, block_n = block_size
    x_padded = torch.zeros(
        (_align(m, block_m), _align(n, block_n)), dtype=x.dtype, device=x.device
    )
    x_padded[:m, :n] = x
    x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
    x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
    _, fp8_max = get_fp8_min_max()
    sf = x_amax / fp8_max
    sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
    x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype)
    return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
        x_view.size(0), x_view.size(2)
    )


def calc_diff(x: torch.Tensor, y: torch.Tensor):
    """Return a global difference metric for unit tests.

    DeepGEMM kernels on Blackwell/B200 currently exhibit noticeable per-element
    error, causing `torch.testing.assert_close` to fail.  Instead of checking
    every element, we compute a cosine-style similarity over the whole tensor
    and report `1 - sim`.  Once kernel accuracy improves this helper can be
    removed.
    """
    x, y = x.double(), y.double()
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return 1 - sim


def should_use_deepgemm_for_fp8_linear(
    output_dtype: torch.dtype,
    weight_shape: tuple[int, int],
    supports_deep_gemm: bool | None = None,
):
    if supports_deep_gemm is None:
        supports_deep_gemm = is_deep_gemm_supported()

    # Verify DeepGEMM N/K dims requirements
    # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul
    # test inside kernels/quantization/test_block_fp8.py
    N_MULTIPLE = 64
    K_MULTIPLE = 128

    return (
        supports_deep_gemm
        and output_dtype == torch.bfloat16
        and weight_shape[0] % N_MULTIPLE == 0
        and weight_shape[1] % K_MULTIPLE == 0
    )


__all__ = [
    "calc_diff",
    "bf16_mega_gate",
    "DeepGemmQuantScaleFMT",
    "fp8_gemm_nt",
    "fp8_einsum",
    "m_grouped_fp8_gemm_nt_contiguous",
    "m_grouped_fp8_fp4_gemm_nt_contiguous",
    "fp8_m_grouped_gemm_nt_masked",
    "fp8_fp4_mqa_logits",
    "fp8_fp4_paged_mqa_logits",
    "get_paged_mqa_logits_metadata",
    "mqa_logits_mx_inputs",
    "mqa_logits_use_mx_scales",
    "native_next_n_supported",
    "per_block_cast_to_fp8",
    "is_deep_gemm_e8m0_used",
    "is_deep_gemm_supported",
    "mega_mhc",
    "get_num_sms",
    "set_num_sms",
    "should_use_deepgemm_for_fp8_linear",
    "get_col_major_tma_aligned_tensor",
    "get_mk_alignment_for_contiguous_layout",
    "get_theoretical_mk_alignment_for_contiguous_layout",
    "pack_ue8m0_to_int",
    "get_mn_major_tma_aligned_packed_ue8m0_tensor",
    "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor",
]