Skip to content

[BugFix] Fallback from FA4->FA2 for Batch Invariance #36059

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yewentao256 merged 3 commits into from
Mar 5, 2026
Merged

[BugFix] Fallback from FA4->FA2 for Batch Invariance #36059

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
46aa13d
Fallback from FA4->FA2 for batch invariance
frankwang28 Mar 4, 2026
71b3dd3
Merge branch 'main' into disable-fa4-batch-inv
frankwang28 Mar 4, 2026
7db7994
Merge branch 'main' into disable-fa4-batch-inv
frankwang28 Mar 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions vllm/v1/attention/backends/fa_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from typing import Any

from vllm.logger import init_logger
from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
from vllm.platforms import current_platform

logger = init_logger(__name__)
Expand Down Expand Up @@ -111,6 +112,16 @@ def get_flash_attn_version(
)
fa_version = 2

# FA4 currently uses batch-shape-dependent scheduling
# heuristics on SM100+, which breaks batch invariance.
if vllm_is_batch_invariant() and fa_version == 4:
logger.warning_once(
"Cannot use FA version 4 with batch invariance, "
"defaulting to FA version 2.",
scope="local",
)
fa_version = 2

# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
# supported head dimensions.
# See: https://github.com/Dao-AILab/flash-attention/issues/1959
Expand Down