[Attn,KV-cache] Use per-head scales in the attention selector

[eldarkurtic]

As requested by @MatthewBonanni and @LucasWilkinson in #30141 attention backends should be filtered during backend selection based on whether they support per-head attention quantization scales.

This enables early failure when a user attempts to load a model that requires per-head scales but no compatible attention backend is available.

[gemini-code-assist]

Code Review

This pull request refactors the check for per-head attention quantization scales support. The changes introduce a mechanism to filter attention backends early during selection based on this capability. This is achieved by adding a requires_per_head_quant_scales parameter that is propagated down to the attention backend selector. The AttentionBackend class is updated with a supports_per_head_quant_scales method to facilitate this check, with FlashAttentionBackend correctly implementing it based on the FlashAttention version. The previous runtime assertion is removed, enabling earlier failure for unsupported configurations, which improves user experience. The changes are well-structured and correctly implemented.

[MatthewBonanni]

LGTM, maybe we could add a case to the attention selector test to verify this is working?

[eldarkurtic]

tests added, let me know if this is what you had in mind @MatthewBonanni

[mergify]

Hi @eldarkurtic, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[MatthewBonanni]

Thanks for adding the test! Just one comment

[MatthewBonanni]

By patching this I think the test isn't actually exercising the backend support, it'll just always pass as long as supports_per_head matches should_succeed in the test case. Can you get rid of this patch?

Suggested change

patch(supports_attr, return_value=supports_per_head),

[MatthewBonanni]

see below

Suggested change

	backend_name: str, supports_per_head: bool, should_succeed: bool
	backend_name: str, should_succeed: bool

[eldarkurtic]

@MatthewBonanni great point, thanks a lot! I wanted to use that to distinguish between FA2 and FA3, but just found out that I can simply pass flash_attn_version when selecting the backend.
The test is updated and fixed now.

[LucasWilkinson]

LGTM, thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @eldarkurtic.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @eldarkurtic, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint