[Bugfix][Attention] Resolve kernel block-size constraints per layer head size

[xutingl]

Purpose

FlashAttentionBackend.get_supported_kernel_block_sizes() decides its FA4 hd256 page constraint ([128] on SM100/SM110) from the model-wide ModelConfig head size. Layers with a different head size, in practice a spec-decode drafter (DFlash, EAGLE-3, DSpark, MTP) with smaller heads than its target, therefore get the target's constraint instead of their own at every place block sizes are chosen:

1. backend selection under a user --block-size (validate_configuration);
2. the default block size (Platform.update_block_size_for_backend, which since [Bugfix] Pick a KV block size supported by every attention backend #49845 consults every backend found in the loaded layers, including the drafter's);
3. each KV cache group's kernel block size (prepare_kernel_block_sizes / select_common_block_size);
4. a sliding-window layer's own block size (_largest_kernel_block_within).

Example: Qwen3.8-27B (GDN hybrid, head size 256, full attention on FlashInfer) + z-lab/Qwen3.8-27B-DFlash2 (head size 128, FlashAttention) on GB200. The Mamba alignment gives an 832-token block, which the drafter's 128-token page constraint does not divide.

• Before [Bugfix] Pick a KV block size supported by every attention backend #49845, KV-cache init failed: ValueError: No common block size for 832. ([Bugfix][Attention] Resolve kernel block sizes per attention group #58425).
• On main today with the default block size, Phase 1 honors the drafter's constraint and picks 896 instead of 832: Mamba page padding 9.5% instead of 1.7% and a 4.5% smaller KV cache.
• On main today with --block-size 64, FlashAttention is precluded for the drafter, which falls back to FlashInfer and, lacking CUDA-graph support for the DFlash2 draft, runs it eagerly: 1.7x slower decode.

#58207 notes the same failure for Qwen3.6-35B-A3B + DFlash on B300.

Fix

Resolve a backend's block-size constraints from the head size of the layers it serves. head_size is FlashAttention's only per-layer input to the decision, and it is available at all four sites.

• AttentionBackend.get_supported_kernel_block_sizes_for_head_size(head_size) and supports_block_size_for_head_size(head_size, block_size), defaulting to the model-wide methods (None means model-wide). Other backends are unchanged, including CPU_MLA's exact-size supports_block_size override. The block-size test shared by supports_block_size moves into supports_kernel_block_size.
• FlashAttentionBackend overrides both, resolving the hd256 page from the given head size.
• validate_configuration checks the block size for the layer's head size.
• Platform._find_non_ssm_backend_head_sizes collects distinct (backend, head size) pairs from the layers; _preferred_block_size_for_backends takes the head sizes as an optional parallel argument. _find_non_ssm_backends, its CPU/XPU callers and the later alignment phases are unchanged.
• select_common_block_size takes optional per-backend head sizes; prepare_kernel_block_sizes passes each attention group's. Other callers are unchanged.
• _largest_kernel_block_within takes the layer's head size.

Only layers whose head size differs from the model-wide one change. A smaller-head drafter behind a head-size-256 target is no longer held to the 128-token page. Conversely, a head-size-256 drafter behind a smaller-head target now makes the default block size 128 model-wide, as a head-size-256 target does, instead of silently running FA2 when the block size is not a multiple of 128.

Supersedes #58425, which only covered site 3.

Test Plan

pytest tests/kernels/attention/test_attention_selector.py -k fa4_hd256_block_size_advertisement
pytest tests/v1/worker/test_gpu_model_runner.py -k "select_common_block_size or preferred_block_size"
pytest tests/v1/worker/test_attn_utils.py -k "kernel_block_size or hisparse_block_size or preserve_global"

E2E on one GB200: Qwen3.8-27B (bf16) + z-lab/Qwen3.8-27B-DFlash2 (num_speculative_tokens=7), default block size and --block-size 64, on main and with this PR; single-request decode benchmark with the acceptance length pinned to 3.

Test Result

Unit tests pass (new: test_fa4_hd256_block_size_advertisement_per_head_size, test_preferred_block_size_resolves_constraints_per_head_size, test_select_common_block_size_resolves_sizes_per_head_size).

	main (8b660ce)	this PR
default block size	block 896, Mamba padding 9.54%, KV cache 578k tokens, TPOT 4.53 ms	block 832, padding 1.71%, KV cache 605k tokens, TPOT 4.53 ms
--block-size 64	drafter precluded from FlashAttention, eager FlashInfer draft, TPOT 7.88 ms	drafter on FlashAttention with CUDA graphs, TPOT 4.51 ms

---

Code written with AI assistance (Claude); reviewed, tested and validated end to end by the author.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[mergify]

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

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

[xutingl]

Replaced by #58630 (redesigned on the spec-keyed hook from #53175).