[Bugfix][Attention] Temporarily disable FA4 head-dim 256 (#52050)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Taneem Ibrahim committed
Aug 16, 2026 at 12:13 UTC
1f0e0bf61210346a6bef4ad75172e62554d1b86c
3 files changed
+15
-14
tests/models/multimodal/pooling/test_colpali.py
+5
@@ -223,7 +223,10 @@ def _run_multimodal_text_query_image_docs_test(
223
max_model_len=4096,
224
enforce_eager=True,
225
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
226
+ attention_backend="FLASH_ATTN",
227
+ kernel_config={"enable_flashinfer_autotune": False},
228
) as vllm_model:
229
+ assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
230
scores = vllm_model.llm.score(query, image_docs)
231
232
assert len(scores) == 2
@@ -297,9 +300,11 @@ def _run_multimodal_image_query_text_docs_test(
300
@pytest.mark.parametrize("dtype", [DTYPE])
301
def test_colpali_multimodal_text_query_image_docs(
302
vllm_runner,
303
+ monkeypatch: pytest.MonkeyPatch,
304
model: str,
305
dtype: str,
306
) -> None:
307
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
308
_run_multimodal_text_query_image_docs_test(vllm_runner, model, dtype=dtype)
309
310
vllm/v1/attention/backends/fa_utils.py
+10
-13
@@ -72,7 +72,6 @@ def get_flash_attn_version(
72
head_size: int | None = None,
73
head_size_v: int | None = None,
74
has_sinks: bool = False,
75
- requires_local_attention: bool = False,
75
) -> int | None:
76
if current_platform.is_xpu():
77
return 2
@@ -169,28 +168,26 @@ def get_flash_attn_version(
168
)
169
fa_version = 2
170
172
- if (
173
- fa_version == 4
174
- and device_capability.major >= 10
175
- and head_size == 256
176
- and requires_local_attention
177
- ):
171
+ # TODO: Restore the `requires_local_attention` restriction when FA4
172
+ # head-dim 256 is re-enabled.
173
+ if fa_version == 4 and device_capability.major >= 10 and head_size == 256:
174
logger.warning_once(
179
- "FA4 on Blackwell does not support local attention with "
180
- "head_size=256, defaulting to FA version 2."
175
+ "FA4 on Blackwell is temporarily disabled for head_size=256, "
176
+ "defaulting to FA version 2."
177
)
178
fa_version = 2
179
180
# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
185
- # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA
186
- # prefill). Development of symmetric 192, 384, and 512 support is being tracked
187
- # in https://github.com/Dao-AILab/flash-attention/issues/2456
181
+ # supported head dimensions to ≤128. The 192/128 MLA prefill case is
182
+ # supported; 256 is temporarily disabled until upstream supports the
183
+ # required features. Development of symmetric 192, 384, and 512 support
184
+ # is tracked in https://github.com/Dao-AILab/flash-attention/issues/2456
185
if (
186
fa_version == 4
187
and device_capability.major >= 10
188
and head_size is not None
189
and head_size > 128
193
- and not (head_size == 256 or (head_size == 192 and head_size_v == 128))
190
+ and not (head_size == 192 and head_size_v == 128)
191
):
192
logger.warning_once(
193
"FA4 on Blackwell does not support head_size=%d due to TMEM "
vllm/v1/attention/backends/flash_attn.py
-1
@@ -905,7 +905,6 @@ class FlashAttentionImpl(AttentionImpl):
905
self.attn_type = attn_type
906
self.vllm_flash_attn_version = get_flash_attn_version(
907
requires_alibi=alibi_slopes is not None,
908
- requires_local_attention=sliding_window is not None,
908
head_size=head_size,
909
has_sinks=sinks is not None,
910
)