[Bug]: Gemma-4 + DFlash unservable on Ampere — non-causal + head_dim=256 has no compatible attention backend

[noonghunna]

Your current environment

2× NVIDIA RTX 3090 (Ampere sm_86), CUDA 12.9, vLLM v0.19.2rc1.dev21+g893611813 (nightly). Full collect-env at the bottom.

🐛 Describe the bug

Serving RedHatAI/gemma-4-31B-it-speculator.dflash (DFlash draft for Gemma-4) against an AutoRound-Int4 Gemma-4 target on TP=2 fails to initialize on every available Ampere-compatible attention backend. DFlash's block-parallel drafting requires non-causal attention; Gemma-4 has head_dim=256. On sm_86, no backend supports both.

Backend matrix

Backend	Error / constraint
FLASH_ATTN (FA2)	head_size not supported — head_dim=256 is not enabled in the stock FA2 wheel on sm_86
FLASH_ATTN_DIFFKV	head_size not supported
FLASHINFER	non-causal attention not supported
TRITON_ATTN	non-causal attention not supported
FLEX_ATTENTION	non-causal attention not supported
TREE_ATTN	non-causal attention not supported

Minor pydantic auto-detect gap

Separate smaller issue: speculators auto-detect doesn't populate num_speculative_tokens from the draft's speculators_config.proposal_methods[0].speculative_tokens. Passing just {\"model\": \"…speculator.dflash\"} raises:

pydantic_core._pydantic_core.ValidationError: 1 validation error for SpeculativeConfig
  Value error, A speculative model was provided, but neither `speculative_token_tree` nor `num_speculative_tokens` was provided

Workaround: set \"num_speculative_tokens\": 8 explicitly (matching the draft's speculative_tokens in its speculators_config). Worth forwarding this field automatically.

Reproduction

docker run --rm --gpus all -p 8000:8000 \\
  -v /models:/models \\
  vllm/vllm-openai:nightly \\
  --model /models/gemma-4-31B-it-int4-AutoRound \\
  --tensor-parallel-size 2 \\
  --disable-custom-all-reduce \\
  --attention-backend FLASHINFER \\
  --speculative-config '{\"method\":\"dflash\",\"model\":\"/models/gemma-4-31B-it-speculator-dflash\",\"num_speculative_tokens\":8}'

Swap --attention-backend across FLASH_ATTN / FLASH_ATTN_DIFFKV / FLASHINFER / TRITON_ATTN / FLEX_ATTENTION / TREE_ATTN to reproduce each row above.

What would unblock Ampere

• Non-causal support added to FLASHINFER / FLEX_ATTENTION / TREE_ATTN, OR
• head_dim=256 enabled in FA2/FA3 wheels for sm_86

Red Hat's card already notes "validated on H100, other hardware validation pending" (https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dflash) — this issue is upstream of their draft, so filing here. Cross-references:

• HF discussion: https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dflash/discussions/1
• llama.cpp DFlash PR comment (parallel finding, different layer): [Speculative decoding] feat: add DFlash support ggml-org/llama.cpp#22105 (comment)

Environment

Collecting environment information...
==============================
        System Info
==============================
OS                           : Ubuntu 22.04.5 LTS (x86_64)
GCC version                  : (Ubuntu 11.4.0-1ubuntu1~22.04.3) 11.4.0
Clang version                : Could not collect
CMake version                : Could not collect
Libc version                 : glibc-2.35

==============================
       PyTorch Info
==============================
PyTorch version              : 2.11.0+cu129
Is debug build               : False
CUDA used to build PyTorch   : 12.9
ROCM used to build PyTorch   : N/A
XPU used to build PyTorch    : N/A

==============================
      Python Environment
==============================
Python version               : 3.12.13 (main, Mar  4 2026, 09:23:07) [GCC 11.4.0] (64-bit runtime)
Python platform              : Linux-6.8.0-110-generic-x86_64-with-glibc2.35
    
==============================
       CUDA / GPU Info
==============================
Is CUDA available            : True
CUDA runtime version         : 12.9.86
CUDA_MODULE_LOADING set to   : 
GPU models and configuration : 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version        : 595.58.03
cuDNN version                : Could not collect
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:                           48 bits physical, 48 bits virtual
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               AuthenticAMD
Model name:                              AMD EPYC 7543 32-Core Processor
CPU family:                              25
Model:                                   1
Thread(s) per core:                      1
Core(s) per socket:                      16
Socket(s):                               1
Stepping:                                1
BogoMIPS:                                5599.65
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 rep_good nopl cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw perfctr_core ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr wbnoinvd arat npt lbrv nrip_save tsc_scale vmcb_clean flushbyasid pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor fsrm
Virtualization:                          AMD-V
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                8 MiB (16 instances)
L3 cache:                                256 MiB (16 instances)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
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:      Mitigation; Safe RET
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; Retpolines; IBPB conditional; IBRS_FW; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Vulnerable: Clear CPU buffers attempted, no microcode
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

==============================
Versions of relevant libraries
==============================
[pip3] flashinfer-python==0.6.7
[pip3] numpy==2.2.6

(Truncated; full dump available on request.)

AI-assisted (Claude) finding, human review & submit.

[Sunt-ing]

I reproduced this on current main and it looks like the backend situation has changed since this was filed. Sharing in case it unblocks you.

TL;DR: On current main, Gemma-4 + DFlash initializes and generates. The DFlash draft's non-causal + head_dim=256 attention is now served by FLEX_ATTENTION, which gained non-causal support in #40394 (merged the same day this issue was filed). If you let vLLM auto-select the attention backend (drop --attention-backend), the drafter picks FLEX_ATTENTION automatically and the engine comes up.

Environment

• GPU: 1x NVIDIA A800-SXM4-80GB (Ampere, sm_80)
• vLLM: source build at commit 4f423bd5b
• torch 2.11.0+cu130
• Draft: RedHatAI/gemma-4-31B-it-speculator.dflash (the model in this issue)
• Target: cyankiwi/gemma-4-31B-it-AWQ-4bit (compressed-tensors W4A16; the RedHatAI/...int4-AutoRound target from the report is no longer public, and RedHatAI/gemma-4-31B-it-FP8-dynamic does not load on Ampere because compressed-tensors w8a8-fp8 requires sm_89+)
• TP=1, enforce_eager=True, num_speculative_tokens=8, parallel_drafting=True (confirms the non-causal drafting path is active)

With auto backend selection the drafter's candidate list filters to exactly ['FLEX_ATTENTION'] and the engine initializes + generates. To make the constraint explicit I forced the drafter's backend via speculative_config.attention_backend and reproduced the per-backend behavior:

forced drafter backend	result	engine reason
FLEX_ATTENTION	init + generate OK	(works)
TRITON_ATTN	init fails	['non-causal attention not supported']
FLASHINFER	init fails	['partial multimodal token full attention not supported', 'non-causal attention not supported']
FLASH_ATTN	init fails	['partial multimodal token full attention not supported']

So among the text-attention backends, FLEX_ATTENTION is currently the only one that satisfies non-causal + head_dim=256. FLEX_ATTENTION is a torch-level backend and is not Ampere-specific, so this should also apply to your 2x3090 (sm_86): on current main, serving without an explicit --attention-backend (or with --attention-backend FLEX_ATTENTION) should let the drafter come up. Could you retry on current main?

(Note: FLASH_ATTN failing above is due to the multimodal Gemma4ForConditionalGeneration target needing mm_prefix, which is a different reason from the head_size not supported you saw on sm_86; on sm_80 FA2 head_dim=256 itself is fine. Either way the practical conclusion is the same.)

---

On the minor auto-detect gap you flagged (passing just {"model": "...speculator.dflash"} without num_speculative_tokens raising a ValidationError): I confirmed it and it has a small fix. When the speculators-format draft is passed via --speculative-config (rather than as the top-level --model), the draft's speculators_config.proposal_methods[0].speculative_tokens was dropped, so num_speculative_tokens could not be defaulted. Forwarding that value onto the draft config lets SpeculativeConfig default it. Happy to open a separate PR for that piece if useful.