Speculative Decoding with Gemma 4

TLDR; We made the attention kernel 6x faster on Gemma-4-26B-A4B, specifically for the unique characteristics of the verification phase of speculative decoding. This led to end-to-end gains as much as 82% in throughput and 1st place on the Artificial Analysis leaderboard on a single AMD MI355X across all context lengths. We see improvements on both MTP and DSpark speculative decoding and we released our own trained DSpark drafter head.

We have developed the fastest inference endpoint for Gemma-4-26B-A4B using a single AMD MI355X GPU. We achieve this through two optimizations: First, we accelerate the verification step of the target model. Second, we train our own DSpark draft model.

Artificial analysis rankings on our AMD MI355X deployment. Numbers may differ from our internal benchmarks due to variation in live traffic.

As the model’s name suggests, Gemma-4-26B-A4B follows a mixture of experts architecture with only 4 billion parameters of the total 26 billion parameters active during inference. Frontier-level efficiency is achieved by setting 26 of their 30 layers to be local sliding window attention, with the remaining four layers set to the usual global attention layers. Instead of a quadratic complexity in token-to-token attention, the local layers only have to attend to the tokens within a neighboring 1024 token window. The head dimension is also reduced from 512 to 256 for these layers.

Why multi-token prediction?

To speed up generation, one of the techniques from speculative decoding [1] is to predict an additional k “draft” tokens for every 1 token that is conventionally generated autoregressively. There are other speculative decoding methods such as the diffusion-based DFlash [2], however, MTP has shown the most robust support to different hardware, which includes AMD’s GPUs.

Google has released their own MTP draft model that is smaller and faster than the base model, but the value of k is limited by the acceptance length determined through the verification of the target model. Tuning the number of speculative tokens, we find k=6 as a sweet spot across concurrencies.

The bottleneck: draft token verification

A useful offline calculation is how many work items are available to fill up the Compute Units of the GPU. Some preliminary profiling points us to the attention kernel responsible for verifying the k+1 draft tokens with a target model forward pass. In Group Query Attention (GQA) [3], the query heads in a group share the same Key-Value(KV) heads, and so computation is only launched once per KV head across the query dimension. This gives a 2D kernel launch grid with size (total_query_blocks, num_kv_heads).

Although the same kernels are launched during prefill, notice that the query sequence length of k+1 is significantly less than the prefill case of one query row per prompt token. One fine detail to consider is that Gemma 4 has different group sizes for full and local attention layers. To be explicit, the number of KV heads is 2 for the full attention and 8 for local attention. For full attention, this is far fewer than other models such as Qwen3.

The grid blocks are then partitioned across the 256 Compute Units of our AMD MI355X. Even at concurrency 4, we have a 14% utilization cap during this stage of our attention pipeline. Through more profiling we confirmed this bottleneck for the full attention layers as each query block must do computation on the whole KV-context. On the other hand, local attention layers only look at the most recent 1024 tokens and also have 4 times as many grid blocks due to the increased number of KV heads.

The fix: 3D split-KV path

By splitting the KV sequence into many segments, we can increase the number of grid blocks to be assigned to Compute Units, increasing GPU occupancy during the token verification step. The launch grid grows from (total_query_blocks, num_kv_heads) to (total_query_blocks, num_kv_heads, num_kv_splits) essentially turning it to a 3D grid with a lightweight reduction kernel to combine the partial results. We find 64 splits is well-tuned for the underlying hardware.

InferenceX-style benchmark with Magpie-Llama-3.1 prompts run on 1xB200

On our internal kernel micro-benchmark, we achieve a kernel speedup factor of 6.6 for C=1 and 6.2 for C=4. Selectively applying this to the verification step of the full attention layer leads to a modest end-to-end throughput improvement up to 7.5% on top of the already parameter-tuned Gemma 4 MTP setup.

To squeeze out even more performance, we decided to train our own speculative drafting head.

