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 4, 06:15 UTC
[Attention] Move FA3→FA4 upgrade into get_flash_attn_version() (#40815)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Canlin Guo committed
May 5, 2026 at 23:43 UTC
2228fe68687d98fb796e46e83de770ade483d73a
2 files changed
+18
-19
vllm/v1/attention/backends/fa_utils.py
+18
-11
@@ -115,22 +115,29 @@ def get_flash_attn_version(
115
)
116
fa_version = 2
117
118
- # The FA3 kernel rejects s_aux (sinks) when hdim != hdim_v; upgrade to
119
- # FA4 on SM90 when available.
118
+ # Some FA3 unsupported SM90 cases can use FA4 when available.
119
if (
120
fa_version == 3
122
- and has_sinks
123
- and head_size is not None
124
- and head_size_v is not None
125
- and head_size != head_size_v
121
and device_capability.major == 9
122
and is_fa_version_supported(4)
123
):
129
- logger.info_once(
130
- "Diff-KV with sinks: upgrading FlashAttention 3 -> 4",
131
- scope="local",
132
- )
133
- fa_version = 4
124
+ upgrade_reason = None
125
+ if head_size is not None and head_size > 256:
126
+ upgrade_reason = f"FA3 does not support head_size={head_size} on SM90"
127
+ elif (
128
+ has_sinks
129
+ and head_size is not None
130
+ and head_size_v is not None
131
+ and head_size != head_size_v
132
+ ):
133
+ upgrade_reason = "Diff-KV with sinks"
134
+ if upgrade_reason:
135
+ logger.info_once(
136
+ "%s: upgrading FlashAttention 3 -> 4",
137
+ upgrade_reason,
138
+ scope="local",
139
+ )
140
+ fa_version = 4
141
142
# FA4 currently uses batch-shape-dependent scheduling
143
# heuristics on SM100+, which breaks batch invariance.
vllm/v1/attention/backends/flash_attn.py
-8
@@ -638,14 +638,6 @@ class FlashAttentionImpl(AttentionImpl):
638
requires_alibi=alibi_slopes is not None,
639
head_size=head_size,
640
)
641
- # head_size > 256 requires FA4 on SM90+; force upgrade from FA3
642
- if (
643
- head_size > 256
644
- and self.vllm_flash_attn_version == 3
645
- and current_platform.is_cuda()
646
- and current_platform.is_device_capability_family(90)
647
- ):
648
- self.vllm_flash_attn_version = 4
641
logger.info_once(
642
"Using FlashAttention version %s",
643
self.vllm_flash_attn_version,