PR #38891: [Gemma4] Allow per-layer attention backend selection for heterogeneou…

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

Description (problem / solution / changelog)

Purpose

Fix #38887: Gemma 4 models are extremely slow on vLLM v0.19.0 (~9 tok/s on RTX 4090 for E4B) because Gemma4Config forces all layers to use TRITON_ATTN, even though ~83% of layers (sliding-window, head_dim=256) are fully compatible with FlashAttention.
Root cause: Gemma 4 has heterogeneous head dimensions — sliding-window layers use head_dim=256 and full-attention layers use global_head_dim=512. Since FlashAttention's kernel limit is head_size <= 256, the previous code forced TRITON_ATTN globally to avoid mixed-backend usage. However, vLLM's get_attn_backend() already supports per-layer backend selection via its @cache-decorated selector
(distinct head_size arguments produce distinct backend choices). The global forcing was unnecessary and
penalized the majority of layers.

What this PR does:

• Removes the forced attention_config.backend = TRITON_ATTN override
• Adds informational logging showing the sliding/full layer count split
• Adds an early warning if the user explicitly sets --attention-backend to a backend that cannot handle
  global_head_dim

Layer breakdown across all Gemma 4 variants:

Not a duplicate: Searched open PRs for 38887 in:body and Gemma4 FlashAttention heterogeneous — no
results. PR #38879 optimizes Gemma4 prefill (YOCO fast prefill) and is complementary; it does not address the decode bottleneck caused by the forced TRITON_ATTN backend.

Test Plan

# Lint (pre-commit)
pre-commit run ruff-check --files vllm/model_executor/models/config.py
pre-commit run ruff-format --files vllm/model_executor/models/config.py
pre-commit run typos --files vllm/model_executor/models/config.py

# Unit tests (requires GPU + model access)
.venv/bin/python -m pytest tests/models/multimodal/processing/test_gemma4.py -v

# Serving benchmark (Gemma4 E4B on RTX 4090)
# Before (forced TRITON_ATTN):
vllm serve google/gemma-4-e4b-it --max-model-len 8192 --dtype bfloat16

# After (per-layer backend selection):
vllm serve google/gemma-4-e4b-it --max-model-len 8192 --dtype bfloat16
# Expected log output:
#   Gemma4 model has heterogeneous head dimensions (head_dim=256, global_head_dim=512).
#   35 sliding-window layers will use FlashAttention;
#   7 full-attention layers will fall back to a compatible backend.

Test Result

Lint: All passed (ruff-check, ruff-format, typos).

## Changed files

- `vllm/model_executor/models/config.py` (modified, +66/-34)