Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions vllm/v1/attention/backends/flashinfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -838,6 +838,23 @@ def __init__(
if can_use_xqa_or_trtllm_gen_decode
else None
)
# The dedicated FlashInfer XQA API accepts head dimensions in
# [16, 256] that are divisible by 16. Some hybrid-attention models
# (for example Gemma 4) use XQA-compatible sliding-attention groups
# alongside global-attention groups with head_dim=512. Resolve the
# backend per KV-cache group so the eligible groups still use XQA and
# the wider groups fall back to native FlashInfer decode.
if (
self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
and not (16 <= self.head_dim <= 256 and self.head_dim % 16 == 0)
):
logger.warning_once(
"FlashInfer XQA decode does not support head_dim=%d; "
"reverting this KV-cache group to native FlashInfer decode.",
self.head_dim,
)
self.use_trtllm_decode_attention = False
self.flashinfer_trtllm_api_decode_kernel = None
if (
self.use_dcp
and self.flashinfer_trtllm_api_decode_kernel == FlashInferDecodeKernel.XQA
Expand Down
Loading