vllm.utils.flashinfer ¶

def _flashinfer_concat_mla_k(
    k: torch.Tensor,
    k_nope: torch.Tensor,
    k_pe: torch.Tensor,
) -> None:
    """Custom op wrapper for flashinfer's concat_mla_k.

    This is an in-place operation that concatenates k_nope and k_pe into k.

    The kernel is optimized for DeepSeek V3 dimensions:
    - num_heads=128
    - nope_dim=128
    - rope_dim=64

    Key optimizations:
    - Warp-based processing with software pipelining
    - Vectorized memory access (int2 for nope, int for rope)
    - L2 prefetching for next row while processing current
    - Register reuse for rope values across all heads

    Args:
        k: Output tensor, shape [num_tokens, num_heads, nope_dim + rope_dim].
            Modified in-place.
        k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim].
        k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim].
              This is broadcast to all heads.
    """
    from flashinfer.concat_ops import concat_mla_k

    concat_mla_k(k, k_nope, k_pe)

def _get_submodule(module_name: str) -> Any | None:
    """Safely import a submodule and return it, or None if not available."""
    try:
        return importlib.import_module(module_name)
    except (ImportError, ModuleNotFoundError):
        return None

def _lazy_import_wrapper(
    module_name: str, attr_name: str, fallback_fn: Callable[..., Any] = _missing
):
    """Create a lazy import wrapper for a specific function."""

    @functools.cache
    def _get_impl():
        if not has_flashinfer():
            return None
        mod = _get_submodule(module_name)
        return getattr(mod, attr_name, None) if mod else None

    def wrapper(*args, **kwargs):
        impl = _get_impl()
        if impl is None:
            return fallback_fn(*args, **kwargs)
        return impl(*args, **kwargs)

    return wrapper

def _missing(*_: Any, **__: Any) -> NoReturn:
    """Placeholder for unavailable FlashInfer backend."""
    raise RuntimeError(
        "FlashInfer backend is not available. Please install the package "
        "to enable FlashInfer kernels: "
        "https://github.com/flashinfer-ai/flashinfer"
    )

def can_use_trtllm_attention(
    num_qo_heads: int, num_kv_heads: int, is_prefill: bool = False
) -> bool:
    """Check if the current configuration supports TRTLLM attention."""
    if force_use_trtllm_attention() is False:
        return False
    return supports_trtllm_attention(is_prefill=is_prefill) and (
        num_qo_heads % num_kv_heads == 0
    )

def flashinfer_mm_mxfp8(
    a: torch.Tensor,
    b: torch.Tensor,
    block_scale_a: torch.Tensor,
    block_scale_b: torch.Tensor,
    out_dtype: torch.dtype,
    backend: str = "cutlass",
) -> torch.Tensor:
    """MXFP8 MM helper - mirrors flashinfer_scaled_fp4_mm API.

    Takes non-transposed weights and handles transpose internally.

    CRITICAL: mm_mxfp8 CUTLASS kernel requires SWIZZLED 1D scales for optimal
    performance and accuracy. Both input and weight scales should be in
    swizzled format from FlashInfer's mxfp8_quantize(is_sf_swizzled_layout=True).
    """
    # a shape [M, K]
    # b shape [K, N]
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[1]  # K dimension must match

    if block_scale_b.ndim != 1:
        raise ValueError(
            "mm_mxfp8 expects 1D swizzled weight scales for CUTLASS; "
            f"got shape={tuple(block_scale_b.shape)}"
        )

    # Output tensor [M, N]
    return mm_mxfp8(
        a,
        b.t(),  # Transpose weight: [N, K] -> [K, N]
        block_scale_a,
        block_scale_b,
        out_dtype,
        backend=backend,
    )

def force_use_trtllm_attention() -> bool | None:
    """
    This function should only be called during initialization stage when vllm config
    is set.
    Return `None` if --attention-config.use_trtllm_attention is not set,
    return `True` if TRTLLM attention is forced to be used,
    return `False` if TRTLLM attention is forced to be not used.
    """
    from vllm.config import get_current_vllm_config

    vllm_config = get_current_vllm_config()
    return vllm_config.attention_config.use_trtllm_attention

@functools.cache
def has_flashinfer() -> bool:
    """Return `True` if flashinfer-python package is available."""
    # Use find_spec to check if the module exists without importing it
    # This avoids potential CUDA initialization side effects
    if importlib.util.find_spec("flashinfer") is None:
        logger.debug_once("FlashInfer unavailable since package was not found")
        return False
    # When not using flashinfer cubin,
    # Also check if nvcc is available since it's required to JIT compile flashinfer
    if not has_flashinfer_cubin() and shutil.which("nvcc") is None:
        logger.debug_once(
            "FlashInfer unavailable since nvcc was not found "
            "and not using pre-downloaded cubins"
        )
        return False
    return True

@functools.cache
def has_flashinfer_b12x_gemm() -> bool:
    """Return True if FlashInfer b12x FP4 GEMM backend is available (SM120+)."""
    if not has_flashinfer_cutedsl():
        return False
    mod = _get_submodule("flashinfer.gemm")
    if mod is None:
        return False
    # FlashInfer 0.6.11 renamed Sm120BlockScaledDenseGemmKernel ->
    # Sm120B12xBlockScaledDenseGemmKernel (commit 223f2a49). Accept either.
    return hasattr(mod, "Sm120B12xBlockScaledDenseGemmKernel") or hasattr(
        mod, "Sm120BlockScaledDenseGemmKernel"
    )

@functools.cache
def has_flashinfer_b12x_moe() -> bool:
    """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available."""
    if not has_flashinfer_moe():
        return False

    required_functions = [
        ("flashinfer.fused_moe", "b12x_fused_moe"),
        ("flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"),
    ]

    for module_name, attr_name in required_functions:
        mod = _get_submodule(module_name)
        if not mod or not hasattr(mod, attr_name):
            return False
    return True

@functools.cache
def has_flashinfer_bf16_fp4() -> bool:
    """Return ``True`` if FlashInfer's CuTe-DSL W4A16 GEMM is available."""
    if not has_flashinfer_cutedsl():
        return False
    mod = _get_submodule("flashinfer.gemm")
    return mod is not None and all(
        hasattr(mod, name) for name in ("mm_bf16_fp4", "prepare_bf16_fp4_weights")
    )

@functools.cache
def has_flashinfer_bf16_gemm() -> bool:
    """Return whether FlashInfer exposes the BF16 dense GEMM API."""
    if not has_flashinfer():
        return False
    mod = _get_submodule("flashinfer")
    return mod is not None and callable(getattr(mod, "mm_bf16", None))

@functools.cache
def has_flashinfer_comm() -> bool:
    """Return `True` if FlashInfer comm module is available."""
    return has_flashinfer() and importlib.util.find_spec("flashinfer.comm") is not None

@functools.cache
def has_flashinfer_cubin() -> bool:
    """Return `True` if flashinfer-cubin package is available."""
    if envs.VLLM_HAS_FLASHINFER_CUBIN:
        return True
    if importlib.util.find_spec("flashinfer_cubin") is not None:
        return True
    logger.debug_once("flashinfer-cubin package was not found")
    return False

@functools.cache
def has_flashinfer_cutedsl() -> bool:
    """Return ``True`` if FlashInfer cutedsl module is available."""
    return (
        has_flashinfer() and importlib.util.find_spec("flashinfer.cute_dsl") is not None
    )

@functools.cache
def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool:
    """Return ``True`` if FlashInfer CUTLASS fused MoE is available."""
    if not has_flashinfer_cutedsl():
        return False

    # Check if all required functions are available
    required_functions = [
        ("flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"),
        ("flashinfer", "scaled_fp4_grouped_quantize"),
        ("flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"),
    ]

    for module_name, attr_name in required_functions:
        mod = _get_submodule(module_name)
        if not mod or not hasattr(mod, attr_name):
            return False
    return True

@functools.cache
def has_flashinfer_cutedsl_moe_nvfp4() -> bool:
    """Return ``True`` if FlashInfer cute_dsl_fused_moe_nvfp4 is available."""
    if not has_flashinfer_cutedsl():
        return False
    mod = _get_submodule("flashinfer")
    return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4")

@functools.cache
def has_flashinfer_cutlass_fused_moe() -> bool:
    """Return `True` if FlashInfer CUTLASS fused MoE is available."""
    if not has_flashinfer_moe():
        return False

    # Check if all required functions are available
    required_functions = [
        ("flashinfer.fused_moe", "cutlass_fused_moe"),
        ("flashinfer", "fp4_quantize"),
        ("flashinfer", "nvfp4_block_scale_interleave"),
        ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),
    ]

    for module_name, attr_name in required_functions:
        mod = _get_submodule(module_name)
        if not mod or not hasattr(mod, attr_name):
            return False
    return True

