[mkornreich]

Summary

Serves Gemma-4 E2B and E4B, from both HF safetensors and GGUF, on current main. The GGUF refusal in models/gemma4/gguf.py is gone. Rebuilt on top of #706 and the GGUF refactor rather than rebased: the dense parser this PR used to carry is superseded, so what is left is the E-series delta.

E2B and E4B add three things to the gemma4 decoder: a per-layer input embedding table, KV-shared layers, and (E2B only) a double-wide MLP on those layers.

Changes

• Per-layer inputs (models/gemma4/model.py): embed_tokens_per_layer plus a per-layer projection and norm, one more residual step before layer_scalar.
• Table placement (engine/engine.py): the table defaults to the GPU. --embed-weights host now moves it to pinned host RAM, replacing the old FREETOKEN_GEMMA4_PLE_CPU env var. A table tied to the LM head stays on the GPU, and host is still refused when that is the only table. CUDA graphs stay on, which the old host path could not do.
• KV sharing (models/gemma4/attention.py): the last num_kv_shared_layers layers project q only and read k, v from the last earlier layer of their attention type. They are handed over in a dict living for one forward, which removes the old persistent stash and its fixed 8192-token size. Each shared layer still stores its own KV copy.
• MLP width (models/config.py, models/blocks.py): ModelConfig.layer_intermediate_sizes is a per-layer override, read from HF's use_double_wide_mlp or GGUF's feed_forward_length list. intermediate_size stays the base width.
• Loaders (models/gemma4/weight.py, models/gemma4/gguf.py, models/gguf/tensors.py): drop the k, v and k_norm tensors that files ship for shared layers, and accept head_count_kv and feed_forward_length as either one value or a per-layer list. The table is read without an extra copy on both paths, so loading a 5.6 GB BF16 table does not double the peak RAM, and the safetensors reader never touches VRAM for it under --embed-weights host.
• Vision (models/register.py, engine/config.py): the E-series tower applies clipped linears, which the tower does not implement. These checkpoints now serve text only with a warning, through a small EncoderSpec.unserved hook, instead of refusing to start. An image request gets a 400.

Three commits left this PR: the Llama GGUF loader (main has its own now), fp8 KV (it has its own PR, #126), and run-freetoken.sh (a local convenience).

Tested models

RTX 5070 Laptop 8 GB, driver 595.99, torch 2.14.1, transformers 5.18, Python 3.13. CUDA graphs on throughout, bs=1, 400 tokens.

File	Per-layer table	Flags	tok/s
google/gemma-4-E2B-it	bf16	--embed-weights host	67.5
E2B GGUF, Q4_0	Q6_K	default	171
E2B GGUF, Q4_0	Q6_K	--embed-weights host	169
community E4B GGUF, Q4_K_M	Q6_K	default	92
E4B Q4_K_M (the weights of Ollama's gemma4:latest, with a corrected tokenizer section)	BF16, 5.25 GiB	--embed-weights host	93

E4B with the BF16 table takes 6.7 GB of VRAM under --embed-weights host, leaving about 78k KV tokens. By the numbers the bf16 safetensors cannot fit an 8 GB card with the table on the GPU (4.4 GiB table plus 4.6 GiB of other weights against 7.3 GiB free), so I did not run that case here. Its default path was last exercised in the H100 runs earlier in this thread, on the old branch.

Results

• vs the transformers reference, same weights: E2B bf16, greedy, 8 prompts × 64 tokens, against transformers on CPU in bf16. 4/8 identical. At each of the 4 divergences FreeToken's token is the reference's second choice in bf16 (logit margin 0.03 to 0.56) and its first choice in fp32, so these are rounding flips at near-ties, not an architecture difference.
• vs llama.cpp (Ollama's build), greedy, 6 prompts × 128 tokens, equal prompt token counts: E4B Q4_K_M, same weights: 2/6 identical, the others diverge at near-ties in wording, as feat(gguf): serve Gemma 4 dense GGUF checkpoints #706 saw on the 12B and 31B. E2B Q4_0, same file: 0/6, with one divergence at the first token. The reference result above rules out the architecture, but I did not separately explain the larger gap on Q4_0.
• Graph replay: prompt-specific greedy answers through the captured decode graph on every row of the table (Au, 51, Guten Morgen, Jupiter), so the per-layer inputs are live under replay. E2B's answer to the France prompt is byte-identical to the earlier H100 run in this thread.
• Staggered requests: E2B and E4B, 48 requests at concurrency 16: 48/48 correct on both, no empty or corrupted output.
• Tokenizer (tok_compare): 0 differing texts on E2B and E4B against the HF tokenizer.json.
• Coverage (gguf_check, with the kernel table), run one base commit earlier:
  • E4B: 52/53 across unsloth, bartowski, ggml-org, lmstudio-community and google.
  • E2B: 135/142 on the stubs fetched. HF rate-limited the wider search, so this sweep is not exhaustive.
  • 12B, 26B-A4B and 31B, as a regression check: 64/64.
• FTW: ft checkpoint on the E2B safetensors converts in 8 s, loads in 20 s with --embed-weights host, and matches the safetensors' greedy output.
• Unit tests: new tests/models/test_gemma4_e_series.py (19 tests). tests/models tests/moe tests/kernels/test_ggml.py tests/server tests/engine, one base commit earlier: 1441 passed, 221 skipped, 3 failed. The same 3 fail on clean main on this machine: test_scheme_for_agrees_with_the_stored_tensors over non-LLM repos in my local HF cache. On the final commit the gemma4, loader, registry and ggml tests give 110 passed, 4 skipped, and the serving, staggered and greedy-answer checks above were repeated.

Not covered

• Image and audio input on E2B / E4B: served text only, see above. Applying clipped linears in the tower would be a follow-up.
• Zero-copy KV sharing: shared layers keep their own KV copy, a memory optimization for a follow-up.
• Refused or failing files in the sweeps:
  • UD-Q2_K_XL on both models: embedding TQ2_0: no get_rows kernel.
  • 8 ik_llama.cpp files with fork-only ggml types.
  • 5 files with no chat template.
  • 1 LoRA adapter.
  • gemma4-assistant MTP draft files (not a served architecture).
• Ollama's own blob as is: the weights load, but it has no chat template and declares a SentencePiece tokenizer, so about half of texts tokenize wrong. It takes a corrected tokenizer section to serve.
• GGUF → FTW: ft checkpoint converts a GGUF, but serving the result fails at model.embed_tokens, because source_metadata.gguf has no tensor table for GGUFConfig. That lookup is not touched here and is not specific to these models, so I left it.
• Not run: the AIME eval, TP > 1 (E-series is TP=1 only, like the rest of gemma4 GGUF), and ROCm, including the Radeon 890M that first motivated this.

One observation outside this PR: after a ggml source change the first server start compiles the GGML kernels during graph capture, with the model already loaded. That compile alone peaked at 11.6 GB of RAM here (24 jobs), which is a lot to add to a loaded model on a consumer machine.

🤖 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.

[void-mckenzie]

Hi @mkornreich, still working on this? The branch has been quiet since August and main has
moved a lot under it.

What changed: dense Gemma-4 GGUF landed in #706 (your dense parser is superseded), and the
GGUF loader was refactored around it (#690, #700, #704, #698). models/gemma4/gguf.py:220
now refuses E-series cleanly, so the rebase builds on that seam. What remains is the
E-series delta: PLE, shared KV layers, double-wide MLP, per-layer FFN width.

Pointers: per-layer FFN width wants a per-layer override, not a widened intermediate_size
(#700's slot pools are the precedent for lists); default the PLE table on-GPU as your
abd3b14 does and make the host-RAM path a flag rather than an env var (it's the only way
E2B fits small cards, so worth keeping); the KV stash is fine to land, zero-copy sharing
can be a follow-up. Verification: scripts/gguf_check sweep, tok_compare, a real-file
serving run.

If you don't have time, I'm happy to pick it up with co-author credit to you and finish the
GGUF side on current main. Thanks!

[mkornreich]

@void-mckenzie I'll check it soon

[mkornreich]

@void-mckenzie Rebuilt on current main, on top of #706, with your review items addressed: the per-layer MLP width is now a per-layer override, the PLE table defaults to the GPU with host RAM behind --embed-weights host (graphs stay on), and the KV stash stays, without its fixed-size buffer. The gguf_check sweep, tok_compare and real-file serving results are in the description, along with what's not covered. Ready for review.