Repository navigation
Gemma 4 on Macs: the dense 31B, the 12B (gemma4_unified) with its MTP assistant, and oQ mixed widths? #322
Description
The question
The gemma4 family serves the 26B-A4B MoE checkpoints at uniform 4-bit with an 8-bit router. TensorFold refuses
three other Gemma 4 cases today:
- the dense 31B (
model_type: gemma4,enable_moe_block: false); - the 12B, which ships as
Gemma4UnifiedForConditionalGeneration(model_type: gemma4_unified) with a separate
MTP assistant (gemma4_unified_assistant); - every oMLX oQ/oQe checkpoint of any size, whose
config.jsonsets widths per layer.
Would you take Mac support for these, under the exactness contract? And should the dense paths live inside gemma4,
or should the 12B be its own family?
What it would take
- Dense layers (31B and 12B).
RowDecoderuns the dense layout with the existingattn_tail, the stacked GeGLU
gate/up, the down projection, and one new kernel,dense_tail. That kernel ismoe_tail's arithmetic with the MLP
output in place of the joined branches: post-feedforward norm, residual, layer scalar, and the next layer's input
norm. - Mixed widths. Projections read every MLX affine width the existing kernels read (2–8 bits in groups of
32/64/128): the 4-bit row matvec,qwen.dense.v1.affine_rowsfor other widths, and the lane matmul where it reads
the width. A stack whose members differ in width (oQ boosts one of q/k/v or gate/up) runs one matmul per run of
equal widths and joins the outputs. Every kernel is row-invariant, so windows, shared rounds and one-row steps keep
the same bits. The router may be any width; experts stay 4-bit.checkrefuses, by tensor name, any width, group or
mode no kernel reads. - The 12B.
- Text only. Prompts run mlx-lm's
gemma4forward, with the patch embedder and audio projection dropped at
load. - Decode runs Gemma 4's v1 attention, RoPE and caches with the dense tail above.
- Drafts come from the
gemma4_unified_assistanthead: four layers attending over the target's last sliding and
last full layer's keys from the last kept row, batched across streams in shared rounds. Drafts only propose, so
the assistant runs MLX's stock kernels. checkrefuses the other Gemma 4 layouts this path does not read: MoE, per-layer input, shared KV.
- Text only. Prompts run mlx-lm's
- Memory. On the rows backend the stacked q/k/v and gate/up become views into mlx_lm's modules, so a stack costs
no second copy.
What it does not change
The uniform 4-bit 26B-A4B checkpoint keeps its kernels: a 4-bit run still takes the fused q/k/v row kernel.
Measured so far (our own scripts, not yet the upstream tools)
M5 Ultra, MLX 0.32.3, greedy, 512-token replies. In every case, drafted replies equal "draft": false, and two- and
four-stream replies equal serial ones, by token hash.
checkpoint (ours, oMLX oq, groups of 64) |
single stream | aggregate | prompt |
|---|---|---|---|
| 26B-A4B oQ4e (4-bit base with 5/6/8-bit layers, experts 4-bit) | 177 tok/s undrafted; 187–193 with z-lab/gemma-4-26B-A4B-it-DFlash |
— | — |
| 31B oQ8e (uniform 8-bit) | 40 tok/s with a DFlash drafter | 83 at four streams (lane backend) | — |
12B oQ8e + mlx-community/gemma-4-12B-it-qat-assistant-4bit |
131.7 tok/s drafted, 52.0 undrafted | 238.7 at four streams | about 2.9–3.2K tok/s |
The three checkpoints are published as unigilby/gemma-4-26B-A4B-it-qat-oQ4e, unigilby/gemma-4-31B-it-qat-oQ8e
and unigilby/gemma-4-12B-it-oQ8e. They use oMLX oq (enhanced) with groups of 64, and one width per q/k/v and
gate/up stack. Other oQ/oQe packs of Gemma 4, whose stacks mix widths, load too: the mixed-width runs above read them.
Questions
- One family or two? Dense layers inside
gemma4(the 31B), and the 12B as its owngemma4_unifiedfamily
reusing them? Or one dense path for both insidegemma4? - Checkpoints. The families'
MODELSwould name our published packs above. Would you ratherTensorFold/
host copies, as it does for its other models?
We have both branches ready and would open them after your answer.
Activity
ashhart commentedon Oct 6, 2026
Thank you, @gilby. Asking first, and #360 and #361, made the Gemma 4 plan concrete. The Python engine is frozen (#286), so the dense 31B, the 12B and the oQ widths won't land in the Python
gemma4family. New families go to the Zig engine now, and Gemma 4 isn't ported there yet. Your two questions, one family or two and where the packs live, we'll answer when Gemma 4 is picked up for the Zig engine. Your branches, with the row-exact dense tail and the mixed-width runs, are where that port starts, with your credit.ashhart commentedon Oct 7, 2026
Thank you, @gilby. Gemma 4 on Macs, the dense 31B, the 12B with its MTP assistant, and the mixed oQ widths, is a family of its own. The Zig tree does not have it.
The port is Metal, on the shared lane core. Dense 31B first, greedy decode matching the Python reference, then the 12B and its MTP head only after plain decode matches. Mixed oQ widths are a loader question and they come with a fixture that fails when a width is guessed. We will not add a private matmul for this family. Thank you for listing the three checkpoints as one request so the order is obvious.
BobClawblaw commentedon Oct 14, 2026
@ashhart @gilby I'm starting the native Gemma 4 port in the order set above: the dense 31B on Metal on the shared lane core first, then the 12B and its MTP head, then mixed oQ widths with a width fixture.
mlx-community/gemma-4-31B-it-4bit(696d436c), uniform affine 4-bit in groups of 64.gemma4-dense(a6ff795e) with lane kernels, on mlx 0.32.3 / mlx_lm 0.32.0. Decode rows have to match its RowDecode bit for bit, and the prompt path has to meet the usual prompt criteria against mlx_lm's forward.core/lane_projection. The new shared pieces go in core: a per-head norm with RoPE from a frequency table (weighted or unweighted, head dim up to 512), sliding-window decode attention at head dim 512, GeGLU, and a softcap on the argmax and sampler.attn_tail,dense_tailand ring cache carry over with their credit.@gilby, if you've already started a Zig port, say so and I'll stop and pick up whatever's useful from yours instead. The PR goes up once greedy decode, the prompt path, streams and copy drafts all pass, with receipts.
gilby commentedon Oct 14, 2026
@BobClawblaw I don't have the resources available right now to do it, so go ahead. I did make a bunch of local changes (on the python branch) that improves the Gemma models along with many other things, so I will push that to my fork soon (though my agent said it's multi-day work to this, which usually means a couple hours).
I've started the consolidation. It's a large job and will run in stages, keeping a checkpoint in TENSORFOLD-FORK-CONSOLIDATION.md.
The plan:
Inventory: collect our commits from every node tree, read-only:
It'll map each commit to a topic and its dependencies, and dedupe changes that exist in several trees.
Rebase onto v0.6.6 and squash into clean, one-change PRs that follow upstream's contributing rules:
Third-party work (perf(exl3 cuda): EXL3 prompts and multi-stream verify for the 27B on CUDA, same bits per commit #260, Mac prefill: keep the chunk start before a long history's boundary (--tail-checkpoint-tokens) #287, our port of Qwen3.8-27B CUDA: --kv-dtype int8 (packed -cq 8 rows; drafted rounds -6% at 90k, -13% at 242k) #247, DeepSeek-V4.1-Flash CUDA family: TP=2 multi-stream serving on 2x DGX Spark #342) keeps its original authors, with credit in the PR body; our adaptations go in separate commits. Our --name-priority port is dropped, since 0.6.6 already has it.
gilby commentedon Oct 14, 2026
@BobClawblaw My fork is now updated: https://github.com/gilby/TensorFold
ashhart commentedon Oct 16, 2026
Thank you, @gilby, for publishing the updated fork and handing the native Gemma work to @BobClawblaw.
#879 now carries the dense 31B Metal implementation through the shared projections and lane core. We checked its loader and serving path; its declared vision tower is still to be connected, and startup cleanup and source attribution need follow-up. The 12B assistant and mixed widths remain later pieces of the plan.
Please point Bob to the commits for the dense tail and ring cache in your updated fork, so the reference and the work carried across have exact pins and your credit.
Posted by Ash's AI Assistant Synthena
gilby commentedon Oct 16, 2026
@BobClawblaw, here are the exact pins in my fork (gilby/TensorFold, Python engine,
mainat tagfleet-v0.6.6-1=3602c9a770f3915422841b659814bfe984549e98). Each change landed as its own PR on the fork:Dense tail (31B and 12B dense layers)
bee717922ab625801fce36adcc6bd00c1ce20562, gilby/TensorFold#6: the dense 31B layout and oQ per-layer mixed widths.dense_tailis insrc/tensorfold/kernels/gemma/v1/glue.py, and the dense branch ofRowDecodeis insrc/tensorfold/kernels/gemma/v1/decode.py.dense_tailismoe_tail's arithmetic with the MLP output in place of the joined branches: post-feedforward norm, residual, layer scalar and the next layer's input norm. The tests aretests/test_gemma4_kernels.pyandtests/test_gemma4_family.py, with the tiny dense fixture intests/gemma4_tiny.py.KV caches
RingKVCache, position p at slot p % ring, rollback by movingoffset) is upstream's, from99efdf4d8a5cb84ff8f613822f09e9508f7eafeb. I didn't change it.ce0f552a13310df5cee39da1d891fd419519075d, gilby/TensorFold#37: the full-attention layers' cache (LinearKVCacheinsrc/tensorfold/families/gemma4/cache.py) is written in place by a Metal kernel. This replaces the alternating spare buffer and halves those layers' cache memory. It's insrc/tensorfold/engine/inplace_kv.py, withtests/test_inplace_kv.py. The ring layers are unchanged.Later pieces of the plan, for reference
6b0ead7ac17ba13d4689fea03c824360cffdf6da, #7: Gemma 4 12B (gemma4_unified), dense row-exact decode at any oQ width, MTP assistant drafts.1bceacf70ef89fbcfcbca5982b733a448cebb62c, #38: the 26B and 31B take a Gemma 4 MTP assistant (gemma4_unified_assistant) as their drafter.0bd40336ea423286e7b2669f8facbab1535fa199, #34: before M5, each decode linear picks the matrix or the row kernels by its shape (feat(gemma4): the dense 31B layout and oQ per-layer mixed widths on the Mac decode kernels #360's matrix backend applied to Gemma 4).All of these are on the Python engine, rebased onto
python-0.6(v0.6.6). The receipts are in each fork PR's body.ashhart commentedon Oct 17, 2026
Thank you, @gilby, the exact dense-tail, in-place KV and assistant pins make the fork's contributions separable, including the unchanged upstream ring.
We checked the native release scope and the new proposals: #879 covers the 31B path, and #975 now proposes the 12B with its assistant. Your Python fork remains useful behavior and provenance evidence for those native adapters.
We will compare those arithmetic/cache choices with the native implementations and preserve your credit, keeping the shared-math and image/draft gates visible. The frozen Python serving line will not take new merges. The pins you supplied answer the requested history question; no replacement measurements are needed merely to restate it.
Posted by Ash's AI Assistant Synthena