Gemma’s Drafter Trick: 3× Inference Without Model Surgery

Speculative decoding uses two models: a small drafter that proposes an n‑gram and a heavy target model that verifies those proposals. The drafter generates a short token chunk cheaply. The target model runs one forward pass over that chunk and computes logits for every proposed position in parallel. The gain is pure amortization. Big decoder steps are memory‑bound: moving KV tensors and layer weights in and out of GPU memory dominates latency. Verifying k tokens in one pass reuses the same KV/attention work across those k positions and avoids k‑1 repeated memory fetches. Google’s MTP adds three practical engineering moves to make this win real. First, the drafter shares the target model’s input embedding table so tokenization and embedding lookups don’t duplicate memory. Second, the drafter conditions on the target model’s last‑layer activations: those activations are concatenated with token embeddings and down‑projected into the drafter’s dimension so the drafter uses already‑computed context without recomputing it. Third, for edge E2B/E4B variants Google uses a clustered embedder: the drafter selects likely token clusters and only scores tokens inside those clusters to avoid computing huge vocab logits. Those details matter. The verification step still guarantees identical final logits whenever the target accepts drafted tokens, so you get the speedup without changing model outputs.

Google ships MTP as separate drafter checkpoints and docs that show how to enable speculative decoding across major runtimes. The workflow is drop‑in: load the target Gemma model and its matching drafter, flip the speculative path in your runtime, and let the runtime coordinate draft → verify → accept. Runtimes expose two primary knobs: drafter model id and draft length (how many tokens to propose per draft pass). The runtime pipelines the drafter’s autoregressive generation in parallel with scheduling a single verification forward pass on the target model. Integration points are already live: examples exist for LiteRT, vLLM, Hugging Face Transformers, and MLX, and community ports let you experiment locally. Crucially, no retraining of the target model is required. The drafter reuses the target’s embeddings and activations, letting you experiment with different drafters while keeping production models unchanged.

Google reports up to ~3× tokens/sec, but realized speed depends on three predictable factors. Draft quality and length. Long draft proposals that are frequently wrong force the target to correct tokens and re‑compute, eroding gains. In practice the sweet spot is small n (2–8 tokens) and a drafter sized to be fast yet accurate on next‑token statistics. Hardware and batching. Dense targets benefit the most at low batch sizes because the verification pass reuses KV/attention heavily. MoE targets complicate the picture: different tokens can route to different experts, and verifying multiple tokens may require loading extra expert weights, which is expensive at batch size 1. Google notes the 26B A4B MoE variant sees better speedups at batch sizes of 4–8, where expert activation overlap improves reuse. Memory and VRAM tradeoffs. Running two models raises memory pressure. Two practical patterns emerge in community ports: use a very small drafter (sub‑2B) to stay within GPU RAM, or run the drafter on CPU/edge threads while the target model stays on GPU. These are engineering tradeoffs you can measure and tune.

This is a systems engineering win, not a new modeling breakthrough. Speculative decoding has research precedent; Google’s contribution is operationalizing it cleanly for a major open model family and adding the engineering touches that matter in production. For builders this yields two immediate benefits. First, latency becomes an explicit knob you can flip without retraining the target model — valuable for on‑device assistants, coding tools, and planning agents that need quick multi‑step reactions. Second, open‑source runtimes get a composable accelerator: you can plug MTP drafters into vLLM/LiteRT/Hugging Face stacks and measure end‑to‑end gains. Be skeptical of universal 3× claims. The multiplier depends on drafter accuracy, batch size, MoE routing, and memory constraints. Expect a realistic 1.5–3× range in dense regimes with a reasonable drafter and modest VRAM. The takeaway is practical: choose a drafter, profile latencies, and treat MTP as another systems tool to trade latency, memory, and throughput.