Gemma 4 E4B-it: transformers MPS vs llama.cpp Metal on M4 Max

bench_cpu.py

import time
import torch
from transformers import AutoProcessor, AutoModelForCausalLM

TARGET_MODEL_ID = "google/gemma-4-E4B-it"
MAX_NEW_TOKENS = 256
DEVICE = "cpu"

print(f"Using device: {DEVICE} | threads={torch.get_num_threads()}")

processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)
target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID, dtype=torch.bfloat16,
).to(DEVICE)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a detailed 300-word explanation of how virtual memory and paging work in modern operating systems. Cover TLBs, page faults, and swap."},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=text, return_tensors="pt").to(DEVICE)
input_len = inputs["input_ids"].shape[-1]

# warmup
_ = target_model.generate(**inputs, max_new_tokens=4, do_sample=False)

torch.manual_seed(0)
t0 = time.perf_counter()
out = target_model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
dt = time.perf_counter() - t0
new_tokens = out.shape[-1] - input_len
print(f"\ntransformers CPU bf16 E4B: {new_tokens} tok in {dt:.2f}s = {new_tokens/dt:.2f} tok/s")
print("---")
print(processor.decode(out[0][input_len:], skip_special_tokens=True)[:300])

gemma4-bench.md

Gemma 4 E-series — transformers vs llama.cpp on Apple M4 Max

Quick benchmark prompted by the question "how can E4B-it be only ~20 tok/s on MPS?"

Setup

• Hardware: Apple M4 Max, 128 GB unified memory
• Prompt: 52 tokens (system + 1 user turn asking for a 300-word explainer on virtual memory / TLBs / page faults / swap)
• Generation: greedy, max_new_tokens=256, thinking disabled
• transformers 5.8.0 / torch 2.11.0 / llama.cpp (homebrew, Metal)

Results

stack	model	precision	gen tok/s	prefill tok/s
transformers (CPU)	gemma-4-E4B-it	bf16	8.70	—
transformers (MPS)	gemma-4-E4B-it	bf16	20.99	—
llama.cpp (Metal)	gemma-4-E4B-it	bf16	21.52	499
transformers (MPS)	gemma-4-E2B-it	bf16	26.77	—
llama.cpp (Metal)	gemma-4-26B-A4B-it	Q4_K_XL	41–44	300–345

Speculative decoding with the official MTP drafter (gemma-4-E{2,4}B-it-assistant) was slower on MPS in both cases (~15–19% regression) — verification overhead and Metal kernel-launch cost outweigh the drafting win at this model size. The README's "up to 2x" assumes a discrete GPU.

Takeaways

1. At identical precision and identical model, transformers MPS ≈ llama.cpp Metal for token generation (20.99 vs 21.52 tok/s). The HF stack is not the bottleneck on Apple Silicon for decode.
2. Quantization, not the framework, is the lever. The earlier ~2x gap I observed was llama.cpp running a Q4 model vs transformers running bf16. A Q4/Q8 E4B GGUF should land in the 50–80 tok/s range on this Mac.
3. Prefill is where llama.cpp wins — ~500 tok/s vs (much slower) HF MPS. Matters a lot on long prompts; doesn't matter on short ones.
4. MTP drafters need a real GPU. The drafter is ~78M params, but each verification step on Metal pays a non-trivial overhead that erases the speedup at E2B/E4B target size. They'd shine on the 26B/31B targets on CUDA.
5. CPU bf16 is ~2.4× slower than MPS bf16 for the same model — fine for offline batch but not interactive.

Gotchas hit along the way

• transformers chat/transformers serve requires transformers[serving] extras (openai, fastapi, uvicorn); without them the CLI errors on import. There is no --assistant-model / speculative-decoding flag in transformers serve.
• Gemma4Processor and Gemma4VideoProcessor require pillow and torchvision even for text-only — install both up front.
• The hf CLI download (xet client) downloaded the BF16 GGUF at 2.1 MB/s. A direct curl -L from the same CDN URL got 26 MB/s — ~12× faster. The curl progress bar's 100.0% was misleading; the file was actually ~14 MB short and llama.cpp errored with tensor 'per_layer_model_proj.weight' data is not within the file bounds. curl -C - resumed cleanly.
• device_map="auto" on a Mac with accelerate is fiddly; explicit .to("mps") was simpler and worked first try.

Reproduce

uv init --python 3.12
uv add transformers torch torchvision accelerate pillow librosa
# bench script: see main.py / bench_cpu.py in the gist files

# llama.cpp side
brew install llama.cpp
curl -L -C - -o ~/models/gemma-4-E4B-it-BF16.gguf \
  https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF/resolve/main/gemma-4-E4B-it-BF16.gguf
llama-server -m ~/models/gemma-4-E4B-it-BF16.gguf -c 4096 -ngl 999 --port 8080

main.py

import time
import torch
from transformers import AutoProcessor, AutoModelForCausalLM

TARGET_MODEL_ID = "google/gemma-4-E4B-it"
ASSISTANT_MODEL_ID = "google/gemma-4-E4B-it-assistant"
MAX_NEW_TOKENS = 256

if torch.backends.mps.is_available():
    device = "mps"
elif torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
print(f"Using device: {device}")

processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)

target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID, dtype=torch.bfloat16,
).to(device)

assistant_model = AutoModelForCausalLM.from_pretrained(
    ASSISTANT_MODEL_ID, dtype=torch.bfloat16,
).to(device)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a detailed 300-word explanation of how virtual memory and paging work in modern operating systems. Cover TLBs, page faults, and swap."},
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=text, return_tensors="pt").to(target_model.device)
input_len = inputs["input_ids"].shape[-1]


def sync():
    if device == "mps":
        torch.mps.synchronize()
    elif device == "cuda":
        torch.cuda.synchronize()


def bench(label, **gen_kwargs):
    # warmup
    _ = target_model.generate(**inputs, max_new_tokens=8, do_sample=False, **gen_kwargs)
    sync()
    torch.manual_seed(0)
    t0 = time.perf_counter()
    out = target_model.generate(
        **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, **gen_kwargs,
    )
    sync()
    dt = time.perf_counter() - t0
    new_tokens = out.shape[-1] - input_len
    print(f"{label:24s}  {new_tokens:4d} tok in {dt:6.2f}s = {new_tokens/dt:6.2f} tok/s")
    return out


print(f"\nPrompt tokens: {input_len} | max_new_tokens: {MAX_NEW_TOKENS} (greedy)\n")
out_plain = bench("target only")
out_spec  = bench("target + assistant", assistant_model=assistant_model)

print("\n--- output (target + assistant) ---")
print(processor.decode(out_spec[0][input_len:], skip_special_tokens=True))