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 7, 05:15 UTC
[BugFix] Fallback from FA4->FA2 for Batch Invariance (#36059)
Signed-off-by: frankwang28 <frank.wbb@hotmail.com>
Frank Wang committed
Mar 5, 2026 at 11:05 UTC
a57c877f18188cb7bafc0fc5309b6c88fe2a8f66
1 file changed
+11
vllm/v1/attention/backends/fa_utils.py
+11
@@ -4,6 +4,7 @@
4
from typing import Any
5
6
from vllm.logger import init_logger
7
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
8
from vllm.platforms import current_platform
9
10
logger = init_logger(__name__)
@@ -111,6 +112,16 @@ def get_flash_attn_version(
112
)
113
fa_version = 2
114
115
+ # FA4 currently uses batch-shape-dependent scheduling
116
+ # heuristics on SM100+, which breaks batch invariance.
117
+ if vllm_is_batch_invariant() and fa_version == 4:
118
+ logger.warning_once(
119
+ "Cannot use FA version 4 with batch invariance, "
120
+ "defaulting to FA version 2.",
121
+ scope="local",
122
+ )
123
+ fa_version = 2
124
+
125
# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
126
# supported head dimensions.
127
# See: https://github.com/Dao-AILab/flash-attention/issues/1959