A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 00:45 UTC
[Attention] FlashAttention ViT support, make default backend (#28763)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Nov 18, 2025 at 23:06 UTC
4c23690f43e51eccf6ce5866ac47adcf39215e4d
5 files changed
+15
-46
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -38,7 +38,7 @@ else()
38
FetchContent_Declare(
39
vllm-flash-attn
40
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
41
- GIT_TAG 58e0626a692f09241182582659e3bf8f16472659
41
+ GIT_TAG 71bb26f6295449be880344b93b51791cc009237d
42
GIT_PROGRESS TRUE
43
# Don't share the vllm-flash-attn build between build types
44
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
tests/kernels/attention/test_flash_attn.py
+2
-2
@@ -13,14 +13,14 @@ from vllm.vllm_flash_attn import (
13
)
14
15
NUM_HEADS = [(4, 4), (8, 2)]
16
-HEAD_SIZES = [128, 256]
16
+HEAD_SIZES = [40, 72, 80, 128, 256]
17
BLOCK_SIZES = [16]
18
DTYPES = [torch.bfloat16]
19
QDTYPES = [None, torch.float8_e4m3fn]
20
# one value large enough to test overflow in index calculation.
21
# one value small enough to test the schema op check
22
NUM_BLOCKS = [32768, 2048]
23
-SOFT_CAPS = [None, 50.0]
23
+SOFT_CAPS = [None]
24
SLIDING_WINDOWS = [None, 256]
25
26
tests/kernels/attention/test_mha_attn.py
+1
-29
@@ -62,38 +62,10 @@ def test_mha_attn_platform(device: str):
62
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
63
64
# Test CUDA with head_size=72 (not divisible by 32)
65
- # - with upstream FA not available
66
- # - should use xformers
67
- with (
68
- patch("vllm.attention.layer.current_platform", CudaPlatform()),
69
- patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
70
- patch(
71
- "vllm.attention.layer.check_upstream_fa_availability",
72
- return_value=False,
73
- ),
74
- ):
75
- attn = MultiHeadAttention(16, 72, scale=1)
76
- assert attn.attn_backend == AttentionBackendEnum.XFORMERS
77
-
78
- # Test CUDA with head_size=72 (not divisible by 32)
79
- # - with upstream FA available
80
- # - should use upstream FA
65
+ # - should use vLLM's FlashAttention
66
with (
67
patch("vllm.attention.layer.current_platform", CudaPlatform()),
68
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
84
- patch(
85
- "vllm.attention.layer.check_upstream_fa_availability", return_value=True
86
- ),
87
- patch.dict(
88
- "sys.modules",
89
- {
90
- "flash_attn": type(
91
- "MockFlashAttn",
92
- (),
93
- {"flash_attn_varlen_func": lambda *args, **kwargs: None},
94
- )()
95
- },
96
- ),
69
):
70
attn = MultiHeadAttention(16, 72, scale=1)
71
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
vllm/platforms/cuda.py
+9
-12
@@ -267,24 +267,21 @@ class CudaPlatformBase(Platform):
267
) -> "AttentionBackendEnum":
268
from vllm.attention.backends.registry import AttentionBackendEnum
269
270
- # For Blackwell GPUs, force TORCH_SDPA for now.
271
- # See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
272
- if cls.has_device_capability(100):
273
- return AttentionBackendEnum.TORCH_SDPA
274
-
275
- if dtype not in (torch.float16, torch.bfloat16):
276
- return AttentionBackendEnum.XFORMERS
277
-
278
- if cls.has_device_capability(80):
270
+ # Try FlashAttention first
271
+ try:
272
backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
273
if backend_class.supports_head_size(
274
head_size
275
) and backend_class.supports_dtype(dtype):
276
return AttentionBackendEnum.FLASH_ATTN
284
- else:
285
- return AttentionBackendEnum.XFORMERS
277
+ except ImportError:
278
+ pass
279
+
280
+ if cls.has_device_capability(100):
281
+ # xFormers doesn't support Blackwell, fall back to SDPA
282
+ # See https://github.com/facebookresearch/xformers/issues/1317#issuecomment-3199392579 # noqa: E501
283
+ return AttentionBackendEnum.TORCH_SDPA
284
else:
287
- # Fallback for Volta/Turing GPUs or FA not supported
285
return AttentionBackendEnum.XFORMERS
286
287
@classmethod
vllm/v1/attention/backends/flash_attn.py
+2
-2
@@ -119,8 +119,8 @@ class FlashAttentionBackend(AttentionBackend):
119
raise ValueError(f"Unrecognized FP8 dtype: {kv_cache_dtype}")
120
121
@classmethod
122
- def get_supported_head_sizes(cls) -> list[int]:
123
- return [32, 64, 96, 128, 160, 192, 224, 256]
122
+ def supports_head_size(cls, head_size: int) -> bool:
123
+ return head_size % 8 == 0 and head_size <= 256
124
125
@classmethod
126
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool: