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
)

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.

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.