Feature Request: Support for Google Gemma 4 model family (PLE architecture, variable head dims, KV cache sharing)

[elbruno]

Summary

Google released the Gemma 4 model family (April 2, 2026) with four variants (E2B, E4B, 26B-A4B MoE, 31B). All are Apache 2.0 licensed and use a novel architecture that introduces three features not currently supported by onnxruntime-genai v0.12.2.

I'm the author of ElBruno.LocalLLMs, a C# library that uses ONNX Runtime GenAI for local LLM inference via IChatClient. I've spent significant time attempting to convert Gemma 4 models and have detailed technical findings on what's blocking support.

What I Tried

1. Patched builder.py to route Gemma4ForConditionalGeneration through the Gemma 3 pipeline — conversion produced a 1.6GB ONNX file (INT4), but the runtime failed at load time:

   ShapeInferenceError: Incompatible dimensions for matrix multiplication
   at /model/layers.4/attn/o_proj/MatMul_Q4
   

   Layer 4 is a full-attention layer with global_head_dim=512, but the Gemma 3 model builder created the graph assuming uniform head_dim=256.

2. Changed genai_config.json model type from gemma4 to gemma3_text — same shape error.

3. Examined onnx-community/gemma-4-E2B-it-ONNX — correct ONNX graph structure, but uses a different I/O contract (separate embed_tokens.onnx + decoder_model_merged.onnx with per_layer_inputs tensor) that's incompatible with GenAI's KV cache management.

4. Attempted Gemma4ForCausalLM loading — weights stored under model.language_model.* prefix (multimodal wrapper), causing weight mismatch.

Three Architectural Blockers

1. Per-Layer Embeddings (PLE)

Gemma 4's embed_tokens produces TWO outputs:

• inputs_embeds: [batch, seq, hidden_size] (standard)
• per_layer_inputs: [batch, seq, num_hidden_layers, hidden_size_per_layer_input] — e.g., [batch, seq, 35, 256] for E2B

Each transformer layer receives its own slice from per_layer_inputs. The GenAI runtime currently expects a single embedding tensor flowing into the decoder stack.

Config reference (E2B):

"hidden_size_per_layer_input": 256,
"num_hidden_layers": 35

2. Variable Attention Head Dimensions

Gemma 4 uses two different head dimensions depending on the attention type:

• Sliding attention (most layers): head_dim = 256
• Full attention (every 5th layer — indices 4, 9, 14, 19, 24, 29, 34): global_head_dim = 512

The genai_config.json schema only supports a single head_size field. The KV cache buffer allocation uses this single value for all layers, causing shape mismatches at full-attention layers.

Config reference:

"head_dim": 256,
"global_head_dim": 512,
"attention_pattern": [0, 0, 0, 0, 1, 0, 0, 0, 0, 1, ...]  // 0=sliding, 1=full

3. KV Cache Sharing

Gemma 4 E2B has 35 decoder layers but only 15 unique KV cache pairs. Multiple layers share the same KV cache through a sharing pattern:

"num_kv_shared_layers": 20

The runtime expects one KV cache I/O pair per layer (past_key_values.0 through past_key_values.N). With sharing, the ONNX model only produces 15 unique KV outputs for 35 layers.

Gemma 4 Model Details

Model	Total Params	Effective Params	Architecture	Context	HuggingFace
E2B IT	5.1B	2.3B	Dense + PLE	128K	google/gemma-4-E2B-it
E4B IT	8B	4.5B	Dense + PLE	128K	google/gemma-4-E4B-it
26B-A4B IT	25.2B	3.8B active	MoE + PLE	256K	google/gemma-4-26B-A4B-it
31B IT	30.7B	30.7B	Dense	256K	google/gemma-4-31B-it

All variants use:

• architectures: ["Gemma4ForConditionalGeneration"] (multimodal — no text-only CausalLM variant exists)
• Same chat template as Gemma 2/3: <start_of_turn>role\ncontent<end_of_turn>
• Apache 2.0 license (no gating)
• RoPE with separate theta values for sliding (10000.0) and full attention (1000000.0)

What Would Need to Change

Based on my investigation, support would require:

1. Builder (builder.py): Add Gemma4ForConditionalGeneration case that handles:

   • Extracting text_config from the multimodal config
   • Generating ONNX graph with per_layer_inputs as an additional input
   • Creating per-layer attention nodes with correct head dimensions (256 vs 512)
   • Mapping shared KV caches correctly

