diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index b3a4a0c76e..1e74e4c48e 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -115,22 +115,29 @@ def get_flash_attn_version( ) fa_version = 2 - # The FA3 kernel rejects s_aux (sinks) when hdim != hdim_v; upgrade to - # FA4 on SM90 when available. + # Some FA3 unsupported SM90 cases can use FA4 when available. if ( fa_version == 3 - and has_sinks - and head_size is not None - and head_size_v is not None - and head_size != head_size_v and device_capability.major == 9 and is_fa_version_supported(4) ): - logger.info_once( - "Diff-KV with sinks: upgrading FlashAttention 3 -> 4", - scope="local", - ) - fa_version = 4 + upgrade_reason = None + if head_size is not None and head_size > 256: + upgrade_reason = f"FA3 does not support head_size={head_size} on SM90" + elif ( + has_sinks + and head_size is not None + and head_size_v is not None + and head_size != head_size_v + ): + upgrade_reason = "Diff-KV with sinks" + if upgrade_reason: + logger.info_once( + "%s: upgrading FlashAttention 3 -> 4", + upgrade_reason, + scope="local", + ) + fa_version = 4 # FA4 currently uses batch-shape-dependent scheduling # heuristics on SM100+, which breaks batch invariance. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index e73954ee74..4b8b86d864 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -638,14 +638,6 @@ class FlashAttentionImpl(AttentionImpl): requires_alibi=alibi_slopes is not None, head_size=head_size, ) - # head_size > 256 requires FA4 on SM90+; force upgrade from FA3 - if ( - head_size > 256 - and self.vllm_flash_attn_version == 3 - and current_platform.is_cuda() - and current_platform.is_device_capability_family(90) - ): - self.vllm_flash_attn_version = 4 logger.info_once( "Using FlashAttention version %s", self.vllm_flash_attn_version,