Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048

Question: is a CUDA backend for the gemma4 family wanted? I'm asking before writing any code. If the answer is "not now", that's fine and this can close.

The plan builds Gemma 4 on the shared CUDA operators from #746, #797 and #798. The extensions those operators need are general, and they land first, each as its own PR.

Why now

• Gemma 4 (31B and 12B, Gemma 4 31B served natively on Metal: the dense family, DFlash drafts, tensor-unit prompts past mlx-lm, and core affine_mm without its epilogue #879, Gemma 4 12B on the native Metal engine: text, images and its MTP assistant (stacked on #879) #975, Gemma 4 at every MLX affine width, uniform or per module: oQ packs on the native Metal engine (stacked on #975) #977) runs on Metal only. On CUDA, a gemma4 or gemma4_unified checkpoint gets "the native CUDA engine has no backend" (zig/src/native/cuda.zig:203).
• No family calls most of the shared CUDA operators yet (Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548). Gemma 4's layer is mostly what they compute already:
  • RMSNorm rows;
  • QK norm and RoPE with the cache write;
  • window attention over a ring;
  • activation rows with 64-value group sums;
  • qlinear's MLX affine paths (4-bit, and 2, 3, 5, 6 and 8 bits, in groups of 64);
  • tree_round's argmax;
  • block_attention for the DFlash drafter (head dim 128, GQA 8, window, causal block).
• RFC: an OpenAI-compatible /v1/embeddings endpoint (first model: EmbeddingGemma-2 text) #1035's EmbeddingGemma-2 family will reuse Gemma 4's layers, Metal first and then CUDA. The layer operators below would be there for it.
• I found no issue, PR, fork branch or Zig CUDA ports: shared layers below lanes.Backend, and one place to coordinate them #548 claim for Gemma on CUDA.

What the Metal family computes, and where it lands on CUDA

The reference is zig/kernels/metal/gemma4/glue.metal, including its rounding points (MLX's bf16 rounding after each op).

Most of these operators have no kernel-contract registration yet. Each step below registers the operator it extends, unless the contract work on 1.1.1's list has already done so.

The family keeps only Gemma glue:

• config admission and refusals (shared with Metal);
• the weight mapping;
• the layer order and the global-every-sixth pattern;
• the global RoPE frequency table;
• the cache layout;
• the lane and graph wiring.

Steps (one PR each, shared parts first)

1. norms: the rounding flag, the sandwich branch norm and the scalar.
2. act_rows: GELU-tanh, and the per-op rounding option.
3. embed_rows: the scale, and 2-6 and 8-bit tables.
4. qk_rows: head dim 512, and the weightless V norm.
5. window_attention: head dim 256, prompt tiles at GQA 2, and the per-row last key.
6. New shared attention rows for head dims 256 and 512.
7. New shared soft-cap rows.
8. The gemma4 CUDA family on the 12B (4-bit), text only:
   • plain decode, prompt chunks, lanes and prompt reuse;
   • images refused by name, with a test of that refusal, until step 12.
9. The 31B, and the other MLX widths (oQ8e, mixed recipes).
10. The 12B's MTP assistant. Its queries attend over the target's own caches, so it reuses steps 4-6.
11. The 31B's DFlash drafter on block_attention and tree_round.
12. Images:
    • first the 12B's patch embedder (LayerNorm, Linear, positions, projection; vit's LayerNorm needs widths up to 3,840);
    • then the 31B tower (27 layers, head dim 72 on vit's attention, 2D RoPE, 3 × 3 pooling).

Steps 1-7 each come with synthetic tests at Gemma's served shapes and tile edges, and with the CUDA SUMMARY lines. A change to an existing kernel leaves every current caller's bytes unchanged: the default path stays byte-equal to v1.1.0, and only the new flags run new code.

Gates

• Per kernel: the float64 bound at every served shape and tile edge. Current callers byte-equal to v1.1.0.
• Per layer: hidden rows, keys and values against a float64 reference of the Metal family's formula. Prompt rows and decode rows are checked separately, and the worst column is reported.
• Against Metal: teacher-forced top-token differences and margins on the same checkpoint. Bits can't match across the two backends (Metal uses fast cos and sin, and its own sum orders), so this is a fidelity check, not an equality check.
• Engine: drafted == plain, concurrent == solo, resumed == fresh, token for token.
• Speed on GB10 for each family step (bench_openai.py, bench_concurrent.py), with receipts.

Hardware

• GB10 (sm_121) for every run and receipt.
• sm_80, sm_86, sm_89 and sm_120: compile and limits records only (gate 5). I can't run them.
• Metal: unchanged.

First checkpoint: mlx-community/gemma-4-12B-it-4bit (MLX affine 4-bit, groups of 64, 6.8 GB). It needs no conversion and no 8-bit table, so step 8 doesn't wait on step 3's 8-bit part. Its patch_dense is quantized, so images on that checkpoint stay refused. For step 9, unigilby/gemma-4-12B-it-oQ8e (the recipe's checked 12B, 12.8 GB) and the 31B checkpoints follow.

Questions

• Is a Gemma 4 CUDA backend wanted now, or should it wait for the EmbeddingGemma-2 work in RFC: an OpenAI-compatible /v1/embeddings endpoint (first model: EmbeddingGemma-2 text) #1035?
• Should global attention be a new shared kernel (step 6), or window_attention with no window and a linear cache?
• Should the MLX rounding points be a flag on norms, qk_rows and act_rows, or Gemma glue in the family?

I'd start with step 1 if this order works for you.