Sparking the next frontier (by training DSpark)

During our optimization cycle of Gemma 4, Deepseek released a new method of speculation called DSpark [4]. Their approach introduces inter-token dependencies in the originally parallel generation from DFlash to increase the average acceptance lengths.

With the online training support from vLLM’s Speculators library [5], we were able to train our DSpark head on 600k training samples for 10 epochs on 4xB200s. We used our own curated dataset with an even split between samples from coding, math and natural chat, but also other buckets to ensure more robust generalization in the wild such as tool-calling and longer contexts (>16k tokens).

InferenceX-style benchmark with Magpie-Llama-3.1 prompts run on 1xB200

On lower concurrencies, DSpark matches the performance of MTP with k=6. By applying the 3D-split KV path on the verifier kernel we achieve even more throughput gains. This shows that our kernel optimization ports over to DSpark as well.

InferenceX-style benchmark with Math-500 prompts run on 1xB200

It intuitively checks out that the performance is higher on math than on natural chat. Speculative tokens can be better predicted due to the constrained output space of “template-y” domains.

Our checkpoint is now available on Hugging Face: https://huggingface.co/makora-ai/gemma4-26b-a4b-dspark. We are excited to see what additional inference optimizations the community builds upon it!

What’s next?

While we have focused mainly on the multi-token prediction extension, there is promise to optimize the attention kernel of the base model itself. For example, flash attention kernels can be swapped in for head dimension at most 256 to give additional inference speed-ups[6]. On the theme of full attention layers being the bottleneck, the 512 head dimension on these layers currently restrict such optimizations, even though it only comprises 13% of total layers. We are also working on training more DSpark speculators for larger models, stay tuned for more on that soon.

This post serves as one more signal for the use-case of the agent development cycle for model optimization through MakoraGenerate and MakoraOptimize. The quick feedback loop from speed and accuracy benchmarks with deep profiling from AMD tools like rocprof-computeand rocprof-sys and Nvidia tools like nsight-compute and nsight-systems gives multi-agents an edge in the optimization search space. The human-in-the-loop retains importance as a high-capacity verifier to ensure any proposed optimization paths do not get led astray, just like the role of the verification step in speculative decoding.

References

[1] Leviathan, Yaniv, Matan Kalman, and Yossi Matias. “Fast Inference from Transformers via Speculative Decoding.” arXiv preprint arXiv:2211.17192 (2023).

[2] Chen, Jian, Yesheng Liang, and Zhijian Liu. “DFlash: Block Diffusion for Flash Speculative Decoding.” arXiv preprint arXiv:2602.06036 (2026).

[3] Ainslie, Joshua, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” arXiv preprint arXiv:2305.13245 (2023).

[4] Cheng, Xin, Xingkai Yu, Chenze Shao, Jiashi Li, Yunfan Xiong, Yi Qian, Jiaqi Zhu, Shirong Ma, Xiaokang Zhang, Jiasheng Ye, Qinyu Chen, Chengqi Deng, Jiping Yu, Damai Dai, Zhengyan Zhang, Yixuan Wei, Yixuan Tan, Wenkai Yang, Runxin Xu, Yu Wu, Zhean Xu, Xuanyu Wang, Muyang Chen, Rui Tian, Xiao Bi, Zhewen Hao, Shaoyuan Chen, Huanqi Cao, Wentao Zhang, Anyi Xu, Huishuai Zhang, Dongyan Zhao, and Wenfeng Liang. “DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation.” arXiv preprint arXiv:2607.05147 (2026).

[5] Red Hat. “Speculators: A Unified Library for Speculative Decoding Algorithms in LLM Serving.” GitHub repository. https://github.com/vllm-project/speculators (2025).

[6] Zadouri, Ted, Markus Hoehnerbach, Jay Shah, Timmy Liu, Vijay Thakkar, and Tri Dao. “FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling.” arXiv preprint arXiv:2603.05451 (2026).