Skip to content

vllm_omni.diffusion.attention.backends.utils.fa

FlashAttnFn module-attribute 

FlashAttnFn = Callable[..., Any]

HAS_FLASH_ATTN module-attribute 

HAS_FLASH_ATTN = (
    flash_attn_func is not None
    or flash_attn_varlen_func is not None
)

IS_AITER module-attribute 

IS_AITER = False

IS_FLASH_ATTN_4 module-attribute 

IS_FLASH_ATTN_4 = True

flash_attn_func module-attribute 

flash_attn_func: FlashAttnFn | None = _fa_func

flash_attn_varlen_func module-attribute 

flash_attn_varlen_func: FlashAttnFn | None = _fa_varlen

logger module-attribute 

logger = init_logger(__name__)

is_flash_attn_4_available cached 

is_flash_attn_4_available() -> bool

Return whether CuTe FlashAttention-4 (flash_attn.cute) can be imported.

FA2/FA3 wheels are often importable on Blackwell but only ship Hopper kernels. Those must not be treated as a runnable FLASH_ATTN implementation here.

is_flash_attn_installed cached 

is_flash_attn_installed() -> bool

Return whether a Flash Attention backend package is importable.

Shared by CUDA/ROCm/MUSA platforms.

resolve_vllm_flash_attn_version cached 

resolve_vllm_flash_attn_version(
    requested: str | int | None = None,
) -> int

Resolve an available vLLM-bundled FA2/FA3/FA4 implementation.

validate_fa4_head_dims 

validate_fa4_head_dims(
    head_dim: int, head_dim_v: int, alignment: int
) -> bool

Use the selected FA4 implementation's constraints during path resolution.

Return False when its validator is unavailable; propagate kernel validation errors for the backend to report. Keep the private FA4 API dependency here, rather than duplicating architecture/version-specific dimension rules. This is a planning-time call, never part of compiled tensor execution.

vllm_flash_attn_varlen_with_lse 

vllm_flash_attn_varlen_with_lse(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    *,
    cu_seqlens_q: Tensor,
    cu_seqlens_k: Tensor,
    max_seqlen_q: int,
    max_seqlen_k: int,
    softmax_scale: float | None = None,
    softcap: float = 0.0,
    causal: bool = False,
    deterministic: bool = False,
    fa_version: int | None = None,
) -> tuple[Tensor, Tensor]

Run packed vLLM FlashAttention and retain LSE for post-processing.