[Bug]: Gemma4 MTP speculative decoding crashes at engine init on 0.25.1 — "a and b must have same reduction dim" (regression from 0.21.0)

[pryingopen3rdi]

Environment

• vLLM 0.25.1 (pip, PyPI aarch64 wheels), torch 2.11.0+cu130, Python 3.11
• NVIDIA DGX Spark (GB10, aarch64, 121 GB unified memory), single GPU
• Target: coder3101/gemma-4-26B-A4B-it-heretic (bf16 checkpoint, --quantization fp8)
• Drafter: google/gemma-4-26B-A4B-it-assistant

Command

vllm serve coder3101/gemma-4-26B-A4B-it-heretic \
  --host 0.0.0.0 --port 8001 \
  --max-model-len 131072 --max-num-batched-tokens 2560 --max-num-seqs 2 \
  --gpu-memory-utilization 0.55 --dtype bfloat16 --quantization fp8 \
  --enable-auto-tool-choice --tool-call-parser gemma4 --reasoning-parser gemma4 \
  --chat-template <gemma4 template> \
  --speculative-config '{"method":"mtp","model":"google/gemma-4-26B-A4B-it-assistant","num_speculative_tokens":4}'

Behavior

Engine core init fails during dynamo fake-tensor tracing:

torch._subclasses/fake_tensor.py:2908
RuntimeError: a and b must have same reduction dim, but got [s47, 3840] X [5632, 1024].
...
torch._dynamo.exc.TorchRuntimeError: RuntimeError when making fake tensor call
RuntimeError: Engine core initialization failed.

Observations:

• Works on vLLM 0.21.0 — same command, same checkpoints: engine starts, MTP
  gives ~54 tok/s vs ~40 no-spec (mean acceptance length ~2.85). Rolling the env
  back to 0.21.0 restores it, so this is a 0.25 regression.
• The shapes suggest the MTP head's pre_projection expects
  concat(inputs_embeds, hidden_states) = 1792 + 3840 = 5632 input features
  (weight [5632, 1024]), but receives a 3840-dim tensor. The concat is
  present in Gemma4MTP.forward (gemma4_mtp.py: torch.cat([inputs_embeds, hidden_states], dim=-1)), so the mismatch appears to come from the proposer's
  hidden-state buffer wiring (backbone- vs draft-dim) rather than the forward
  itself.
• Adding "enforce_eager": true to the speculative config does NOT help — the
  same shape error moves to graph [0/0], consistent with a real shape bug
  rather than a compile artifact.
• Only tested with --quantization fp8 (on-the-fly quant of the bf16
  checkpoint); happy to re-run plain bf16 if useful.

Expected

Engine starts with the gemma4 assistant MTP drafter, as on 0.21.0.

Possibly related

• [Spec Decode] Add Gemma4 MTP speculative decoding support #41745 (original gemma4 MTP support)
• [Bug][Spec Decode] num_speculative_tokens_per_batch_size + MTP speculator fails full CUDA graph decode capture (InputBatch.make_dummy assert) #48494 (0.25 MTP speculator fails full CUDA graph decode capture)
• [Bug]: DiffusionGemma crashes under tensor-parallel (TP>1) and pipeline-parallel (PP>1) — multi-GPU is unusable #45719 (same "reduction dim" error text, DiffusionGemma TP>1)

[davidscarth]

i noticed a crash and just forced VLLM_USE_V2_MODEL_RUNNER=1 and it seemed to bypass it

this works, i'm currently running it (on the GB10, on vLLM 0.25.1):

sudo docker run --rm -it \
  --gpus all \
  --ipc=host \
  --network host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -v ~/models/templates:/mnt/templates \
  -e VLLM_USE_V2_MODEL_RUNNER=1 \
  ghcr.io/timothystewart6/vllm-gb10:latest \
  vllm serve RedHatAI/gemma-4-26B-A4B-it-NVFP4 \
  --host 0.0.0.0 --port 8000 \
  --gpu-memory-utilization 0.25 \
  --enable-prefix-caching \
  --speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 3}' \
  --kv-cache-dtype fp8 \
  --max-model-len 32768 \
  --max-num-seqs 6 \
  --max-num-batched-tokens 8192 \
  --chat-template /mnt/templates/chat_template.jinja \
  --default-chat-template-kwargs '{"enable_thinking": true}' \
  --enable-auto-tool-choice \
  --tool-call-parser gemma4 \
  --reasoning-parser gemma4

edit: i use the timstewart vllm build in docker, and i updated my chat-template to the latest one google just dropped, but i dont think those have any impact on the issue, the environmental flag is the one to try.