Skip to content

vllm.v1.attention.backends.fa_utils

Classes:

Functions:

FlashAttentionCuTeDSLCompileSpec dataclass

High-level FA4 compile-only request used by vLLM warmup.

This is not the CuTeDSL cache key. FA4 owns the selector that maps these serving inputs to the actual compile-static fields: tile sizes, q_stage, Split-KV, scheduler choice, layout-presence booleans, dtype/head dims, arch, and related fields.

Source code in vllm/v1/attention/backends/fa_utils.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
@dataclass(frozen=True)
class FlashAttentionCuTeDSLCompileSpec:
    """High-level FA4 compile-only request used by vLLM warmup.

    This is not the CuTeDSL cache key. FA4 owns the selector that maps these
    serving inputs to the actual compile-static fields: tile sizes, q_stage,
    Split-KV, scheduler choice, layout-presence booleans, dtype/head dims,
    arch, and related fields.
    """

    q_shape: tuple[int, ...]
    k_shape: tuple[int, ...]
    v_shape: tuple[int, ...]
    q_dtype: torch.dtype
    max_seqlen_q: int
    max_seqlen_k: int
    softmax_scale: float
    causal: bool
    fa_version: int
    v_stride: tuple[int, ...] | None = None
    cu_seqlens_q_shape: tuple[int, ...] | None = None
    cu_seqlens_k_shape: tuple[int, ...] | None = None
    window_size: tuple[int, int] | None = None
    return_softmax_lse: bool = False
    num_splits: int = 0

    def compile(self) -> None:
        assert compile_flash_attn_varlen_func_from_specs is not None
        window_size = list(self.window_size) if self.window_size is not None else None
        compile_flash_attn_varlen_func_from_specs(
            q_shape=self.q_shape,
            k_shape=self.k_shape,
            v_shape=self.v_shape,
            q_dtype=self.q_dtype,
            v_stride=self.v_stride,
            cu_seqlens_q_shape=self.cu_seqlens_q_shape,
            cu_seqlens_k_shape=self.cu_seqlens_k_shape,
            max_seqlen_q=self.max_seqlen_q,
            max_seqlen_k=self.max_seqlen_k,
            softmax_scale=self.softmax_scale,
            causal=self.causal,
            window_size=window_size,
            return_softmax_lse=self.return_softmax_lse,
            fa_version=self.fa_version,
            num_splits=self.num_splits,
        )

    def request_key(self) -> tuple[object, ...]:
        return (
            self.q_shape,
            self.k_shape,
            self.v_shape,
            self.q_dtype,
            self.max_seqlen_q,
            self.max_seqlen_k,
            self.softmax_scale,
            self.causal,
            self.fa_version,
            self.v_stride,
            self.cu_seqlens_q_shape,
            self.cu_seqlens_k_shape,
            self.window_size,
            self.return_softmax_lse,
            self.num_splits,
        )

is_flash_attn_varlen_func_available()

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: upstream flash_attn.flash_attn_varlen_func (if available)

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

    True if a working flash_attn_varlen_func implementation is available.

Source code in vllm/v1/attention/backends/fa_utils.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
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: upstream flash_attn.flash_attn_varlen_func (if available)

    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