[Bugfix] Temporarily disable FA4 head-dim 256

[taneem-ibrahim]

Purpose

PR #42669 enabled FA4 head-dim 256 on Blackwell, but the specialized SM100 2-CTA kernel still rejects seqused_q/k, which vLLM decoder attention supplies. Temporarily resolve FA2 for head-dim 256 on Blackwell until upstream adds the required support. FA4 remains enabled for head-dim 128 and the supported MLA 192/128 case. ColPali under MRV2 exposed the failure (#48290).

Reproducer

CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest \
  tests/models/multimodal/pooling/test_colpali.py::test_colpali_multimodal_text_query_image_docs \
  -v

Output on main / on branch

main:
AssertionError: SM100 forward with head_dim=256 does not support seqused_q/seqused_k

branch:
PASSED
1 passed, 16 warnings in 116.81s

Test Plan

.venv/bin/python -m pytest tests/kernels/attention/test_attention_selector.py -q
31 passed, 10 skipped

.venv/bin/python -c "import vllm.v1.attention.backends.mla.prefill.flash_attn; import pytest; raise SystemExit(pytest.main([\"tests/v1/attention/test_mla_prefill_selector.py\", \"-q\"]))"
21 passed

B300 forced-MRV2 ColPali test
1 passed

.venv/bin/pre-commit run --files vllm/v1/attention/backends/fa_utils.py vllm/v1/attention/backends/flash_attn.py tests/kernels/attention/test_attention_selector.py tests/models/multimodal/pooling/test_colpali.py
All hooks passed

AI assistance disclosure

OpenAI Codex (GPT-5) assisted with drafting the code. The submitter reviewed every changed line and ran the tests and model evaluation reported above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yewentao256]

Thanks for the work! Also CC @MatthewBonanni

[yewentao256]

Will this import fail in Rocm?

[taneem-ibrahim]

Good catch, fixed it.

[github-actions]

✅ @taneem-ibrahim, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[yewentao256]

Could you also take a deeper look what is the root cause? I believe fixing accordingly would be better than fall back

[taneem-ibrahim]

Could you also take a deeper look what is the root cause? I believe fixing accordingly would be better than fall back

Root cause is here where it explicitly has this as a TODO .

# Keep parity with FlashAttentionForwardSm100.__call__ interface.
        # (TODO@wangsiyu) Implement these features.
        assert mSeqUsedQ is None and mSeqUsedK is None, (
            "SM100 forward with head_dim=256 does not support seqused_q/seqused_k"
        )

We could wait until FA4 adds the support or do the fall back option for now.

[MatthewBonanni]

Ah thanks, I missed this when I landed #42669. Since this precludes its use for all decoders, I think we should just disable FA4 hdim 256 entirely until we can fix this upstream, rather than piping a new support argument through everything

[MatthewBonanni]

cc @simon-veitner-redhat

[taneem-ibrahim]

Ah thanks, I missed this when I landed #42669. Since this precludes its use for all decoders, I think we should just disable FA4 hdim 256 entirely until we can fix this upstream, rather than piping a new support argument through everything

Thanks Matt! I will disable FA4 for head-dim 256 centrally. Just to confirm: should this disable FA4 head-dim 256 for encoder attention as well?

[simon-veitner-redhat]

i'll take a look at this

[MatthewBonanni]

LGTM, thanks!

[taneem-ibrahim]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #83975.

[taneem-ibrahim]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #83975.

[taneem-ibrahim]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #83975.

[taneem-ibrahim]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #83975.

[taneem-ibrahim]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84037 for commit 2e62addce01b.

[taneem-ibrahim]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84098 for commit de4ca9e18bee.

[yewentao256]

Thanks for the work!

[kzos]

Upstream support is now proposed in Dao-AILab/flash-attention#2810.

It implements seqused_q / seqused_k for the SM100/SM110 head-dim-256 2-CTA forward kernel and normalizes the paged-KV decoder contract (max_seqlen_k as the current batch maximum with a potentially wider block table), which was the second interface guard preventing typical vLLM decoder calls after the seqused assertion.

Validation against the PR source on B300/SM103 covered paged decode with max_seqlen_k=257 and a wider page table, the page-aligned capacity path, q=120 causal paged attention, and CUDA Graph replay after seqused_k grew from 64 to 129. All matched the fp32 reference within BF16 error.

Linking here because this PR intentionally disabled the path pending an upstream fix.