safemlx: gemma-4 E4B decode ~1.5 tok/s vs ~36 tok/s for mlx-lm on the same M4 — per-token cost scales with prompt length

[texchi2]

Thanks for the crate — the vendored MLX with checksums and the offline build are exactly what I wanted from a Rust MLX stack, and safemlx-lm built cleanly first try on two accounts (3 m 3x s each) with Xcode 26.6 / Metal toolchain 17F109.

I measured it against mlx-lm on the same model, same snapshot, same machine, and the decode path looks like it is recomputing the whole context every step. Reporting in case it is a known gap for gemma-4, or a misuse on my side.

Environment

• Mac mini M4, 16 GB, macOS 26.5.2, Xcode 26.6 (clang 21.0.0), Metal toolchain 17F109 (metal version 32023.883), cmake 4.4.3, cargo 1.98.0
• safemlx-sys 0.2.3 / safemlx 0.3.1 / safemlx-lm 0.4.1 from crates.io, cargo build --release --locked -j 4, default features
• Model: mlx-community/gemma-4-e4b-it-4bit, snapshot 475b9088d29754a3379866cf5aeb6b41acd313c2
• Comparison: mlx-lm 0.31.3 on mlx 0.31.2, same snapshot, python -m mlx_lm.generate --temp 0 --max-tokens 128
• Both runs: greedy, 128 generated tokens, two other idle model servers resident (~4.6 GB free); pausing them changes wall time by ~2 s, so this is not memory pressure

safemlx-lm — examples/qwen35_moe_bench.rs pointed at the gemma-4 snapshot (it loads through LoadedModel::load; profiling off)

$ ./target/release/examples/qwen35_moe_bench <snapshot> 128 prefill_128
load_s=2.512
warmup_s=0.469    case,prompt_tokens,prefill_s,generated_tokens,decode_s,decode_tok_s,first_id,last_id 

prefill_128,152,0.489611,128,85.149649,1.491,107,106

$ ./target/release/examples/qwen35_moe_bench <snapshot> 128 prefill_512
load_s=2.239
warmup_s=0.458
prefill_512,529,1.665251,128,241.580618,0.526,107,106

So: prefill ~310 tok/s (fine), but decode 1.49 tok/s at a 152-token prompt and 0.53 tok/s at 529 tokens. Per generated token that is 0.66 s and 1.89 s — almost exactly the measured prefill time of the whole prompt in each case, which is what made me suspect the cache is not being extended for this architecture. The gemma4_e4b_probe example shows the same rate (90 tokens in ~27 s including a
2.5 s load), and it uses the same generate_input_with_cache path.

mlx-lm on the same snapshot

Prompt: 160 tokens, 253.398 tokens-per-sec
Generation: 128 tokens, 36.350 tokens-per-sec
Peak memory: 4.448 GB

36 tok/s is ~110 GB/s of weight traffic, close to the M4's 120 GB/s bus, so mlx-lm is at the bandwidth wall; safemlx-lm is at ~4 % of it.

Questions

1. Is incremental KV-cache decode expected to work for the gemma-4 loader in 0.4.1, or is the per-step recompute known?
2. Is there a flag or a different generation entry point I should be using for gemma-4 (the bench example is the Qwen3.5-MoE one; I used it only because it is the one that prints prefill/decode separately)?

Happy to re-run any diagnostic build on this machine and post the numbers.

[texchi2]

gemma-4-12B (gemma4_unified) decodes correctly for ~3 tokens then degenerates; prefill matches the reference exactly

Follow-up to #4, on eredu 0.5.0 (workspace at df0af94) rather than the old safemlx-lm crate. Two mlx-community gemma-4 checkpoints, same machine, same binary:

• mlx-community/gemma-4-e4b-it-4bit (model_type: gemma4) — loads and generates coherently, ~39 tok/s.
• mlx-community/gemma-4-12B-it-4bit (model_type: gemma4_unified) — loads, answers short prompts, but longer
generations degenerate after a handful of tokens.

(Both needed a small loader patch first — the mlx-vlm physical layout language_model.model.* and the stub vision_config / audio_config of a text-only conversion. I can open that as a separate PR if useful; it is unrelated to the decode behaviour below.)

Environment

• Mac mini M4, 16 GB, macOS 26.5.2, Xcode 26.6, cargo 1.98.0, cargo build --release --locked -j4 -p eredu-cli
• Reference for comparison: mlx-vlm 0.6.15 on mlx 0.32.2, same snapshots, greedy, step-by-step with its own cache

What degenerates

$ eredu --model mlx-community/gemma-4-12B-it-4bit -t 0 -n 28 \
    "Write a short paragraph about how cotton denim is woven."
Cotton denim is-produced using a **twill weave**, specifically aas-typeaaee-one-over-onee-under-twill.
In thise-tw. no-_.. no-. no-. no-. no-.

mlx-vlm on the same checkpoint and prompt:

Cotton denim is created through a specific weaving process known as **twill weaving**,

The two agree on tokens 0-2 (Cotton, denim, is) and diverge at token 3.

What is NOT wrong (measured)

1. Prefill is exact. First generated token vs the mlx-vlm reference over six prompts: 6/6 identical, including one where the reference's top-1/top-2 gap is only 2.62 logits. So the forward pass, the tokenizer, the chat template and the sampler agree on a fresh prompt.
2. Not quantisation noise. gemma-4-12B-it-6bit degenerates the same way
   (Cotton denim is primarily-produced using a twill weave, specifically a technique characterized byarer and/or/s/s/).
