[Bug]: Gemma-4 fails when forcing FLASHINFER attention backend on Blackwell SM120 (head_size not supported)

[dhayanesh]

Your current environment

The output of python collect_env.py

Collecting environment information... ============================== System Info ============================== OS : Ubuntu 24.04.3 LTS (x86_64) GCC version : (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0 Clang version : Could not collect CMake version : version 3.28.3 Libc version : glibc-2.39

==============================
PyTorch Info

PyTorch version : 2.10.0+cu128
Is debug build : False
CUDA used to build PyTorch : 12.8
ROCM used to build PyTorch : N/A
XPU used to build PyTorch : N/A

==============================
Python Environment

Python version : 3.12.3 (main, Aug 14 2025, 17:47:21) [GCC 13.3.0] (64-bit runtime)
Python platform : Linux-6.8.0-106-generic-x86_64-with-glibc2.39

==============================
CUDA / GPU Info

Is CUDA available : True
CUDA runtime version : 12.8.93
CUDA_MODULE_LOADING set to :
GPU models and configuration : GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition
Nvidia driver version : 580.126.09
cuDNN version : Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.8.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.8.0
HIP runtime version : N/A
MIOpen runtime version : N/A
Is XNNPACK available : True

==============================
CPU Info

Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 52 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 128
On-line CPU(s) list: 0-127
Vendor ID: AuthenticAMD
Model name: AMD EPYC 9355 32-Core Processor
CPU family: 26
Model: 2
Thread(s) per core: 2
Core(s) per socket: 32
Socket(s): 2
Stepping: 1
BogoMIPS: 7089.26
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx_vnni avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin cppc amd_ibpb_ret arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid bus_lock_detect movdiri movdir64b overflow_recov succor smca fsrm avx512_vp2intersect flush_l1d debug_swap srso_user_kernel_no
Virtualization: AMD-V
L1d cache: 3 MiB (64 instances)
L1i cache: 2 MiB (64 instances)
L2 cache: 64 MiB (64 instances)
L3 cache: 512 MiB (16 instances)
NUMA node(s): 2
NUMA node0 CPU(s): 0-31,64-95
NUMA node1 CPU(s): 32-63,96-127
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; STIBP always-on; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected

==============================
Versions of relevant libraries

[pip3] flashinfer-python==0.6.6
[pip3] numpy==2.1.2
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cudnn-frontend==1.18.0
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-cufile-cu12==1.13.1.3
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-cutlass-dsl==4.5.0.dev0
[pip3] nvidia-cutlass-dsl-libs-base==4.5.0.dev0
[pip3] nvidia-ml-py==13.595.45
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvshmem-cu12==3.4.5
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] pyzmq==27.1.0
[pip3] torch==2.10.0
[pip3] torch_c_dlpack_ext==0.1.5
[pip3] torchaudio==2.10.0
[pip3] torchvision==0.25.0
[pip3] transformers==5.6.0
[pip3] triton==3.6.0
[conda] Could not collect

==============================
vLLM Info

ROCM Version : Could not collect
vLLM Version : 0.19.1
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; XPU: Disabled
GPU Topology:
GPU0 NIC0 NIC1 CPU Affinity NUMA Affinity GPU NUMA ID
GPU0 X SYS SYS 32-63,96-127 1 N/A
NIC0 SYS X PIX
NIC1 SYS PIX X

Legend:

X = Self
SYS = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
PHB = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
PXB = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
PIX = Connection traversing at most a single PCIe bridge
NV# = Connection traversing a bonded set of # NVLinks

NIC Legend:

NIC0: mlx5_0
NIC1: mlx5_1

==============================
Environment Variables

