vllm.utils.deep_gemm

Compatibility wrapper for DeepGEMM API changes.

Users of vLLM should always import only these wrappers.

Classes:

• DeepGemmQuantScaleFMT –

Functions:

• bf16_mega_gate –

  Run DeepGEMM Mega Gate.

• calc_diff –

  Return a global difference metric for unit tests.

• fp8_fp4_mqa_logits –

  Compute MQA logits for a single sequence without KV paging.

• fp8_fp4_paged_mqa_logits –

  Compute MQA logits using a paged KV-cache.

• get_col_major_tma_aligned_tensor –

  Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor.

• get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor –

  Grouped (3D, expert-batched) variant of

• get_mn_major_tma_aligned_packed_ue8m0_tensor –

  Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the

• get_paged_mqa_logits_metadata –

  Build scheduling metadata for paged MQA logits.

• get_theoretical_mk_alignment_for_contiguous_layout –

  Per-call optimal M alignment for grouped contiguous GEMMs.

• is_deep_gemm_e8m0_used –

  Return True if vLLM is configured to use DeepGEMM "

• is_deep_gemm_supported –

  Return True if DeepGEMM is supported on the current platform.

• mega_mhc –

  Run DeepGEMM Mega mHC with caller-owned output tensors.

• mqa_logits_mx_inputs –

  Adapt indexer Q scales and weights to the SM100 MQA logits kernels.

• mqa_logits_use_mx_scales –

  Whether the MQA logits kernels take only MX-scaled inputs.

• native_next_n_supported –

  Whether the paged MQA logits kernel takes next_n Q rows per request.

• pack_ue8m0_to_int –

  Pack 4 UE8M0 (uint8) scales into one int32.

DeepGemmQuantScaleFMT

Bases: Enum

Methods:

• from_oracle –

  Return the oracle decision, initializing it on first use.

• init_oracle_cache –

  Initialize the oracle decision and store it in the class cache.

Source code in vllm/utils/deep_gemm.py

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

from_oracle() classmethod

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

Source code in vllm/utils/deep_gemm.py

@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

init_oracle_cache() classmethod

Initialize the oracle decision and store it in the class cache.

Source code in vllm/utils/deep_gemm.py

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

_import_deep_gemm() cached

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.

Source code in vllm/utils/deep_gemm.py

@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

_lazy_init()

Import deep_gemm and resolve symbols on first use.

Source code in vllm/utils/deep_gemm.py

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

_missing(*_, **__)

Placeholder for unavailable DeepGEMM backend.

Source code in vllm/utils/deep_gemm.py

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."
    )

_paged_mqa_logits_schedule_slots(num_sms, next_n)

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.

Source code in vllm/utils/deep_gemm.py

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

bf16_mega_gate(x, weight, num_topk, *, scoring_func, routed_scaling_factor, ep_rank, bias=None, image_bias=None, image_token_mask=None, fix_routing_mask=None, unmapped_topk_idx=None)

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.

Source code in vllm/utils/deep_gemm.py

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

calc_diff(x, y)

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.

Source code in vllm/utils/deep_gemm.py

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

fp8_fp4_mqa_logits(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, max_seqlen_k)

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

Parameters:

• q

  (tuple[Tensor, Tensor | None]) –

  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[Tensor, Tensor]) –

  Tuple (k_values, k_scales) of shape [N, D] and [N]; k_scales is int32 packed UE8M0 with MX scales, fp32 otherwise.

• weights

  (Tensor) –

  Tensor of shape [M, H], bf16 on SM100 and fp32 elsewhere.

• cu_seqlen_ks

  (Tensor) –

  Start indices (inclusive) for valid K per query position, shape [M], dtype int32.

• cu_seqlen_ke

  (Tensor) –

  End indices (exclusive) for valid K per query position, shape [M], dtype int32.

• max_seqlen_k

  (int) –

  Upper bound on cu_seqlen_ke - cu_seqlen_ks.

Returns:

• Tensor –

  Compressed logits of shape [M, max_seqlen_k]: row i holds K

• Tensor –

  positions [cu_seqlen_ks[i], cu_seqlen_ke[i]) starting at column

• Tensor –

  zero, and the tail is not cleaned. bf16 on SM100, fp32 elsewhere.

Source code in vllm/utils/deep_gemm.py

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

fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_tables, schedule_metadata, max_model_len, indices=None)

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

Parameters:

• q

  (tuple[Tensor, Tensor | None]) –

  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

  (Tensor) –

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

  Tensor of shape [B * next_n, H], bf16 on SM100 and fp32 elsewhere.

• context_lens

  (Tensor) –

  Tensor of shape [B], dtype int32; effective context length for each batch element.

• block_tables

  (Tensor) –

  Tensor of shape [B, max_blocks], dtype int32; maps logical block indices to physical blocks in the paged cache.

• schedule_metadata

  (Tensor) –

  Returned by get_paged_mqa_logits_metadata; used to distribute work across SMs.

• max_model_len

  (int) –

  Maximum sequence length used to size the logits output.

• indices

  (Tensor | None, default: None ) –

  Optional request index for each varlen row.

Returns:

• Tensor –

  Logits tensor of shape [B * next_n, max_model_len]; columns past each

• Tensor –

  row's context length are not cleaned. bf16 on SM100, fp32 elsewhere.

Source code in vllm/utils/deep_gemm.py

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

fp8_fp4_paged_sparse_mqa_logits(q, kv_cache, weights, metadata, num_max_sparse_blocks, sparse_block_kv)

Compute MQA logits only at the candidate blocks (paged decode path).

Parameters:

• q

  (tuple[Tensor, Tensor]) –

  (q_values, q_scale); q_values is [num_q_tokens, 1, H, D].

• kv_cache

  (Tensor) –

  Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, page stride 512B-aligned.

• weights

  (Tensor) –

  [num_q_tokens, H] torch.bfloat16.

• metadata

  (Tensor) –

  From get_paged_sparse_mqa_logits_metadata.

• num_max_sparse_blocks

  (int) –

  Candidate blocks per row.

• sparse_block_kv

  (int) –

  Tokens per sparse block, 8 or 16.

Returns:

• Tensor –

  bf16 logits of shape

• Tensor –

  [num_q_tokens, num_max_sparse_blocks * sparse_block_kv].

Source code in vllm/utils/deep_gemm.py

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 512B-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,
    )

fp8_fp4_sparse_mqa_logits(q, kv, weights, metadata, num_max_sparse_blocks, sparse_block_kv)

Compute MQA logits only at the candidate blocks (prefill path).

Parameters:

• q

  (tuple[Tensor, Tensor]) –

  (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

  (tuple[Tensor, Tensor]) –

  (kv_values, kv_scale) — the packed KV workspace with UE8M0 int32 scales.

• weights

  (Tensor) –

  [M, H] torch.bfloat16 (the Q scale is NOT folded in).

• metadata

  (Tensor) –

  From get_sparse_mqa_logits_metadata.

• num_max_sparse_blocks

  (int) –

  Width of each sparse-index row.

• sparse_block_kv

  (int) –

  Tokens per sparse block, 8 or 16.

Returns:

• Tensor –

  bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv];

• Tensor –

  column j * sparse_block_kv + o scores the token at

• Tensor –

  sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o.

Source code in vllm/utils/deep_gemm.py

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

get_col_major_tma_aligned_tensor(x)

Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor.

Source code in vllm/utils/deep_gemm.py

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)

get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf, ks_tensor, ks, gran_k)

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

Source code in vllm/utils/deep_gemm.py

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
    )

get_mn_major_tma_aligned_packed_ue8m0_tensor(x)

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.

Source code in vllm/utils/deep_gemm.py

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)

get_paged_mqa_logits_metadata(context_lens, block_size, num_sms, indices=None)

Build scheduling metadata for paged MQA logits.

Parameters:

• context_lens

  (Tensor) –

  Tensor of shape [B, next_n], dtype int32; effective context length per Q row.

• block_size

  (int) –

  KV-cache block size in tokens (e.g., 64).

• num_sms

  (int) –

  Number of SMs available. 132 for Hopper

• indices

  (Tensor | None, default: None ) –

  Optional request index for each varlen row.

Returns:

• Tensor –

  Tensor of shape [slots + 1, 2] consumed by fp8_fp4_paged_mqa_logits

• Tensor –

  to schedule work across SMs.

Source code in vllm/utils/deep_gemm.py

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
    )

get_paged_mqa_page_sizes()

KV page sizes (in cache entries) the paged-MQA logits kernels take (only 64 on SM120); larger storage blocks are split into one of these.

Source code in vllm/utils/deep_gemm.py

def get_paged_mqa_page_sizes() -> tuple[int, ...]:
    """KV page sizes (in cache entries) the paged-MQA logits kernels take (only
    64 on SM120); larger storage blocks are split into one of these."""
    return (64,) if current_platform.is_device_capability_family(120) else (32, 64)

