[Bugfix][Attention] Fall back to native FlashInfer decode when XQA cannot serve a KV-cache group's head_dim (#53111)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Andrii Skliar <askliar@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
Stefano Castagnetta committed
Aug 21, 2026 at 20:41 UTC
d9e0ace7a07d38095d8b34bad58852cdffa27744
1 file changed
+17
vllm/v1/attention/backends/flashinfer.py
+17
@@ -839,6 +839,23 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
839
if can_use_xqa_or_trtllm_gen_decode
840
else None
841
)
842
+ # The dedicated FlashInfer XQA API accepts head dimensions in
843
+ # [16, 256] that are divisible by 16. Some hybrid-attention models
844
+ # (for example Gemma 4) use XQA-compatible sliding-attention groups
845
+ # alongside global-attention groups with head_dim=512. Resolve the
846
+ # backend per KV-cache group so the eligible groups still use XQA and
847
+ # the wider groups fall back to native FlashInfer decode.
848
+ if (
849
+ self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
850
+ and not (16 <= self.head_dim <= 256 and self.head_dim % 16 == 0)
851
+ ):
852
+ logger.warning_once(
853
+ "FlashInfer XQA decode does not support head_dim=%d; "
854
+ "reverting this KV-cache group to native FlashInfer decode.",
855
+ self.head_dim,
856
+ )
857
+ self.use_trtllm_decode_attention = False
858
+ self.flashinfer_trtllm_api_decode_kernel = None
859
if (
860
self.use_dcp
861
and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA