Fix Gemma 4 for upcoming Transformers version

[hmellor]

Transformers v5.15.0 will introduce heterogeneous config machinery that has been adopted by Gemma 4.

This PR updates the Gemma 4 implementations and the Transformers modelling backend to be compatible with these new heterogeneous configs.

It also required making some small changes to the model arch converter so that it plays nicely with heterogeneous configs.

Supersedes #48432

---

With huggingface/transformers#47547, here is a performance comparison of vLLM vs Transformers backend for this model (google/gemma-4-12B-it-qat-w4a16-ct, 1xH100, in/out = 1024/512):

Metric	Native (auto)	Transformers	Native / TF
Prompt tokens	65536	65536	—
Generated tokens	32768	32768	—
Wall time (s)	15.039	16.313	0.92×
Gen throughput (tok/s)	2178.8	2008.7	1.08×
Total throughput (tok/s)	6536.5	6026.1	1.08×
Requests/s	4.26	3.92	1.09×
Mean latency/req (s)	0.235	0.255	0.92×

This performance difference comes from the Transformers backend fusions not being generic enough for Gemma 4's shape, I will follow up to fix that.

Note

Model loading might be a little slow on Transformers main due to a known issue with per_layer_config lookups. This will be resolved in huggingface/transformers#47539 before the release of v5.15.0.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[hmellor]

cc @charlotte12l for model arch converter
cc @lucianommartins for Gemma
cc @gante for Diffusion Gemma (no direct diff but it inherits the Gemma 4 changes)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @hmellor.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[lucianommartins]

thanks for taking this on, @hmellor! I ran the PR's logic against Transformers main with the actual Hub
configs and it holds up - including on the two points that were disputed on #48432. Details below, since
I think the evidence is worth having on the record... then a few smaller things.

The gate is correct, and the proposed alternative would break

the concern raised on #48432 was that hasattr(config, "is_heterogeneous") is a version check
rather than a layout check, so a currently-published (old-layout) checkpoint would take the
per-layer branch, find no overrides, and silently resolve full-attention layers to head_dim=256.

that does not reproduce. Gemma4TextConfig.__post_init__ back-fills the legacy keys
(configuration_gemma4.py, from #47384 itself):

global_head_dim = kwargs.pop("global_head_dim", 512)
num_global_key_value_heads = kwargs.pop("num_global_key_value_heads", None)
if "per_layer_config" not in kwargs:
    ...

so old-layout checkpoints come out already migrated. Loading the published configs on 5.15:

model	sliding	full
gemma-4-12B-it	head_dim=256, kv=8	head_dim=512, kv=1
gemma-4-26B-A4B-it	head_dim=256, kv=8	head_dim=512, kv=2
gemma-4-E4B-it	head_dim=256, kv=2	head_dim=512, kv=2

all three report is_heterogeneous=True, so hasattr and getattr(..., False) agree in every
real case.

the suggested alternative - keying off global_head_dim presence - would actively break on 5.15:
hasattr(text_config, "global_head_dim") is False there (the key is consumed by __post_init__),
so it would route 5.15 configs into the legacy branch, which does config.head_dim and raises.
So the gate as written is the right one...

getattr(cfg, "head_dim", default) is not a protection

this is the part I think is worth calling out in the PR, because it changes how the rest of the
codebase has to be audited. AmbiguousGlobalPerLayerAttributeError is not an AttributeError
subclass, so on a heterogeneous config:

hasattr(t, "head_dim")          -> RAISES AmbiguousGlobalPerLayerAttributeError
getattr(t, "head_dim", None)    -> RAISES AmbiguousGlobalPerLayerAttributeError
"head_dim" in t.to_dict()       -> True          # the safe probe

the usual defensive idiom silently stops being defensive - two consequences:

1. the new get_total_num_kv_heads() override is load-bearing, not cosmetic... without it the
   base implementation's getattr_iter(hf_text_config, [..., "num_key_value_heads", ...]) raises...
   worth a comment saying so, or someone will "simplify" it away later on the grounds that the max
   is a no-op for Gemma 4 (which it is - see below)
2. the or super().get_head_size() / or super().get_total_num_kv_heads() fallbacks are
   landmines - ie. if the per-layer max is ever falsy, super() does
   getattr(self.hf_text_config, "head_dim", None) and raises instead of returning None...
   unreachable today (512 and 8 are truthy), but it turns "couldn't determine" into a hard crash.

I swept vLLM for global reads of the two per-layer attributes (head_dim, num_key_value_heads

• confirmed via per_layer_attributes) and the PR looks complete: get_quant_config,
  Gemma4Config.verify_and_update_config, the three model files and the convertor are all covered...
  the remaining hits are unreachable for Gemma 4 (mla_attention.py, the is_deepseek_mla() branch,
  mamba2.py).

allso worth noting the MTP drafters are heterogeneous too - gemma-4-12B-it-assistant and
-26B-A4B-it-assistant both report per_layer_attributes=['head_dim', 'num_key_value_heads'] -
so gemma4_mtp.py being in this PR is necessary, not just tidy,..

the per-layer accessor: correct, but please pin it down

get_head_size() / get_total_num_kv_heads() are written as aggregates ("return the largest"),
but get_model_arch_config(layer_idx=...) also hands them a single per-layer config, where the
answer should be that layer's value. I expected that to misbehave - it doesn't:

12B  layer 0 (sliding) -> head_size=256 kv=8    (config: 256 / 8)   OK
12B  layer 5 (full)    -> head_size=512 kv=1    (config: 512 / 1)   OK
26B  layer 5 (full)    -> head_size=512 kv=2    (config: 512 / 2)   OK

looks like it works because a per-layer config object reports is_heterogeneous=False but still exposes a
per_layer_config view with no overrides, so every entry returns that layer's own value and the
max collapses to it.

that is three non-obvious behaviours deep, and nothing documents or tests it - a comment on the two
methods saying they're valid for both a whole-model and a single-layer config, plus a test pinning
the numbers above, would keep a future refactor from quietly breaking it.

Performance: not the problem it was reported to be

I couldn't reproduce on main the 47 ms get_head_size() / +8.3 s startup figures from #48432... masured
on the real 12B config (48 layers):

single per_layer_config[i] access (by index)	0.032 ms
single per_layer_config["full_attention"] access (by type)	0.035 ms
get_head_size() by type, set()-deduped	0.07 ms
get_head_size() by type, un-deduped	1.36 ms
full convert()	0.204 ms
create_attention_instances total (2 accessors × 48 layers)	~20 ms

by-type versus by-index is a wash per access, so there's nothing to fix there. ~20 ms of added
startup is fine.

one trivial leftover: Gemma4Config.verify_and_update_config iterates hf_text_config.layer_types
un-deduped - that list is one entry per layer, so 48 lookups where 2 would do (1.36 ms vs
0.07 ms). model_arch_config_convertor.py already uses set(...) for the identical job. Worth
matching, purely for consistency.

total_num_kv_heads is now doing two jobs

switching get_quant_config to the arch converter is a real improvement - it picks up Falcon's
n_head_kv, ChatGLM's multi_query_group_num etc. that the old raw getattr missed, and it's
required anyway per the raising-getattr point above.

But note what it feeds: the whole-model value is total_num_kv_heads=8 for the 12B, while the
full-attention layers actually have kv=1, and the consumer is "TP-aware loading of attn_head
scales" - where the max is not obviously the right quantity, unlike the KV-cache buffer sizing the
new docstring cites as its rationale... this is pre-existing (the old code also read 8), so not a
regression, but routing it through a method documented for a different purpose cements the
conflation. gemma-4-12B-it-qat-w4a16-ct (the model in your benchmark table) is a
compressed-tensors checkpoint, so this path is live... so a comment at minimum would be great;
ideally the per-head-scale consumer asks for per-layer counts.

smaller things

• Both attribute removals in gemma4.py check out: self.use_k_eq_v has no readers and
  self.is_full_attention is only used inside __init__. gemma4_dspark.py correctly keeps its
  own self.use_k_eq_v, which is read in forward and asserted on. no issue - just recording it
• != to is not in _get_transformers_backend_cls is right and matches what get_hf_text_config
  already does for the nesting test... worth one line of comment on why (a deep config comparison
  raises on heterogeneous configs), since is not on configs otherwise reads as a bug
• transformers/base.py makes head_size and num_kv_heads per-layer but leaves num_heads
  global... correct for Gemma 4 today; a comment would help if anything ever varies
  num_attention_heads per layer.

Tests

this is my one substantive ask. There are none, and this PR silently changes per-layer attention
geometry for the whole Gemma 4 family on a path CI can't exercise until 5.15.0 ships. The three
behaviours above that work for non-obvious reasons - legacy back-fill, per-layer-config collapse,
and raising getattr - are each one upstream refactor away from breaking silently... synthetic
configs asserting the per-layer table (256/8 sliding, 512/1 full for the 12B) would be cheap and
would have settled the whole #48432 disagreement in CI rather than in review. #48432 had
tests/config/test_multimodal_config.py and tests/model_executor/test_gemma4_layer_geometry.py
which should mostly drop straight in.

I can run the full family (12B/26B/31B/E2B/E4B) end to end once there's a build against 5.15.0 -
happy to do that before or after merge, whichever is more useful.

my apologies for the huge comment, @hmellor :) thanks for pushing it forward!

[hmellor]

Thanks for reviewing @lucianommartins

---

looks like it works because a per-layer config object reports is_heterogeneous=False but still exposes a
per_layer_config view with no overrides, so every entry returns that layer's own value and the
max collapses to it.

This is very much intentional otherwise any model that could be heterogeneous would need to be littered with many if config.is_heterogeneous. The behaviour where per_layer_config always exists allows us to write the logic once and if the model is homogeneous, that isn't a problem.

But note what it feeds: the whole-model value is total_num_kv_heads=8 for the 12B, while the
full-attention layers actually have kv=1, and the consumer is "TP-aware loading of attn_head
scales" - where the max is not obviously the right quantity

Hmm, I'm not really sure what should be done in this case then. Perhaps a util on the Transformers side that gets all the possible values for a field per layer type as {"full_attention": 2, "sliding_attention": 4} and then the consumer decides what to do with it?

transformers/base.py makes head_size and num_kv_heads per-layer but leaves num_heads
global... correct for Gemma 4 today; a comment would help if anything ever varies
num_attention_heads per layer.

Yes I am still working on a way to generalise this properly that doesn't require me to add layer_idx to every getter in ModelArchConfig. I am drafting something locally and will push soon.