@functools.cache
def has_flashinfer_fp8_blockscale_gemm() -> bool:
    """Return `True` if FlashInfer block-scale FP8 GEMM is available."""
    return (
        has_flashinfer()
        and current_platform.is_device_capability(90)
        and hasattr(_get_submodule("flashinfer.gemm"), "fp8_blockscale_gemm_sm90")
    )

@functools.cache
def has_flashinfer_moe() -> bool:
    """Return `True` if FlashInfer MoE module is available."""
    return (
        has_flashinfer()
        and importlib.util.find_spec("flashinfer.fused_moe") is not None
    )

@functools.cache
def has_flashinfer_nvlink_one_sided() -> bool:
    """Return `True` if FlashInfer trtllm_moe_alltoall module is available."""
    if not has_flashinfer_comm():
        return False
    return importlib.util.find_spec("flashinfer.comm.trtllm_moe_alltoall") is not None

@functools.cache
def has_flashinfer_nvlink_two_sided() -> bool:
    """Return `True` if FlashInfer mnnvl all2all is available."""
    if not has_flashinfer_comm():
        return False

    # Check if all required functions are available
    required_functions = [
        ("flashinfer.comm", "Mapping"),
        ("flashinfer.comm.mnnvl", "MnnvlMemory"),
        ("flashinfer.comm.trtllm_alltoall", "MnnvlMoe"),
        ("flashinfer.comm.trtllm_alltoall", "MoEAlltoallInfo"),
    ]

    for module_name, attr_name in required_functions:
        mod = _get_submodule(module_name)
        if not mod or not hasattr(mod, attr_name):
            return False
    return True

@functools.cache
def has_flashinfer_sm90_nope_mla() -> bool:
    """FlashInfer SM90 NoPE MLA (FP8 KV with in-kernel dequant, kpe=0).

    Feature-detected via the ``ckv_scale_arr`` run() kwarg introduced with
    the SM90 NoPE support (FlashInfer >= 0.6.18), so dev builds carry the
    gate without a version parse.
    """
    if not has_flashinfer():
        return False
    try:
        import inspect

        from flashinfer.mla import BatchMLAPagedAttentionWrapper
    except ImportError:
        return False
    try:
        params = inspect.signature(BatchMLAPagedAttentionWrapper.run).parameters
    except (TypeError, ValueError):
        return False
    return (
        "ckv_scale_arr" in params
        and params["ckv_scale_arr"].kind is inspect.Parameter.KEYWORD_ONLY
    )

@functools.cache
def has_flashinfer_sparse_mla_sm120() -> bool:
    """Return ``True`` if FlashInfer sparse MLA decode support is available."""
    if not has_flashinfer():
        return False
    try:
        from flashinfer.autotuner import autotune
        from flashinfer.decode import (
            trtllm_batch_decode_sparse_mla_dsv4,
            trtllm_batch_decode_with_kv_cache_mla,
        )
    except ImportError:
        return False
    return (
        callable(trtllm_batch_decode_sparse_mla_dsv4)
        and callable(trtllm_batch_decode_with_kv_cache_mla)
        and callable(autotune)
    )

@functools.cache
def has_flashinfer_sparse_mla_sm120_config(num_q_heads: int, top_k: int) -> bool:
    """Return whether FlashInfer ships an SM120 DSV4 decode specialization.

    The public sparse MLA API predates some DSV4 shapes, so checking only that
    the callable exists can select a package that later aborts or rejects a
    valid vLLM configuration. Inspect FlashInfer's dispatch table until it
    exposes a public capability query.
    """
    if not has_flashinfer_sparse_mla_sm120():
        return False
    mod = _get_submodule("flashinfer.mla._sparse_mla_sm120")
    dispatch = getattr(mod, "_DECODE_DSV4_DISPATCH", None) if mod else None
    return dispatch is not None and (int(num_q_heads), int(top_k)) in dispatch

@functools.cache
def has_flashinfer_trtllm_fused_moe() -> bool:
    """Return `True` if FlashInfer TRTLLM fused MoE is available."""
    if not has_flashinfer_moe():
        return False
    required_functions = [
        ("flashinfer.fused_moe", "trtllm_fp8_block_scale_moe"),
        ("flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe"),
        ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),
        ("flashinfer.fused_moe", "trtllm_mxint4_block_scale_moe"),
        ("flashinfer.fused_moe", "trtllm_bf16_moe"),
    ]
    for module_name, attr_name in required_functions:
        mod = _get_submodule(module_name)
        if not mod or not hasattr(mod, attr_name):
            return False
    return True

@functools.cache
def has_nvidia_artifactory() -> bool:
    """Return `True` if NVIDIA's artifactory is accessible.

    This checks connectivity to the kernel inference library artifactory
    which is required for downloading certain cubin kernels like TRTLLM FHMA.
    """
    # If we have pre-downloaded cubins, we can assume the cubins are available.
    if has_flashinfer_cubin():
        return True

    try:
        # Use a short timeout to avoid blocking for too long
        response = requests.get(FLASHINFER_CUBINS_REPOSITORY, timeout=5)
        accessible = response.status_code == 200
        if accessible:
            logger.debug_once("NVIDIA artifactory is accessible")
        else:
            logger.warning_once(
                "NVIDIA artifactory returned failed status code: %d",
                response.status_code,
            )
        return accessible
    except Exception as e:
        logger.warning_once("Failed to connect to NVIDIA artifactory: %s", e)
        return False

@functools.cache
def is_flashinfer_bf16_gemm_supported(
    backend: str,
    compute_capability: int | None = None,
) -> bool:
    """Return whether an exact FlashInfer BF16 backend is available."""
    if not current_platform.is_cuda() or not has_flashinfer_bf16_gemm():
        return False

    mod = _get_submodule("flashinfer")
    mm_bf16 = getattr(mod, "mm_bf16", None) if mod is not None else None
    backend_supported = getattr(mm_bf16, "is_backend_supported", None)
    if not callable(backend_supported):
        return False

    if compute_capability is None:
        device_capability = current_platform.get_device_capability()
        if device_capability is None:
            return False
        compute_capability = device_capability.to_int()

    try:
        return bool(backend_supported(backend, compute_capability))
    except (AttributeError, ImportError, RuntimeError, TypeError, ValueError):
        return False

@functools.cache
def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
    """Check if FP8 ViT attention is supported on this platform.

    Requires Blackwell (SM 100) or newer, the FlashInfer cuDNN backend,
    and cuDNN >= 9.17.1.

    cuDNN's FP8 SDPA forward path with bf16/fp16 output (used by
    ``MMEncoderAttention._forward_flashinfer``) gates internally on
    ``prop.major >= 10``; on Hopper it raises a misleading
    ``cudnnGraphNotSupportedError: ... cuDNN version 9.13.0 and newer``
    even when the installed cuDNN is new enough. See PR #38065 for the
    original Blackwell-only design intent.
    """
    from vllm.v1.attention.backends.registry import AttentionBackendEnum

    # cuDNN SDPA FP8 with bf16/fp16 output requires Blackwell (SM 100) or newer.
    if not current_platform.has_device_capability(100):
        return False

    try:
        supported = current_platform.get_supported_vit_attn_backends()
        if AttentionBackendEnum.FLASHINFER not in supported:
            return False
    except (ImportError, AttributeError):
        return False

    try:
        import torch.backends.cudnn as cudnn

        if cudnn.is_available() and cudnn.version() < _MIN_CUDNN_FP8:
            return False
    except (ImportError, AttributeError):
        pass

    return True

@functools.cache
def is_flashinfer_cutedsl_bf16_gemm_supported() -> bool:
    """Return whether the CuTeDSL BF16 dense GEMM backend is available."""
    if not is_flashinfer_bf16_gemm_supported("cute-dsl"):
        return False
    try:
        from flashinfer.cute_dsl.utils import is_cute_dsl_available
        from flashinfer.utils import is_sm100a_supported
    except (ImportError, ModuleNotFoundError):
        return False
    try:
        return is_cute_dsl_available() and is_sm100a_supported(torch.device("cuda"))
    except (RuntimeError, TypeError, ValueError):
        return False

@functools.cache
def is_flashinfer_fp8_blockscale_gemm_supported() -> bool:
    """Return `True` if FlashInfer block-scale FP8 GEMM is supported."""
    return (
        envs.VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER
        and has_flashinfer_fp8_blockscale_gemm()
    )

def pin_host_range_buf(length: int) -> None:
    """Pin flashinfer's cached host-side arange buffer covering `length`.

    flashinfer's decode plan() builds its host qo_indptr from the cached
    `_get_range_buf` entry and copies it to the GPU with non_blocking=True on
    every call; the cache entry is unpinned, so the copy silently blocks the
    host. Replace the entry with a pinned copy so the copies stay async.
    """
    if not PIN_MEMORY:
        return
    mod = _get_submodule("flashinfer.utils")
    if mod is None:
        return
    get_range_buf = getattr(mod, "_get_range_buf", None)
    cache_buf = getattr(mod, "_cache_buf", None)
    ceil_pow2 = getattr(mod, "_ceil_pow2", None)
    if get_range_buf is None or cache_buf is None or ceil_pow2 is None:
        return
    get_range_buf(length, "cpu")  # Ensure the cache entry exists.
    key = (f"range_{ceil_pow2(length)}", "cpu")
    buf = cache_buf[key]
    if not buf.is_pinned():
        cache_buf[key] = buf.pin_memory()

