Gemma 4 on Macs: the dense 31B, the 12B (`gemma4_unified`) with its MTP assistant, and oQ mixed widths?

Closed
  • Is something wrong?
#322 2 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python
Domain
ai-infra-agents

Research direction

Start by reading the existing gemma4 path, RowDecode, moe_tail, and the kernels named in the issue (attn_tail, dense_tail, and qwen.dense.v1.affine_rows). Check the proposed gemma4 and gemma4_unified layouts against the stated check refusals. Done means the agreed family design supports the specified dense, mixed-width, and 12B cases while preserving the exactness and multi-stream token-hash results described.

Written by the indexing model from the issue text.

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.json sets 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). RowDecode runs the dense layout with the existing attn_tail, the stacked GeGLU
    gate/up, the down projection, and one new kernel, dense_tail. That kernel is moe_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_rows for 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. check refuses, by tensor name, any width, group or
    mode no kernel reads.
  • The 12B.
    • Text only. Prompts run mlx-lm's gemma4 forward, 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_assistant head: 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.
    • check refuses the other Gemma 4 layouts this path does not read: MoE, per-layer input, shared KV.
  • 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

  1. One family or two? Dense layers inside gemma4 (the 31B), and the 12B as its own gemma4_unified family
    reusing them? Or one dense path for both inside gemma4?
  2. Checkpoints. The families' MODELS would name our published packs above. Would you rather TensorFold/
    host copies, as it does for its other models?

We have both branches ready and would open them after your answer.

Dominant language
Python
Stars
995
Forks
169
Avg merge
1d 11h
Merged PRs (30d)
425

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from ashhart/TensorFold

All issues in ashhart/TensorFold

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.