[Bugfix] detect alibi and revert to FA2 (#15231)

Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>
This commit is contained in:
Travis Johnson 2025-03-20 20:20:16 -06:00 committed by GitHub
parent 47195057e9
commit 0032903a5b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 4 deletions

View File

@ -630,7 +630,8 @@ class FlashAttentionImpl(AttentionImpl):
self.sliding_window = ((sliding_window - 1,
0) if sliding_window is not None else (-1, -1))
self.kv_cache_dtype = kv_cache_dtype
self.vllm_flash_attn_version = get_flash_attn_version()
self.vllm_flash_attn_version = get_flash_attn_version(
requires_alibi=self.alibi_slopes is not None)
if (is_quantized_kv_cache(self.kv_cache_dtype)
and self.vllm_flash_attn_version != 3):
raise NotImplementedError(

View File

@ -7,7 +7,7 @@ from vllm.logger import init_logger
logger = init_logger(__name__)
def get_flash_attn_version() -> Optional[int]:
def get_flash_attn_version(requires_alibi: bool = False) -> Optional[int]:
# import here to avoid circular dependencies
from vllm.platforms import current_platform
try:
@ -28,8 +28,14 @@ def get_flash_attn_version() -> Optional[int]:
# 3. fallback for unsupported combinations
if device_capability.major == 10 and fa_version == 3:
logger.warning("Cannot use FA version 3 on Blackwell platform",
"defaulting to FA version 2.")
logger.warning_once(
"Cannot use FA version 3 on Blackwell platform "
"defaulting to FA version 2.")
fa_version = 2
if requires_alibi and fa_version == 3:
logger.warning_once("Cannot use FA version 3 with ALiBi, "
"defaulting to FA version 2.")
fa_version = 2
if not is_fa_version_supported(fa_version):