@functools.cache
def supports_trtllm_attention(is_prefill: bool = False) -> bool:
    """Return whether TRTLLM attention is available on the current platform
    for the given attention phase.

    SM90 (Hopper) and SM12x support the XQA decode kernel but not TRTLLM
    prefill. SM100+ supports TRTLLM for both phases. All others are unsupported.
    """
    # Batch-invariant mode disables TRTLLM attention
    if envs.VLLM_BATCH_INVARIANT:
        return False

    # Requires NVIDIA artifactory to be accessible to download cubins
    if not has_nvidia_artifactory():
        return False

    # SM90 and SM12x have XQA decode only.
    if current_platform.is_device_capability(
        90
    ) or current_platform.is_device_capability_family(120):
        return not is_prefill

    # SM100/SM103 has both prefill and decode TRTLLM kernels.
    return current_platform.is_device_capability_family(100)

def use_trtllm_attention(
    num_qo_heads: int,
    num_kv_heads: int,
    num_tokens: int,
    max_seq_len: int,
    dcp_world_size: int,
    kv_cache_dtype: str,
    q_dtype: torch.dtype,
    is_prefill: bool,
    # None means auto-detection, True means force on, False means force off
    force_use_trtllm: bool | None = None,
    has_sinks: bool = False,
    has_spec: bool = False,
) -> bool:
    """Return `True` if TRTLLM attention is used."""

    # CLI argument is set to 0 - respect it
    if force_use_trtllm is not None and not force_use_trtllm:
        return False

    # TRTLLM prefill attends only the DCP-local KV shard and has no
    # cross-rank LSE combine, so it cannot be used with DCP; fall back to
    # FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected
    # separately (all-gathered query heads + LSE combine in forward).
    if dcp_world_size > 1:
        logger.warning_once(
            "TRTLLM prefill does not support DCP, reverting to FlashInfer"
        )
        return False

    # The platform is not supported
    if not supports_trtllm_attention(is_prefill=is_prefill):
        if force_use_trtllm:
            logger.warning_once(
                "TRTLLM attention is not supported on this platform for %s, "
                "but --attention-config.use_trtllm_attention is set to 1",
                "prefill" if is_prefill else "decode",
            )
        return False

    # The combination of query and key heads is not supported
    if num_qo_heads % num_kv_heads != 0:
        if force_use_trtllm:
            logger.warning_once(
                "TRTLLM attention is not supported for this combination of "
                "query and key heads, but --attention-config.use_trtllm_attention is "
                "set to 1"
            )
        return False

    if has_spec and not is_prefill:
        # Speculative decoding requires TRTLLM attention for decodes
        logger.info_once("Using TRTLLM attention (enabled for speculative decoding).")
        return True

    # Must use TRTLLM attention if query is FP8 quantized
    if q_dtype == current_platform.fp8_dtype():
        logger.info_once("Using TRTLLM attention (query is quantized).")
        return True

    # If sinks are being used, we must use TRTLLM attention as it's
    # the only backend that supports them
    if has_sinks:
        logger.info_once("Using TRTLLM attention (required for attention sinks).")
        return True

    if force_use_trtllm is None:
        # CLI argument not set - use auto-detection
        if is_prefill:
            # Prefill auto-detection
            use_trtllm = kv_cache_dtype == "auto"
        elif (
            current_platform.is_device_capability(90)
            or current_platform.is_device_capability_family(120)
        ) and kv_cache_dtype.startswith("fp8"):
            # SM90/SM12x + FP8 KV cache: prefer the XQA decode kernel.
            use_trtllm = True
        else:
            # Decode auto-detection
            use_trtllm = num_tokens <= 256 and kv_cache_dtype == "auto"
        if use_trtllm:
            logger.warning_once(
                "Using TRTLLM %s attention (auto-detected).",
                "prefill" if is_prefill else "decode",
            )
        return use_trtllm

    # CLI argument is set to 1 - respect it
    logger.info_once(
        "Using TRTLLM attention (--attention-config.use_trtllm_attention is set to 1)"
    )
    return True