RESMP-DEV/gemma-4-E4B-it-qat-mlx-megakernel

Gemma 4 E4B & E2B — single-launch Metal megakernels + lossless speculative decoding (MLX / Apple Silicon)

This repo is code, not weights. Two self-contained MLX implementations that run the full Gemma 4 text decode as one custom Metal kernel launch per token, plus a lossless speculative-decoding path driven by Gemma 4's own MTP draft head. Bring your own MLX 4-bit weights (see below).

file	model	architecture handled
gemma4_spec.py	E4B (42 layers, hidden 2560)	GQA 8:2, 5:1 sliding:global, 18 shared-KV layers
gemma4_spec_e2b.py	E2B (35 layers, hidden 1536)	MQA 8:1 (backbone and draft head), 4:1 pattern, 20 shared-KV layers with double-wide (12288) MLPs

The two files are generated from the same source; only the architecture block differs. All arch-dependent kernel code (GQA head mapping, stage job counts, per-layer MLP width, layer pattern, KV staging offsets) is derived from that block.

Everything here was measured on an Apple M4 Max, not extrapolated. The speculative path is bit-exact to greedy decoding (288/288 across prompts, 640/640 across the sliding-window wrap, deterministic) — it only changes speed, never output.

What it does

Two Metal kernels over the whole text stack, written with mx.fast.metal_kernel:

• decode (M=1): one token per launch. All 42 layers, attention, GeGLU MLP, and per-layer embeddings run inside a single kernel; threadgroups are kept resident and synchronized with device-scope atomic grid barriers, so there is no per-layer launch overhead and no round trip to Python between layers.
• verify (M=k+1): checks a whole draft in one launch. Every 4-bit matmul becomes an M-column GEMM (qmm4), so the quantized weights stream from DRAM once to score all k+1 candidate tokens. Because batch-1 decode is bandwidth-bound, verifying k drafts costs almost the same DRAM traffic as generating one token.
• draft (fused): all k MTP draft steps run as one Metal launch too — in-kernel backbone embedding dequant, the 4 cross-attending draft layers, a 4-bit 262k-vocab LM head with in-kernel argmax, and the surrogate-hidden projection chaining step j into step j+1. Validated 60/60 draft tokens against the reference implementation on live decode states.
• commit: candidate KV is staged in scratch during verification and only accepted rows are committed to the KV cache by a tiny third kernel. Rejected drafts never touch the cache — at wrapped sliding-window context their ring slots alias live window entries, so eager writes would silently corrupt attention (found by adversarial review, confirmed, fixed, and regression-tested across the 512-token wrap).

Drafts come from Gemma 4's MTP assistant head (google/gemma-4-E4B-it-*-assistant): a 4-layer, hidden-256 stack whose queries cross-attend directly into the megakernel's own KV cache — no separate draft cache is allocated.

Measured performance (M4 Max, QAT 4-bit, group size 64)

model	greedy decode	speculative (k=2), lossless	speedup
E4B	64–68 tok/s	84–107 (avg ~94)	~1.4×–1.5×
E2B	~105 tok/s	132–153 (avg ~144)	~1.4×

E2B lossless verification: batch verify is bitwise-exact vs sequential (max abs diff 0.0), 288/288 short-context and 640/640 long-context (across the 512 sliding-window wrap) tokens identical to greedy, 50/50 fused-vs-reference draft agreement.

Losslessness is literal and tested: 288/288 tokens identical to greedy across three prompts and 640/640 on a long generation crossing the 512-token sliding-window wrap, deterministic across runs. This required making it true by construction: the batch verify GEMM is arithmetic-identical to the sequential path, the greedy reference runs through the same compiled kernel (Metal fast-math reassociates float ops differently per kernel compilation — two "identical" kernels disagree at ~1e-4 and near-ties eventually flip), and all host-side pre/post ops run at a fixed batch shape (MLX dispatches different matmul kernels for M=1 vs M=k+1, which also flips near-ties).

Kernel-level: the M=5 verify kernel costs ~1.5× the M=1 decode kernel — five tokens scored for the price of ~1.5, the weight-streaming win the design is built around. The end-to-end multiplier is set by draft acceptance (≈0.6–0.9 accepted per draft), which is why k=2 is the default; larger k over-drafts for this head.

Usage

Requirements: an Apple-silicon Mac, mlx>=0.32, mlx-lm, and transformers<5 (mlx-lm 0.31.x does not import against transformers 5.x).

Bring MLX 4-bit weights at group size 64 (the kernel's dequant path assumes it). The quantization-aware-trained checkpoint is recommended — it survives 4-bit noticeably better than a naive post-training quant:

# convert Google's QAT-unquantized release to MLX 4-bit, group size 64
# (for E2B swap E4B -> E2B everywhere and use gemma4_spec_e2b.py)
python -m mlx_lm.convert \
  --hf-path google/gemma-4-E4B-it-qat-q4_0-unquantized \
  --mlx-path gemma-4-e4b-qat-mlx-4bit -q --q-group-size 64 --q-bits 4

# grab the MTP draft head (bf16; its LM head is 4-bit-quantized at load for drafting)
huggingface-cli download google/gemma-4-E4B-it-qat-q4_0-unquantized-assistant \
  --local-dir gemma-4-e4b-assistant

# lossless speculative decode
python gemma4_spec.py gemma-4-e4b-qat-mlx-4bit \
  --assistant gemma-4-e4b-assistant --k 2 \
  --prompt "Explain PCIe lane bifurcation in one paragraph."

from gemma4_spec import SpecDecoder
dec = SpecDecoder("gemma-4-e4b-qat-mlx-4bit", "gemma-4-e4b-assistant", k=2)
tokens, stats = dec.generate(prompt_ids, max_new=256)   # identical to greedy, faster

Implementation notes (probed on M4 Max — the silicon, not the docs)

• 31 buffer-binding cap on a single metal_kernel: all layer weights are packed into arena buffers addressed by a u32 offset table.
• Grid barriers across resident threadgroups need explicit device-scope atomic_thread_fences on both sides of the arrive/spin; MSL atomics are relaxed-only and the naive version fails at 2 threadgroups.
• Residency cliff: persistent spin barriers require every threadgroup co-resident. 80×256 is the tuned launch shape; oversubscribing collapses into scheduler preemption.
• Cold-DRAM GEMV uses 4 output rows per simdgroup for memory-level parallelism; the verify GEMM amortizes each weight read across M activation columns. The verify inner loop is kept arithmetic-identical to the sequential GEMV (same lane mapping, same fma grouping) so batch and sequential results match bitwise; wider vector loads are reserved for paths without an exactness contract (the drafter's LM head).
• The megakernel mutates its KV cache in place, which MLX does not track as a graph dependency — each launch is forced to complete before the next reads the cache.
• Speculative KV correctness: candidates are staged, never eagerly written; only accepted rows are committed. Eager candidate writes corrupt the sliding-window ring once P + k >= 512 (in-batch successor positions alias the oldest window slots of earlier batch positions).
• Bit-exactness across "identical" kernels is not real under fast math: guarantee equality by using the same compiled kernel and fixed batch shapes, not by writing the same arithmetic twice.

Provenance and license

Derived from google/gemma-4-E4B-it; drafting uses google/gemma-4-E4B-it-qat-q4_0-unquantized-assistant. Model weights are governed by the Gemma Terms of Use — this repo ships no weights; you download them from Google under those terms. The code in this repository is released under Apache-2.0.