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 15, 18:30 UTC
[Bugfix] detect alibi and revert to FA2 (#15231)
Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>
Travis Johnson committed
Mar 20, 2025 at 20:20 UTC
0032903a5bb7c7c655f52f4efdfcc221947e9ca8
2 files changed
+11
-4
vllm/attention/backends/flash_attn.py
+2
-1
@@ -630,7 +630,8 @@ class FlashAttentionImpl(AttentionImpl):
630
self.sliding_window = ((sliding_window - 1,
631
0) if sliding_window is not None else (-1, -1))
632
self.kv_cache_dtype = kv_cache_dtype
633
- self.vllm_flash_attn_version = get_flash_attn_version()
633
+ self.vllm_flash_attn_version = get_flash_attn_version(
634
+ requires_alibi=self.alibi_slopes is not None)
635
if (is_quantized_kv_cache(self.kv_cache_dtype)
636
and self.vllm_flash_attn_version != 3):
637
raise NotImplementedError(
vllm/fa_utils.py
+9
-3
@@ -7,7 +7,7 @@ from vllm.logger import init_logger
7
logger = init_logger(__name__)
8
9
10
-def get_flash_attn_version() -> Optional[int]:
10
+def get_flash_attn_version(requires_alibi: bool = False) -> Optional[int]:
11
# import here to avoid circular dependencies
12
from vllm.platforms import current_platform
13
try:
@@ -28,8 +28,14 @@ def get_flash_attn_version() -> Optional[int]:
28
29
# 3. fallback for unsupported combinations
30
if device_capability.major == 10 and fa_version == 3:
31
- logger.warning("Cannot use FA version 3 on Blackwell platform",
32
- "defaulting to FA version 2.")
31
+ logger.warning_once(
32
+ "Cannot use FA version 3 on Blackwell platform "
33
+ "defaulting to FA version 2.")
34
+ fa_version = 2
35
+
36
+ if requires_alibi and fa_version == 3:
37
+ logger.warning_once("Cannot use FA version 3 with ALiBi, "
38
+ "defaulting to FA version 2.")
39
fa_version = 2
40
41
if not is_fa_version_supported(fa_version):