SYSTEM NOTICE

Auto translation by AI. Be sure, accuracy, nuances and authorial intent may not be fully reflected.
見出し画像

Gemma 4 is up to 3x faster: How MTP drafters accelerate inference

What you will learn in this article

  • The nature of the "MTP (Multi-Token Prediction) drafter" released by Google for Gemma 4

  • Where LLM inference is actually slowing down

  • The logic behind how speculative decoding creates speed

  • Points to keep in mind when trying it in your own environment

What happened

For Gemma 4, which was released a few weeks ago, Google has released an additional component called an "MTP drafter." According to the announcement, using this mechanism can increase inference speed by up to 3 times. The key point is that this is achieved without sacrificing quality.

Gemma 4 itself is an open model family with a 256K token context length, supporting over 140 languages, and is available in both Dense and MoE (Mixture of Experts) configurations. What was added this time is a small "draft model" that works alongside the main model.

Incidentally, MTP was intentionally omitted from the initial release of Gemma 4. It was a topic of discussion in the developer community as to "why it wasn't included from the start," but it appears it has now been enabled after prioritizing quality assurance.

Why is standard LLM inference slow?

Many people have likely felt, "GPUs should be fast, so why am I waiting for token generation?" In fact, an essential bottleneck is hidden here.

Standard LLMs generate tokens one by one in an autoregressive manner. Every time a single token is produced, billions to tens of billions of parameters must be moved from VRAM to the compute units.

In short, to summarize the situation:

  • It is memory bandwidth, not computational performance, that is the limiting factor

  • Most of the compute units are left in a "waiting" state

  • Even so, tokens only advance one at a time

This is the result. This tendency is even stronger with consumer-grade GPUs. It is quite frustrating that output is slow even though computational resources are sitting idle.

The idea of speculative decoding

A technique called speculative decoding is used to utilize these idle computational resources. The principle is simple:

  1. A lightweight "draft model" predicts multiple future tokens at once

  2. The heavy 'target model (main body)' verifies those predictions in bulk

  3. The parts that pass verification are adopted all at once, and for the parts that don't, the target model outputs the correct tokens and starts over

The point is that the target model can verify multiple tokens in the time it takes to calculate one token. It takes advantage of the fact that 'memory bandwidth is the bottleneck' mentioned at the beginning, using idle calculation units to perform batch verification.

When verification passes, in addition to the several tokens proposed by the draft, the target model itself outputs one more token as a continuation. This means that several tokens are processed in 'the time it would normally take to output just one token'.

Since the final output is ultimately what the target model has approved, the quality is exactly the same as normal generation. This is a property that cannot be overlooked, as there is no compromise where it 'got faster but became slightly less intelligent'.

What is new about the Gemma 4 MTP drafter?

Speculative decoding itself is not a new idea. Even so, the Gemma 4 MTP drafter incorporates several eye-catching innovations.

1. Shared input embeddings

The draft model is not an independent, separate model; it shares the target model's input embedding table as is. Since it does not have duplicate weights, it is also memory-friendly.

2. Reuse of target model activations

The draft model borrows the final layer activations (intermediate representations) of the target model. It concatenates this with token embeddings and reduces it to its own dimension before proceeding with inference. It is like using the understanding that 'the main body already understands the context' as a foundation.

3. KV cache sharing

Instead of recalculating the context every time, it shares the target model's KV cache. This is a subtle but effective optimization.

4. Efficient embedding layers (for E2B / E4B)

In small models for the edge, the final logit calculation for the entire vocabulary itself becomes a heavy burden. Therefore, an optimization is included where similar tokens are grouped into clusters, and the final calculation is performed only using tokens within the cluster after first narrowing down 'which cluster is promising'. This is likely to be effective in environments like smartphones.

Benchmarks for speed improvement

According to the official announcement, here is a rough idea for each intended use case.

  • 31B Dense model: Significant effects can be expected on workstations and high-performance GPUs

  • 26B MoE model: With a batch size of about 4-8, a speed improvement of about 2.2x has been reported even in local environments

  • E2B / E4B edge models: Response is visibly improved for on-device use, resulting in reduced battery consumption

I would like to add a note here about a point where beginners often stumble.MoE models tend to see less benefit at a batch size of 1. The reason is that MoE has a structure where it calls different "experts" for each token, and when verifying multiple tokens in a draft, it becomes necessary to additionally load expert weights from memory. As the batch size increases, the probability of experts overlapping across multiple requests rises, and efficiency improves significantly.

If you find that "I tried it with MoE but it didn't get as fast as I thought," the first thing to do is to review whether your configuration allows for increasing the batch size, which makes it easier to isolate the cause.

Where can it be run?

The MTP drafter, like the Gemma 4 model itself, is released under the Apache 2.0 license. These are terms that make it easy to move forward with commercial use.

As for supported execution environments,

  • Hugging Face Transformers

  • vLLM

  • SGLang

  • MLX (for Apple Silicon)

  • Ollama

  • LiteRT-LM (for on-device)

  • Google AI Edge Gallery (can be tried via Android / iOS apps)

are all available options. The fact that there is a high probability you can use your existing inference runtime as-is significantly lowers the psychological hurdle for adoption.

Things to keep in mind when trying it out

There are a few points to keep in mind when actually integrating it.

First, always match the target model and drafter pair. Since a corresponding drafter is provided for each size of Gemma 4, if you get the combination wrong, you may not achieve the intended speed.

Next, the optimal batch size varies depending on the hardware. The behavior changes for RTX-series GPUs, A100s, and Apple Silicon. If you are running an MoE model, it is easier to get a clear outlook if you try adjusting the batch size early on.

Finally, I want to emphasize that it does not mean "it gets faster at the cost of intelligence". Because speculative decoding is a mechanism where the final model is always involved in verification, the output distribution remains consistent. It is safe to say this is the type of optimization you can adopt with peace of mind.

Summary

The bottleneck in LLM inference is not "thinking power" but "the speed of moving weights from memory." The MTP drafter cleverly avoids this and, by using idle computational resources for batch verification, achieves a significant speed increase while maintaining quality.

The fact that such optimizations are officially provided for open models like Gemma 4 is genuinely welcome news for those considering local development or on-device AI. Even on local machines, sizes that previously felt 'a bit too demanding' might now become a realistic option.

It seems like a good idea to start by trying to plug in a drafter with the inference runtime you currently use.

Reference links

いいなと思ったら応援しよう!