NVIDIA_VISIBLE_DEVICES=void
NVIDIA_REQUIRE_CUDA=cuda>=12.8 brand=unknown,driver>=470,driver<471 brand=grid,driver>=470,driver<471 brand=tesla,driver>=470,driver<471 brand=nvidia,driver>=470,driver<471 brand=quadro,driver>=470,driver<471 brand=quadrortx,driver>=470,driver<471 brand=nvidiartx,driver>=470,driver<471 brand=vapps,driver>=470,driver<471 brand=vpc,driver>=470,driver<471 brand=vcs,driver>=470,driver<471 brand=vws,driver>=470,driver<471 brand=cloudgaming,driver>=470,driver<471 brand=unknown,driver>=535,driver<536 brand=grid,driver>=535,driver<536 brand=tesla,driver>=535,driver<536 brand=nvidia,driver>=535,driver<536 brand=quadro,driver>=535,driver<536 brand=quadrortx,driver>=535,driver<536 brand=nvidiartx,driver>=535,driver<536 brand=vapps,driver>=535,driver<536 brand=vpc,driver>=535,driver<536 brand=vcs,driver>=535,driver<536 brand=vws,driver>=535,driver<536 brand=cloudgaming,driver>=535,driver<536 brand=unknown,driver>=550,driver<551 brand=grid,driver>=550,driver<551 brand=tesla,driver>=550,driver<551 brand=nvidia,driver>=550,driver<551 brand=quadro,driver>=550,driver<551 brand=quadrortx,driver>=550,driver<551 brand=nvidiartx,driver>=550,driver<551 brand=vapps,driver>=550,driver<551 brand=vpc,driver>=550,driver<551 brand=vcs,driver>=550,driver<551 brand=vws,driver>=550,driver<551 brand=cloudgaming,driver>=550,driver<551 brand=unknown,driver>=560,driver<561 brand=grid,driver>=560,driver<561 brand=tesla,driver>=560,driver<561 brand=nvidia,driver>=560,driver<561 brand=quadro,driver>=560,driver<561 brand=quadrortx,driver>=560,driver<561 brand=nvidiartx,driver>=560,driver<561 brand=vapps,driver>=560,driver<561 brand=vpc,driver>=560,driver<561 brand=vcs,driver>=560,driver<561 brand=vws,driver>=560,driver<561 brand=cloudgaming,driver>=560,driver<561 brand=unknown,driver>=565,driver<566 brand=grid,driver>=565,driver<566 brand=tesla,driver>=565,driver<566 brand=nvidia,driver>=565,driver<566 brand=quadro,driver>=565,driver<566 brand=quadrortx,driver>=565,driver<566 brand=nvidiartx,driver>=565,driver<566 brand=vapps,driver>=565,driver<566 brand=vpc,driver>=565,driver<566 brand=vcs,driver>=565,driver<566 brand=vws,driver>=565,driver<566 brand=cloudgaming,driver>=565,driver<566
NCCL_VERSION=2.25.1-1
NVIDIA_DRIVER_CAPABILITIES=compute,utility
NVIDIA_PRODUCT_NAME=CUDA
CUDA_VERSION=12.8.1
LD_LIBRARY_PATH=/usr/local/cuda/lib64
NVIDIA_CTK_LIBCUDA_DIR=/usr/lib/x86_64-linux-gnu
PYTORCH_NVML_BASED_CUDA_CHECK=1
TORCHINDUCTOR_COMPILE_THREADS=1
TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor_root

🐛 Describe the bug

On a RTX PRO 6000 GPU, forcing the FLASHINFER attention backend for Gemma-4 causes vLLM to fail during engine initialization with:

ValueError: Selected backend AttentionBackendEnum.FLASHINFER is not valid for this configuration. Reason: ['head_size not supported']

Without the override, the same model loads and runs correctly with TRITON_ATTN.

This reproduces for:

• the base model: google/gemma-4-31b-it
• quantized FP8 checkpoint
• quantized NVFP4 checkpoint

Gemma-4 config on this model reports:

• head_dim=256
• global_head_dim=512

Minimal repro

from vllm import LLM

LLM(
    model="google/gemma-4-31b-it",
    tokenizer="google/gemma-4-31b-it",
    dtype="bfloat16",
    tensor_parallel_size=1,
    max_model_len=2048,
    gpu_memory_utilization=0.9,
    enforce_eager=True,
    disable_log_stats=True,
    attention_config={"backend": "FLASHINFER"},
)

Observed failure:

Using AttentionBackendEnum.FLASHINFER backend.
ValueError: Selected backend AttentionBackendEnum.FLASHINFER is not valid for this configuration. Reason: ['head_size not supported']

What I expected

Either:

1. FLASHINFER should work for Gemma-4 on this configuration, or
2. vLLM should reject the override earlier or more explicitly for Gemma-4 with a clearer compatibility message.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[novelsk]

Check out this code: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/config.py#L101-L107

(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299] 
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299]        █     █     █▄   ▄█
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299]  ▄▄ ▄█ █     █     █ ▀▄▀ █  version 0.19.2rc1.dev174+g100c7b65e
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299]   █▄█▀ █     █     █     █  model   google/gemma-4-E2B-it
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299]    ▀▀  ▀▀▀▀▀ ▀▀▀▀▀ ▀     ▀
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:299] 
(APIServer pid=1) INFO 04-24 16:44:31 [utils.py:233] non-default args: {'model_tag': 'google/gemma-4-E2B-it', ...
(APIServer pid=1) INFO 04-24 16:44:34 [model.py:554] Resolved architecture: Gemma4ForConditionalGeneration
(APIServer pid=1) INFO 04-24 16:44:34 [model.py:1672] Using max model len 8192
(APIServer pid=1) INFO 04-24 16:44:34 [cache.py:247] Using fp8 data type to store kv cache. It reduces the GPU ...
(APIServer pid=1) INFO 04-24 16:44:34 [config.py:101] Gemma4 model has heterogeneous head dimensions (head_dim=256, global_head_dim=512). Forcing TRITON_ATTN backend to prevent mixed-backend numerical divergence.

[hsjlyj]

Adding a second Blackwell datapoint, from a consumer RTX 5090 rather than RTX PRO 6000.

Environment:

• 1x NVIDIA GeForce RTX 5090, 32 GiB, sm_120
• Ubuntu 24.04 container, CUDA 13.0 runtime/toolchain path from pip nvidia-* packages
• vLLM 0.22.1 era stack
• FlashInfer 0.6.11 era stack
• Model: Gemma 4 12B Unified IT QAT, compressed-tensors W4A16 checkpoint

Observed behavior:

1. Basic serving could reach API startup only after avoiding the problematic accelerated path:

python3 -m vllm.entrypoints.openai.api_server \
  --model /root/rivermind-data/gemma-4-12B-it-qat-w4a16-ct \
  --quantization compressed-tensors \
  --max-model-len 32768 \
  --gpu-memory-utilization 0.5 \
  --enforce-eager \
  --trust-remote-code \
  --port 8000

With that configuration, /health and /v1/models responded.

2. First generation then failed in FlashInfer paged attention with a kernel-template configuration error:

FlashInfer Internal Error: Invalid configuration:
NUM_MMA_Q=1 NUM_MMA_D_QK=32 NUM_MMA_D_VO=32 NUM_MMA_KV=1 ...

3. Forcing/using FlashInfer therefore appears fragile for this Gemma4 Unified + sm_120 combination, even after the server can boot. This is separate from the earlier setup-level sm120 / CUDA arch string issues; those were only needed to get JIT/compilation far enough to exercise the runtime path.

4. The same deployment also hit CUDA graph / memory pressure when trying the TurboQuant path, but this issue's FlashInfer failure reproduces without needing TurboQuant: compressed-tensors target + no TurboQuant + eager startup still hits FlashInfer on first chat completion.

Net: this looks like the same class of backend/head-size/template support problem reported here, reproduced on RTX 5090 with Gemma4 Unified 12B QAT. A safe fallback route that avoids FlashInfer for this configuration would be useful; otherwise users can get a green /health but fail at first real generation.

[weimin023]

I retested Gemma4 31B on 2x RTX PRO 6000 Blackwell (SM120) with PR #49797 and FlashInfer 0.6.16.post3.
Forcing FLASHINFER still fails, now with:

partial multimodal token full attention not supported

This differs from the original head_size not supported error, but appears related to the same Gemma4 + FlashInfer limitation on Blackwell. Automatic backend selection falls back to TRITON_ATTN, which successfully runs inference.

[ExtReMLapin]

they just merged a pr that fixes it