[Bugfix][Attention] Fall back to native FlashInfer decode when XQA cannot serve a KV-cache group's head_dim

[stecasta]

Problem

Serving a hybrid-attention model with the FlashInfer TRTLLM/XQA decode path crashes at engine init:

ValueError: Invalid head_dim: 512, must be divisible by 16 and in range [16, 256]

raised from profile_cudagraph_memory() during determine_available_memory().

The dedicated FlashInfer XQA decode API accepts head dimensions in [16, 256] that are divisible by 16. Gemma 4 is a hybrid-attention model that combines sliding-attention groups at head_dim=256 with global-attention groups at head_dim=512. FlashInferMetadataBuilder resolves the decode kernel once, so a single group the XQA API cannot serve makes the server unstartable instead of degrading to a supported kernel.

This became reachable on SM12x via #49718, which enabled the dedicated XQA decode path there. The same crash occurs on any platform where XQA is selected for such a model.

Fix

Resolve XQA eligibility per KV-cache group. When the resolved kernel is XQA and a group's head_dim falls outside the supported range, that group reverts to native FlashInfer decode and warns once. Eligible groups keep XQA. This mirrors the existing DCP guard immediately below in the same function.

Reproducer

vllm serve nvidia/Gemma-4-26B-A4B-NVFP4 \
  --attention-backend FLASHINFER \
  --attention-config '{"use_trtllm_attention": true}' \
  --kv-cache-dtype fp8 --max-model-len 262144 --language-model-only

Before: ValueError: Invalid head_dim: 512 at startup. After: the server starts, the log records the fallback, and the sliding-attention groups still run XQA.

Testing

Hardware: RTX PRO 6000 Blackwell Max-Q (SM120), FlashInfer 0.6.16.post3.

Backend selection confirmed from the serve log, showing both kernels in use as intended:

AttentionBackendEnum.FLASHINFER
decode_backend=xqa
decode_backend=flashinfer-native

End-to-end benchmark at ISL 150k / OSL 4k / concurrency 1:

config	output tok/s	ITL	request latency
TRITON_ATTN baseline	35.25	22.04 ms	113,463 ms
this change	86.37	5.40 ms	46,304 ms

Lint: pre-commit run --files vllm/v1/attention/backends/flashinfer.py passes, including mypy-3.10.

No new unit test is included. The guard lives in FlashInferMetadataBuilder.__init__, and the sibling DCP guard directly below it is likewise untested. Happy to add coverage if reviewers prefer.

Why this is not a duplicate

Searched open issues and PRs for XQA/head_dim, FlashInfer head_dim, and Gemma attention. #48162 routes backends at batch level, which is orthogonal to per-KV-cache-group kernel eligibility. No open PR addresses the head_dim gate.

AI assistance

AI assistance (Claude) was used to diagnose the failure, locate the fix, and prepare this PR, as disclosed in the commit trailers.

[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.

[mgoin]

/ci run

[github-actions]

✅ @stecasta, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /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.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[github-actions]

✅ Triggered Buildkite CI #85019 for commit ad4a6802b8bf.