2. Runtime (C++):

   • Support per_layer_inputs tensor as model input alongside input_ids/inputs_embeds
   • Handle variable-size KV cache buffers per layer (or allocate to max and slice)
   • Support KV cache sharing (N layers → M unique caches where M < N)

3. Config (genai_config.json):

   • Support per-layer head_size (or head_size + global_head_size)
   • Support num_kv_shared_layers or a KV sharing map
   • New model type gemma4_text with PLE-aware inference loop

Environment

• onnxruntime-genai: 0.12.2 (latest)
• transformers: 5.5.0
• Python: 3.13.9
• OS: Windows 11
• GPU: NVIDIA A10-24Q (24GB VRAM)

Related

• Any plans for Gemma 4 support? #2059 — General question about Gemma 4 support (closed, no technical details)
• onnx-community/gemma-4-E2B-it-ONNX — Transformers.js ONNX export (not GenAI-compatible)

[vgallegob]

I want this too

[Durkee]

me too

[jaburges]

Thanks for the detailed write-up @elbruno — your three-blocker decomposition matched exactly what I hit when I dug into this independently from a C# / Windows ML angle. I opened a Draft PR taking a first pass at the builder side: #2088.

What's in the PR

A Python-only scaffold for Gemma 4 support, following the existing Gemma-family pattern (GemmaModel → Gemma2Model → Gemma3Model → new Gemma4TextModel). It covers everything that can be done without C++ runtime changes:

• Dispatcher branch for Gemma4ForConditionalGeneration and Gemma4ForCausalLM with text_config flattening and exclude_embeds=true default, matching the Gemma 3 branch.
• is_local(layer_id) driven by config.layer_types — on E2B this correctly identifies full-attention layers at indices 4, 9, 14, 19, 24, 29, 34 (7 full + 28 sliding out of 35 total), covered by unit tests.
• make_key_value_cache_shape(layer_id, shape) override using the existing per-layer extension point in base.py to emit head_dim=256 on sliding layers and global_head_dim=512 on full layers.
• make_rotary_embedding_multi_cache() building both caches from config.rope_parameters (sliding θ=10_000, full θ=1_000_000 + partial_rotary_factor=0.25 + rope_type=proportional).
• 10 unit tests that run without torch/transformers/onnxruntime (stubbed parent class), all passing locally.

What's deliberately NOT in the PR

Each of the three blockers you identified needs a paired C++ runtime change, so they're out of scope for Part 1:

Blocker	Runtime file
Per-Layer Embeddings routing	src/models/model.cc + new side-channel input in the graph
KV cache sharing (20 of 35 layers alias)	src/models/kv_cache.cpp
Per-layer head_size	src/models/kv_cache.cpp:17 reads one model.decoder.head_size today

The scaffold raises NotImplementedError with a specific message naming which of these gaps was tripped, so anyone running builder.py against a Gemma 4 checkpoint today gets an actionable error (or a valid genai_config.json under config_only=true) rather than a graph the runtime can't load.

One new data point worth capturing

Your list had attention_pattern: [0, 0, 0, 0, 1, 0, 0, 0, 0, 1, ...] — the current HF config field is actually layer_types: ["sliding_attention", ..., "full_attention", ...] with 35 entries for E2B. Same information, just worth noting because the PR's is_local dispatch reads it in string form.

Also for anyone watching who wants to experiment now: the onnx-community/gemma-4-E2B-it-ONNX export does run end-to-end against plain Microsoft.ML.OnnxRuntime (not GenAI) if you implement tokenization + KV bookkeeping yourself and accept the different I/O contract. It's a useful reference for verifying that per-layer variable head dims + KV sharing are correct before wiring them into GenAI, and I've written one up in C# that I'm happy to share privately with anyone working on the runtime side.

Happy to split Parts 2/3/4 (runtime PLE routing, KV sharing, per-layer head_size) into follow-up PRs against the same issue if maintainers prefer that cadence, or to iterate further on this one.

[justinchuby]

We are iterating on https://huggingface.co/collections/justinchuby/gemma-4-onnx. Will require the source build for onnxruntime-genai. Modeling code will be published when we finish evaluation.