PR #40391: Fix Gemma4 KV cache page-size alignment for per-token-head quantization

• Repository: vllm-project/vllm
• Author: lisp19
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/40391

Description (problem / solution / changelog)

Purpose

Fix Gemma4 initialization failure in V1 when per-token-head KV cache quantization is enabled (int8_per_token_head / fp8_per_token_head), as reported in #40388.

Root cause: Gemma4 hybrid attention mixes local (head_dim=256) and global (head_dim=512) layers. With per-token-head KV quantization, per-token scale metadata changes page-size factors to:

• local: (256 * 1) * 2 + 8 = 520
• global: (512 * 1) * 2 + 8 = 1032 1032 is not divisible by 520, so V1 page-size unification can fail.

Additionally, when page_size_padded is present, block-size scaling and KV cache shape reconstruction paths needed consistent padded-size handling.

Fix:

• Add AttentionSpec.copy_with_new_block_size handling to correctly scale and validate page_size_padded when block size changes.
• In Gemma4, set kv_cache_page_size_padded for global-attention layers (1040-based padded factor) under per-token-head cache dtypes, and propagate this through Attention KV spec construction.
• Add shared helper adjust_kv_cache_shape_for_padded_page_size and use it consistently in:
  • vllm/v1/worker/gpu/attn_utils.py
  • vllm/v1/worker/gpu_model_runner.py
  • vllm/v1/worker/kv_connector_model_runner_mixin.py

Fixes issue #40388.

Test Plan

Runtime reproduction for issue validation on Gemma4 environment:

vllm serve <gemma4-model> --kv_cache_dtype int8_per_token_head

Because this change was tested on Turing hardware, #39018 should be merged first before this PR.

Test Results

(APIServer pid=1) INFO:     Started server process [1]
(APIServer pid=1) INFO:     Waiting for application startup.
(APIServer pid=1) INFO:     Application startup complete.
(APIServer pid=1) INFO:     172.17.0.1:33632 - "GET /v1/models HTTP/1.1" 200 OK
(APIServer pid=1) INFO:     172.17.0.1:33640 - "POST /v1/chat/completions HTTP/1.1" 200 OK
(APIServer pid=1) INFO 04-20 17:59:23 [loggers.py:271] Engine 000: Avg prompt throughput: 2.4 tokens/s, Avg generation throughput: 10.2 tokens/s, Running: 0 reqs, Waiting: 0 reqs, GPU KV cache usage: 0.0%, Prefix cache hit rate: 0.0%
(APIServer pid=1) INFO 04-20 17:59:33 [loggers.py:271] Engine 000: Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 0 reqs, Waiting: 0 reqs, GPU KV cache usage: 0.0%, Prefix cache hit rate: 0.0%
(APIServer pid=1) INFO:     172.17.0.1:40910 - "GET /v1/models HTTP/1.1" 200 OK
(APIServer pid=1) INFO:     172.17.0.1:40920 - "POST /v1/chat/completions HTTP/1.1" 200 OK
(APIServer pid=1) INFO 04-20 18:00:33 [loggers.py:271] Engine 000: Avg prompt throughput: 2.4 tokens/s, Avg generation throughput: 9.5 tokens/s, Running: 0 reqs, Waiting: 0 reqs, GPU KV cache usage: 0.0%, Prefix cache hit rate: 0.0%

As a relatively new contributor in this area, I truly appreciate detailed review comments and suggestions, and I will actively iterate on this PR based on feedback. Plz let me know if there are any questions on this change.

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Changed files

• tests/v1/core/test_kv_cache_utils.py (modified, +28/-0)
• tests/v1/worker/test_kv_cache_shape_utils.py (added, +75/-0)
• vllm/model_executor/layers/attention/attention.py (modified, +5/-0)
• vllm/model_executor/models/gemma4.py (modified, +24/-0)
• vllm/v1/core/kv_cache_utils.py (modified, +1/-1)
• vllm/v1/kv_cache_interface.py (modified, +17/-0)
• vllm/v1/worker/gpu/attn_utils.py (modified, +17/-6)
• vllm/v1/worker/gpu_model_runner.py (modified, +24/-6)
• vllm/v1/worker/kv_cache_shape_utils.py (added, +90/-0)
• vllm/v1/worker/kv_connector_model_runner_mixin.py (modified, +25/-6)