[Bug]: Gemma 4 fails to initialize with per-token-head KV cache quantization

[lisp19]

Your current environment

The output of python collect_env.py

Use the builtin dockerfile to build & manually merge PR https://github.com/vllm-project/vllm/pull/39018 for turing devices base compatibility
Main commit hash:726efe177bf22874743d11dfdfef9247dbfb5ff0

🐛 Describe the bug

Gemma 4 fails to initialize in vLLM V1 when --kv_cache_dtype int8_per_token_head is enabled. The same issue also affects fp8_per_token_head.
The failure is caused by Gemma 4's hybrid attention layout using two different KV head dimensions:

• sliding/local layers use head_dim = 256
• global/full-attention layers use head_dim = 512

With per-token-head KV quantization, each KV head also stores scale metadata per token. That extra scale payload breaks the exact 2:1 page-size ratio between the 512-dim and 256-dim layers, so V1 KV page-size unification fails during initialization.

Minimal reproduction:

vllm serve <gemma4-model> \
  --kv_cache_dtype int8_per_token_head

Observed result:

(Worker_TP1 pid=248) ERROR 04-19 14:16:08 [multiproc_executor.py:971]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/kv_cache_utils.py", line 942, in unify_kv_cache_spec_page_size
(Worker_TP1 pid=248) ERROR 04-19 14:16:08 [multiproc_executor.py:971]     raise NotImplementedError(
(Worker_TP1 pid=248) ERROR 04-19 14:16:08 [multiproc_executor.py:971] NotImplementedError: The page size of the layer is not divisible by the maximum page size. Cannot unify by adjusting block_size.

Why this happens:

• local layer page-size factor: (256 * 1 byte * 2) + 8 = 520
• global layer page-size factor: (512 * 1 byte * 2) + 8 = 1032
1032 is not divisible by 520, so the current page-size unification logic cannot reconcile the two layer types by changing block_size alone.

Relevant code paths:

• vllm/model_executor/models/gemma4.py
• vllm/model_executor/layers/attention/attention.py
• vllm/v1/core/kv_cache_utils.py
• vllm/v1/kv_cache_interface.py
• vllm/v1/worker/gpu/attn_utils.py
• vllm/v1/worker/gpu_model_runner.py

Environment

Item	Value
Commit	726efe177bf22874743d11dfdfef9247dbfb5ff0
GPU	2× NVIDIA GeForce RTX 2080 Ti (22 GB, CC 7.5, Turing)
CUDA Driver	580.126.18, CUDA 13.0
Python	3.12 (Docker)
Model	Gemma4-31B AWQ 4-bit (google/gemma-4-31b-it), also reproduced with cyankiwi/gemma-4-31B-it-AWQ-4bit

I'm still new to this project, so if this report doesn't match the community's preferred format or I've misunderstood something basic, corrections are very welcome.

Potential direction for a fix:
I want to post a temporary workaround PR later for reference later, which as I tested, can solve this problem. But it certainly needs more review and discussion.

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.

[lisp19]

Temporary workaround: #40391

[noonghunna]

Just posted a longer cross-rig validation finding on the PR thread: #40391 (comment).

Short version: 2× RTX 3090 (Ampere consumer) + Gemma 4 31B + int8_per_token_head. We hit the same bug, explored two alternatives to lisp19's PR (generic spec-level unifier alone, and worker-runtime logic alone), and both failure modes structurally confirm the PR's full-package design is correct:

• Generic spec-level fix without the worker-side logical-shape view → boots but suffers severe decode-TPS decay (turn-1 33 → turn-5 10 TPS, 30% retention) — vLLM's generic strided-view path uses wrong dimension ordering for standard attention with page_size_padded set
• Worker-only overlay (no spec changes) → strict unifier still rejects → no boot
• Hybrid (generic unifier + worker view, no spec pre-padding) → boots but corrupts output

Adds Ampere consumer + INT8 per-token-head as a third cross-rig validation point alongside @cferra's earlier Blackwell consumer + FP8 validation. Worth surfacing for any reviewer eyeing the PR — the design is well-supported by the alternatives' failure modes. Full diagnostic at our perheadkv-overlay-comparison.md.

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[sergeykad]

Still reproduces on v0.26.0, but the symptom has changed — it no longer fails to initialize, it silently costs KV capacity. Worth noting since the title no longer matches what users on current versions will see.

Config: gemma-4-26B-A4B-it FP8 (30 layers = 25 sliding head_dim 256 / 8 KV heads, 5 global head_dim 512 / 2 KV heads), single L40S (sm89), TRITON_ATTN (forced by the heterogeneous head dims), TP=1, --max-model-len 32768, prefix caching on.

Measured, same 15.18 GiB KV budget, only --kv-cache-dtype changed:

dtype	KV pool	max concurrency @32k
fp8	417,022 tokens	12.73x
fp8_per_token_head	325,335 tokens	9.93x

−22% capacity. The mechanism is the one in the original report: the inline fp32 scale pad makes the per-token page factors 520 (sliding) and 1032 (global). Without it they're 512 and 1024 — an exact 2:1 that unifies for free by doubling block_size on the sliding layers. With it, unification can't reconcile them that way and pads the sliding page up to the global one instead, so 25 of 30 layers pay ~2×. The scales themselves are only ~1.5% (4/head_dim); the cost is entirely in the unification fallback.

#40391 is still open. Flagging in case the init-failure path being resolved elsewhere makes this look fixed when the underlying page-size mismatch isn't.