vllm.v1.attention.backends.fa_utils ¶

def is_flash_attn_varlen_func_available() -> bool:
    """Check if flash_attn_varlen_func is available.

    This function determines whether the flash_attn_varlen_func imported at module
    level is a working implementation or a stub.

    Platform-specific sources:
    - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func
    - XPU: xpu_ops.flash_attn_varlen_func
    - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) or
    upstream flash_attn.flash_attn_varlen_func

    Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py)
    which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is
    handled separately via _aiter_ops.is_aiter_found_and_supported().

    Returns:
        bool: True if a working flash_attn_varlen_func implementation is available.
    """
    if current_platform.is_cuda() or current_platform.is_xpu():
        # CUDA and XPU always have flash_attn_varlen_func available
        return True

    if current_platform.is_rocm():
        # Use the flag set during module import to check if
        # upstream flash-attn was successfully imported
        return _ROCM_FLASH_ATTN_AVAILABLE

    return False

def uses_fa4_hd256_kernel(
    head_size: int | None, head_size_v: int | None = None
) -> bool:
    """Return whether FA4 uses its dedicated hd256 kernel."""
    if head_size != 256:
        return False
    if head_size_v is not None and head_size_v != 256:
        return False
    capability = current_platform.get_device_capability()
    return capability is not None and capability.major in (10, 11)