3. Not the planner. --no-auto produces byte-identical output, and --auto plan reports
   residency.mode = fully_resident with no compressed KV for both checkpoints.
4. Not context loss. A code word placed 134 words before the question is recalled correctly.
5. E4B decodes correctly through the same code path.

Where that leaves it
The divergence appears only in decode, only on the unified 12B. The geometry that checkpoint exercises and E4B never does is the full-attention layers with attention_k_eq_v: true and num_global_key_value_heads: 1 (16 query heads : 1 KV head, global_head_dim 512, values taken from the un-roped k_proj output and normalised without scale). Prefill computes those values directly; decode re-reads them from the cache — which fits a
decode-only, 12B-only divergence — but I have not proved that, and I would rather report the measurements than guess the mechanism.

Happy to run a per-layer activation dump on this machine if you can point me at the right capture plan, or to test a patch.

[texchi2]

The loader side of this is fixed on a branch, if it is useful to you:

https://github.com/jbg/eredu/compare/main...texchi2:mlx-community-gemma4?expand=1
https://github.com/jbg/eredu/compare/main...texchi2:mlx-community-gemma4?expand=1 (2 commits, 4 files, +261/-27; main merges cleanly)

• mlx-community / mlx-vlm SafeTensors conversions nest the decoder under language_model.model.* and keep the media towers and projections at the root without model.; every Gemma 4 constraint now also carries that physical alias, so one plan admits the released, canonical and mlx layouts.
• Text-only conversions keep vision_config / audio_config as stubs without encoder geometry; such a stub is treated as an absent tower and its placeholder ids are cleared (a placeholder with no media config at all is still rejected, as before). The orphan embed_vision / embed_audio / vision_embedder projections are admitted when the tower is absent.
• Smaller: quantised media-to-text projections take the decoder's format; the audio processor accepts the *_token_index spelling of the framing ids (some configs carry both spellings, so it is a fallback, not a serde alias); assistant_safetensors_plan selects the neutral alias explicitly instead of aliases.first(), since plan normalisation sorts aliases.

With it, on the M4 from above: gemma-4-e4b-it-4bit loads and decodes coherently at ~39 tok/s (mlx-lm 36.4 on the same box), and gemma-4-12B-it-4bit loads and answers short prompts — the remaining 12B divergence is the comment above, untouched by the branch.

cargo test --locked -p eredu-architectures --lib: 535 passed (two new tests); rustfmt --check clean on the four files.

Note the reference_realtime_construction test target does not compile on current main for me (PreparedMoshiRealtimeSource::from_reference_source missing at df0af94), so --all-targets was not usable — pre-existing, unrelated.

The repository currently only lets collaborators open pull requests, so I could not open one — happy to if that is lifted, or feel free to cherry-pick the two commits.

[texchi2]

Found the decode divergence I reported above, and it is a one-line policy choice rather than anything numeric.

gemma4/graph.rs** maps AttentionValueSource::ReuseKey (i.e. attention_k_eq_v) to a key-only cache**, and ConcatKeyValueCache::update_and_fetch returns the stored keys in place of the values. But Gemma 4's k_eq_v reuses the key projection, not the cached key tensor: the value branch takes the same k_proj output, normalizes it without a learnable scale (RMSNormNoScale) and is never rotated, while the cached keys carry k_norm and RoPE. So every full-attention layer of a k_eq_v model attends over rotated keys instead of its values.

How it was localized. Model mlx-community/gemma-4-12B-it-4bit (48 layers, full attention every sixth layer, one global KV head, global_head_dim 512); reference mlx-vlm 0.6.15; same prompt, same 24 tokens, greedy; relative L2 at the last position:

• per-layer hidden states, layers 0-4 (sliding): about 6e-3, i.e. bf16 noise
• layer 5, the first full-attention layer: 3.647e-01, and a spike at every sixth layer after it
• inside layer 5, the tensors entering the cache: q after RoPE 8.7e-3, k 1.0e-2, v 1.0e-2 — all fine
• layer 5 attention output: rel 9.4e-1, with norm 326.75 for the reference against 19.77 here, a factor of 16.5
• swapping tensors between the two engines: only v moves the result
• v_all against k_all as returned by the cache: rel 0.000e+00, identical — and the k/v norm ratio is 16.5, the same factor

Fix: store both tensors for every state-owning layer. Branch (on top of the loader work already linked in this issue): https://github.com/texchi2/eredu/compare/main...texchi2:mlx-community-gemma4?expand=1 https://github.com/texchi2/eredu/compare/main...texchi2:mlx-community-gemma4?expand=1 — the last commit is the one-file change; the upstream test that asserted KeyOnly for that layer is corrected in the same commit, since it encoded the same assumption. Cost is one value tensor per k_eq_v layer (here 1 KV head × 512).

After the fix, same prompt and comparison: layer 5 3.647e-01 → 1.337e-02, worst of all 48 layers 1.7e-2, mean 8.2e-3, and the greedy continuation matches mlx-vlm token for token — Cotton denim is created through a specific weaving process known as twill weaving,.

Timings on the M4 (16 GB, 128 tokens greedy, wall time including load): gemma-4-12B-it-4bit 17.8 s / 7.0 GB peak · gemma-4-12B-it-qat-4bit 28.4 s / 7.3 GB · gemma-4-e4b-it-4bit 7.3 s / 6.1 GB — all three coherent. cargo test --locked -p eredu-architectures --lib: 536 passed.

(The QAT checkpoint needed one more fix to load at all — its quantization block carries 144 nested per-tensor overrides, 8-bit on every MLP projection, which the config parser dropped silently. That is the commit before this one on the same branch.)