get_paged_sparse_mqa_logits_metadata(context_lens, block_table, indices, page_kv, sparse_kv_block_indices, qk_dtype, sparse_block_kv)

Build scheduling metadata for fp8_fp4_paged_sparse_mqa_logits.

Parameters:

• context_lens

  (Tensor) –

  Per-row context length, shape [num_q_tokens], int32.

• block_table

  (Tensor) –

  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

  (Tensor) –

  Request index for each query row, shape [num_q_tokens].

• page_kv

  (int) –

  Tokens per KV page; must be a multiple of sparse_block_kv.

• sparse_kv_block_indices

  (Tensor) –

  Per-row logical (context-relative) candidate block ids, [num_q_tokens, num_max_sparse_blocks], int32.

• qk_dtype

  (dtype) –

  dtype of the packed Q values.

• sparse_block_kv

  (int) –

  Tokens per sparse block, 8 or 16.

Source code in vllm/utils/deep_gemm.py

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

get_sparse_mqa_logits_metadata(cu_seqlen_ks, cu_seqlen_ke, num_kv_tokens, sparse_kv_block_indices, qk_dtype, sparse_block_kv)

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.

Parameters:

• cu_seqlen_ks

  (Tensor) –

  Per-row K range start bounds in the packed KV workspace, shape [num_q_tokens], dtype int32.

• cu_seqlen_ke

  (Tensor) –

  Per-row K range end bounds in the packed KV workspace, shape [num_q_tokens], dtype int32.

• num_kv_tokens

  (int) –

  Total KV tokens in the packed workspace.

• sparse_kv_block_indices

  (Tensor) –

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

  dtype of the packed Q values (torch.int8 for MXFP4, torch.float8_e4m3fn for FP8).

• sparse_block_kv

  (int) –

  Tokens per sparse block, 8 or 16.

Source code in vllm/utils/deep_gemm.py

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

get_theoretical_mk_alignment_for_contiguous_layout(expected_m=None, num_groups=None)

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

Source code in vllm/utils/deep_gemm.py

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)

has_deep_gemm_sparse_mqa()

Whether the installed DeepGEMM provides the sparse-indexer kernels (fp8_fp4_(paged_)sparse_mqa_logits, added in DeepGEMM 2.8, SM100-only).

Source code in vllm/utils/deep_gemm.py

def has_deep_gemm_sparse_mqa() -> bool:
    """Whether the installed DeepGEMM provides the sparse-indexer kernels
    (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only)."""
    _lazy_init()
    return (
        _fp8_fp4_sparse_mqa_logits_impl is not None
        and _fp8_fp4_paged_sparse_mqa_logits_impl is not None
        and _get_sparse_mqa_logits_metadata_impl is not None
        and _get_paged_sparse_mqa_logits_metadata_impl is not None
    )

is_deep_gemm_e8m0_used() cached

Return True if vLLM is configured to use DeepGEMM " "E8M0 scale on a Hopper or Blackwell-class GPU.

Source code in vllm/utils/deep_gemm.py

@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

is_deep_gemm_supported() cached

Return True if DeepGEMM is supported on the current platform. Currently, only Hopper and Blackwell GPUs are supported.

Source code in vllm/utils/deep_gemm.py

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

mega_mhc(*args, **kwargs)

Run DeepGEMM Mega mHC with caller-owned output tensors.

Source code in vllm/utils/deep_gemm.py

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)

mk_alignment_scope(value)

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

Source code in vllm/utils/deep_gemm.py

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

mqa_logits_mx_inputs(q, q_scale, weights)

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.

Source code in vllm/utils/deep_gemm.py

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)

mqa_logits_use_mx_scales() cached

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.

Source code in vllm/utils/deep_gemm.py

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

native_next_n_supported(next_n)

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.

Source code in vllm/utils/deep_gemm.py

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

pack_ue8m0_to_int(x)

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

Source code in vllm/utils/deep_gemm.py

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)

set_mk_alignment_for_contiguous_layout(value)

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

Source code in vllm/utils/deep_gemm.py

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)

should_auto_disable_deep_gemm(model_type)

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.

Source code in vllm/utils/deep_gemm.py

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

tf32_hc_prenorm_gemm(x, fn, out, sqrsum, num_split)

Perform the following computation

out = x.float() @ fn.T sqrsum = x.float().square().sum(-1)

See the caller function for shape requirement

Source code in vllm/utils/deep_gemm.py

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