Add Gemma-4 E-series (E2B/E4B) support: PLE, KV-layer sharing, double-wide MLP, and a dense GGUF parser

[mkornreich]

Summary

Adds support for the Gemma-4 E-series (E2B / E4B) — the small dense models — which previously could not load on FreeToken. The gemma4 module targeted only the dense/MoE line (12B / 26B-A4B / 31B) and the GGUF parser was MoE-only, so E2B/E4B failed: the safetensors path silently ignored the E-series features (Per-Layer Embeddings, KV-layer sharing, double-wide MLP) and the GGUF path raised KeyError on expert_count.

This implements those three features (math ported from the transformers gemma4 reference) plus a dense (non-MoE) GGUF parser.

What's added

• Per-Layer Embeddings (PLE) — embed_tokens_per_layer + per_layer_model_projection + a per-layer gate → gelu → ⊙ per_layer_input → proj → norm step inserted after the feed-forward sandwich and before layer_scalar. The large per-layer table (~4.7 GB bf16) can be kept in host RAM via a new generic cpu_offloaded_weight_keys() model hook that _materialize_loaded_weight_state_dict honors — so E2B fits an 8 GB GPU. In the GGUF path it loads as a compact Q6_K GGUFEmbedding on-GPU instead.
• KV-layer sharing — the last num_kv_shared_layers layers carry q_proj only and reuse the K/V of the last non-shared layer of the same attention type (sliding vs full), via a CUDA-graph-safe per-type stash (the source layer publishes post-rope K / post-norm V into a persistent buffer; shared layers read it). Verified correct with CUDA graphs enabled.
• Double-wide MLP on the KV-shared layers (use_double_wide_mlp).
• Dense Gemma-4 GGUF parser — parse_gguf_config detects dense (no expert metadata) and reads PLE / shared_kv_layers / per-layer feed_forward_length (→ double-wide) / k_eq_v (via value_length); iter_gguf_weights maps the PLE tensors and drops shared-layer k/v; convert_gemma4_to_gguf swaps the PLE table and shared-layer q_proj to native GGUF ops.
• ModelConfig gains per_layer_hidden_size, per_layer_vocab_size, num_kv_shared_layers, use_double_wide_mlp, plus is_kv_shared_layer / is_kv_source_layer helpers.

Scope / safety

Every new code path gates on per_layer_hidden_size > 0 / num_kv_shared_layers > 0 / moe_enabled, so the existing dense/MoE Gemma-4 checkpoints (12B / 26B / 31B) and all other models are unaffected.

Verification

google/gemma-4-E2B-it generates coherent, correct output on an 8 GB RTX 5070, both ways:

• bf16 safetensors (PLE offloaded to host RAM), ~65 tok/s, CUDA graphs off.
• Q4_0 GGUF, ~168 tok/s with CUDA graphs on, ~22 ms TTFT.

Samples: "capital of France" → "Paris"; "15 + 27" → "42"; three primes → 2, 3, 5.

Notes / follow-ups

• The bf16 PLE-CPU-offload path requires CUDA graphs off (the host-side gather can't be captured); the GGUF path keeps the table on-GPU and runs with graphs on.
• KV-sharing currently gives each shared layer its own KV cache populated from the stash (correct; small extra KV for the tiny E-series) rather than true zero-copy sharing — a straightforward memory optimization for a follow-up.
• E-series is TP=1 only (mirrors the existing gemma4 TP=1 restriction).

🤖 Generated with Claude Code

[gdevenyi]

Triaged while assembling a merged deployment branch for a 2× RTX 6000 Ada / 2× Xeon Gold 6526Y Linux box serving DeepSeek-V4-Flash with offloaded experts, to benchmark the open PRs together.

Gemma-4 E-series. Different model family, not merged or tested here.

Flagging only so the absence of a report from me is not read as a problem found — I merged and benchmarked #30, #48, #56, #69, #70, #71 and #81, and left this one out deliberately.

[mkornreich]

@gdevenyi Are you planning on trying to test support for that model? I made the change since it is the only way to get Gemma-4 fully running on my AMD Radeon 890M GPU and not spill into the CPU

[rakhimovv]

@mkornreich I had a spare H100 box and google/gemma-4-E2B-it already downloaded, so I gave this a run on NVIDIA hardware. Short version: it works, and the only thing standing in the way of a default-flags boot is one CPU round-trip that CUDA graph capture will not allow.

It generates correctly with CUDA graphs off.

ft serve --model <gemma-4-E2B-it snapshot> --served-model-name gemma4-e2b \
  --host 127.0.0.1 --port 21941 --cuda-graph-max-bs 0

boots to API server is ready to serve, and a greedy request comes back clean:

prompt:  "Name the capital of France, then count from 1 to 5."
content: "The capital of France is **Paris**.\n\nHere is the count from 1 to 5:\n\n1. One\n2. Two\n3. Three\n4. Four\n5. Five"
finish_reason: stop   usage: 24 prompt / 41 completion

With graphs enabled (the default) it fails during capture, in the per-layer-embedding path:

File "python/freetoken/engine/graph.py", line 179, in _capture_graphs
    self.buffer.logits[:bs] = model.forward()
File "python/freetoken/models/gemma4/model.py", line 209, in forward
    get_global_ctx()._gemma4_ple = self._compute_per_layer_inputs(input_ids, x)
File "python/freetoken/models/gemma4/model.py", line 196, in _compute_per_layer_inputs
    tok = self.embed_tokens_per_layer.forward(input_ids).to(inputs_embeds.device).view(
File "python/freetoken/models/gemma4/model.py", line 50, in forward
    idx = input_ids.to("cpu", dtype=torch.long)
RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture
unless the CPU tensor is pinned. Please use tensor.pin_memory() or allocate the
tensor with pin_memory=True.

model.py:50 moves input_ids to the host to index the PLE table. That is fine in eager mode but illegal inside a capture unless the staging tensor is pinned — which may be why this did not show up for you if graph capture is off or unsupported on the 890M. A persistent pinned staging buffer for idx, or keeping the lookup on device, would presumably let the default path capture.

Separately, the double-wide layer split matches what you implemented. On main this checkpoint dies at layer 15 in layers/base.py with a messageless AssertionError; adding a message locally names it as model.layers.15.feed_forward.shared_mlp.gate_up_proj.weight, model (12288, 1536) vs checkpoint (24576, 1536). Reading the safetensors header directly, mlp.gate_proj.weight rows are 6144 for layers 0–14 and 12288 for layers 15–34, over 35 layers with num_kv_shared_layers: 20 — so the wide layers are exactly the 20 KV-shared ones, as this PR assumes. With the PR applied the load completes without that assert.

Tested on: H100 80GB HBM3 (sm_90), driver 580.126.16, Xeon Platinum 8462Y+, Linux 5.15.0-157-generic, nvcc 13.1, torch 2.11.0+cu130, Python 3.12.13, PR branch at 6486907 on top of bd372b6.

Thanks for putting this together — happy to re-test if you push a fix for the capture path.

[mkornreich]

@rakhimovv thanks for the H100 run and the precise trace — that pinpointed it exactly.

You're right that capture dies on the host round-trip in the per-layer-embedding path. One thing worth noting: pinning the idx staging buffer wouldn't actually be a correct fix — the PLE lookup is a host gather (F.embedding over a CPU table), so even if capture accepted the copy, graph replay wouldn't re-run the gather and the per-layer inputs would be frozen at capture-time values. The lookup has to stay on device to be graph-safe.

Pushed abd3b14:

• Default is now an on-GPU VocabParallelEmbedding for the per-layer table (graph-safe), so the default graphs-on path captures — same as what the GGUF path already did with its Q6_K GGUFEmbedding. On your H100 the ~4.7 GB table fits VRAM comfortably, so a plain ft serve --model <gemma-4-E2B-it> (no --cuda-graph-max-bs 0) should now capture graphs directly.
• Host-RAM offload of the table is now opt-in via FREETOKEN_GEMMA4_PLE_CPU=1 (for VRAM-tight boxes — it's the only way E2B fits my 8 GB card in bf16). In that mode the model reports supports_cuda_graph=False and the engine disables graphs with a logged reason instead of crashing on the default flags.

Verified on an 8 GB RTX 5070:

• FREETOKEN_GEMMA4_PLE_CPU=1 + default flags → boots (CUDA graphs disabled: ... host-resident weight gather), greedy output matches yours (The capital of France is **Paris**. ... 1. One 2. Two ...).
• Q4_0 GGUF + graphs on → captures, ~171 tok/s.

Would appreciate a re-test of the default (no-env) path on the H100. And thanks for independently confirming the double-wide split lines up with the 20 KV-shared layers.

[rakhimovv]

Re-tested abd3b14 on the H100. Both paths are good, and thanks for the correction — you're right that pinning the staging buffer would have been the wrong fix. I was reading it as a transfer problem; that replay wouldn't re-run a host gather at all is the part I missed, and it makes the on-device table the only sound answer.

Default path, no env, graphs on — captures and serves.

Start capturing CUDA graphs with sizes: [1, 2, 4]
Free GPU memory before capturing CUDA graphs: 7.96 GiB
Capturing graphs: bs = 1 | avail_mem = 7.85 GiB: 100%|##########| 3/3 [00:07<00:00, 2.65s/batch]
Free GPU memory after capturing CUDA graphs: 7.84 GiB
API server is ready to serve on 127.0.0.1:21950

Greedy output is byte-identical to the --cuda-graph-max-bs 0 run from my earlier comment:

'The capital of France is **Paris**.\n\nHere is the count from 1 to 5:\n\n1. One\n2. Two\n3. Three\n4. Four\n5. Five'

The per-layer inputs are live under replay, not frozen at capture. Since a stale PLE was the failure mode you were guarding against, I sent four unrelated greedy prompts through the captured decode graph rather than relying on the one:

prompt	answer
chemical symbol for gold	Au
17 × 3	51
"good morning" in German	Guten Morgen
largest planet	Jupiter

All correct and all prompt-specific, which a capture-time-frozen table could not produce.

Opt-in host path also behaves as described. FREETOKEN_GEMMA4_PLE_CPU=1 on default flags declines graphs with the logged reason instead of crashing, and generates the same text:

CUDA graphs disabled: the model's forward uses a host-resident weight gather that graph capture cannot record.
CUDA graph is disabled.
API server is ready to serve on 127.0.0.1:21951

Same environment as before: H100 80GB HBM3 (sm_90), driver 580.126.16, Xeon Platinum 8462Y+, Linux 5.15.0-157-generic, nvcc 13.1, torch 2.11.0+cu130, Python 3.12.13, checkpoint google/gemma-4-E2B-it at revision 3e22461, PR branch at abd3b14.

No blockers left from my side on NVIDIA hardware.