From 2228fe68687d98fb796e46e83de770ade483d73a Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Tue, 5 May 2026 23:43:03 +0800 Subject: [PATCH] =?UTF-8?q?[Attention]=20Move=20FA3=E2=86=92FA4=20upgrade?= =?UTF-8?q?=20into=20get=5Fflash=5Fattn=5Fversion()=20(#40815)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: gcanlin --- vllm/v1/attention/backends/fa_utils.py | 29 +++++++++++++++--------- vllm/v1/attention/backends/flash_attn.py | 8 ------- 2 files changed, 18 insertions(+), 19 deletions(-) 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,