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 5, 11:00 UTC
[Attention] relax the head dim 512 and paged kv for sm90+FA4 (#38835)
Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Lain committed
Apr 8, 2026 at 11:23 UTC
e24e0a43a4f6138c8cf611f880c3609e3edf7624
4 files changed
+29
-10
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -39,7 +39,7 @@ else()
39
FetchContent_Declare(
40
vllm-flash-attn
41
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42
- GIT_TAG c0ec424fd8a546d0cbbf4bf050bbcfe837c55afb
42
+ GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
43
GIT_PROGRESS TRUE
44
# Don't share the vllm-flash-attn build between build types
45
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/v1/attention/backends/fa_utils.py
+11
@@ -154,6 +154,17 @@ def get_flash_attn_version(
154
return None
155
156
157
+def is_fa_version_supported(fa_version: int) -> bool:
158
+ try:
159
+ from vllm.vllm_flash_attn.flash_attn_interface import (
160
+ is_fa_version_supported as _is_fa_version_supported,
161
+ )
162
+
163
+ return _is_fa_version_supported(fa_version)
164
+ except ImportError:
165
+ return False
166
+
167
+
168
def flash_attn_supports_fp8() -> bool:
169
return (
170
get_flash_attn_version() == 3
vllm/v1/attention/backends/flash_attn.py
+17
-2
@@ -10,6 +10,7 @@ import numpy as np
10
import torch
11
12
from vllm.model_executor.layers.attention import Attention
13
+from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import is_quantized_kv_cache
15
from vllm.v1.attention.backend import (
16
AttentionBackend,
@@ -20,6 +21,7 @@ from vllm.v1.attention.backend import (
21
from vllm.v1.attention.backends.fa_utils import (
22
flash_attn_supports_fp8,
23
get_flash_attn_version,
24
+ is_fa_version_supported,
25
is_flash_attn_varlen_func_available,
26
)
27
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
@@ -45,7 +47,6 @@ from vllm.config import (
47
from vllm.config.cache import CacheDType
48
from vllm.distributed.parallel_state import get_dcp_group
49
from vllm.logger import init_logger
48
-from vllm.platforms import current_platform
50
from vllm.platforms.interface import DeviceCapability
51
from vllm.utils.math_utils import cdiv, round_up
52
from vllm.v1.attention.backend import (
@@ -170,7 +171,13 @@ class FlashAttentionBackend(AttentionBackend):
171
172
@classmethod
173
def supports_head_size(cls, head_size: int) -> bool:
173
- return head_size % 8 == 0 and head_size <= 256
174
+ if head_size % 8 != 0:
175
+ return False
176
+ if head_size <= 256:
177
+ return True
178
+ if is_fa_version_supported(4):
179
+ return head_size <= 512
180
+ return False
181
182
@classmethod
183
def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
@@ -618,6 +625,14 @@ class FlashAttentionImpl(AttentionImpl):
625
requires_alibi=alibi_slopes is not None,
626
head_size=head_size,
627
)
628
+ # head_size > 256 requires FA4 on SM90+; force upgrade from FA3
629
+ if (
630
+ head_size > 256
631
+ and self.vllm_flash_attn_version == 3
632
+ and current_platform.is_cuda()
633
+ and current_platform.is_device_capability_family(90)
634
+ ):
635
+ self.vllm_flash_attn_version = 4
636
logger.info_once(
637
"Using FlashAttention version %s",
638
self.vllm_flash_attn_version,
vllm/vllm_flash_attn/flash_attn_interface.py
-7
@@ -366,14 +366,7 @@ def flash_attn_varlen_func(
366
)
367
elif fa_version == 4:
368
assert alibi_slopes is None, "Alibi is not supported in FA4"
369
- # FA4 on SM90 doesn't support paged KV; SM100+ does
370
- from vllm.platforms import current_platform
369
372
- if block_table is not None and current_platform.is_device_capability_family(90):
373
- raise NotImplementedError(
374
- "FA4 with paged KV is not supported on SM90 (Hopper). "
375
- "Use FA3 or upgrade to Blackwell (SM100+)."
376
- )
370
from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
371
372
out, softmax_lse = _flash_attn_fwd(