Skip to content

gemma4-assistant.md

Latest commit

 

History

History
134 lines (109 loc) · 7.48 KB

gemma4-assistant.md

File metadata and controls

134 lines (109 loc) · 7.48 KB
Raw

Gemma 4 Multi-Token Prediction (MTP) drafter

This page documents the gemma4-assistant GGUF architecture, which is the file-format counterpart to Google's official Multi-Token Prediction (MTP) drafters released alongside the Gemma 4 family on 2026-05-05.

What it is

Each Gemma 4 model ships with a paired ~78 M-param "assistant" — a 4-layer transformer that proposes several tokens ahead of the target model. The target then verifies the proposal in parallel, the same way as classical speculative decoding (Leviathan et al. 2022), but with an architecture co-designed with the target backbone:

  • Q-only attention. The drafter does not have its own K/V projections; every layer reads K/V from the target backbone's last layer of the matching attention type (one entry for full-attention, one for sliding-window).
  • 2× backbone input. The forward signature is [target_last_hidden, embed(next_token)] projected through pre_projection to the small assistant hidden size (256 for both released drafters).
  • Centroid-clustered output head. Instead of dense logits over the full 262 144-row vocabulary, the assistant scores 2 048 centroids, picks the top 32, and only computes fine-grained logits inside those clusters. The chosen positions are scattered into the vocab; everything else is masked.
  • Per-layer scalar gate. A single layer_scalar weight per layer rescales the residual stream — same shape as the Gemma 4 backbone uses.

Status of this implementation

Capability Status
Convert Hugging Face checkpoint → GGUF ✅
gguf-dump inspection ✅
Model load (llama_model_load_from_file) ✅
Forward pass / llama_decode ✅ foreign-KV external-assistant path shipped
Speculative-decoding integration ✅ external-assistant MTP shipped; validated 47.3% accept on Gemma4-26B-A4B-it (external drafter, --temp 0, chat-templated, ROCm gfx1150; upstream PR #23398 CUDA reference: 58.8%); bundled-MTP qwen3.5 path: 63.4%; accept varies with draft budget and prompt

The runtime path is intentionally gated: completing it requires plumbing the target context's last-layer K/V tensors into the drafter context, which is a new piece of public API surface. Landing the model architecture first lets the GGUF format stabilise so the ecosystem (Ollama, LM Studio, the quantization tools, etc.) can ingest the released checkpoints while the runtime piece is reviewed separately.

Conversion

# Download the assistant checkpoint (E2B; swap for `E4B-it-assistant` to get the
# E4B drafter). Both are ~158 MB at bf16.
hf download google/gemma-4-E2B-it-assistant --local-dir gemma-4-E2B-it-assistant

python3 convert_hf_to_gguf.py gemma-4-E2B-it-assistant \
    --outfile gemma-4-E2B-it-assistant.gguf \
    --outtype bf16

The converter is registered against the Gemma4AssistantForCausalLM architecture, mirrors the gemma4 tokenizer (262 144 tokens) and writes the following metadata:

  • gemma4-assistant.backbone_hidden_size — 1536 for E2B, 2560 for E4B.
  • gemma4-assistant.assistant.num_centroids — 2048.
  • gemma4-assistant.assistant.centroid_intermediate_top_k — 32.
  • gemma4-assistant.assistant.use_ordered_embeddings — true.

Tensor schema (per-block tensors are written for blk.0..3):

GGUF tensor Shape Source name
token_embd.weight [hidden, vocab] model.embed_tokens.weight
output_norm.weight [hidden] model.norm.weight
assist_pre_proj.weight [2*backbone, hidden] pre_projection.weight
assist_post_proj.weight [hidden, backbone] post_projection.weight
assist_embed_centroids.weight [hidden, num_centroids] masked_embedding.centroids.weight
assist_token_ordering.weight [vocab] (F32, integer values) masked_embedding.token_ordering
blk.{i}.attn_q.weight [hidden, n_head*head_dim] model.layers.{i}.self_attn.q_proj.weight
blk.{i}.attn_q_norm.weight [head_dim] model.layers.{i}.self_attn.q_norm.weight
blk.{i}.attn_output.weight [n_head*head_dim, hidden] model.layers.{i}.self_attn.o_proj.weight
blk.{i}.attn_norm.weight [hidden] model.layers.{i}.input_layernorm.weight
blk.{i}.post_attention_norm.weight [hidden] model.layers.{i}.post_attention_layernorm.weight
blk.{i}.ffn_norm.weight [hidden] model.layers.{i}.pre_feedforward_layernorm.weight
blk.{i}.post_ffw_norm.weight [hidden] model.layers.{i}.post_feedforward_layernorm.weight
blk.{i}.ffn_gate.weight [hidden, intermediate] model.layers.{i}.mlp.gate_proj.weight
blk.{i}.ffn_up.weight [hidden, intermediate] model.layers.{i}.mlp.up_proj.weight
blk.{i}.ffn_down.weight [intermediate, hidden] model.layers.{i}.mlp.down_proj.weight
blk.{i}.layer_output_scale.weight [1] model.layers.{i}.layer_scalar

token_ordering keeps F32 storage because the converter pipeline force-stores 1-D tensors as F32; every value is below 2^18 so the cast is lossless.

Variants supported

  • E2B: 4 layers, hidden 256, KV-heads 1, backbone 1536. ~78 M params.
  • E4B: 4 layers, hidden 256, KV-heads 2, backbone 2560. ~78 M params.
  • 26B-A4B and 31B drafters use the same architecture and convert with the same command (only the backbone size changes); they are not exercised by this PR but should work — please file an issue if not.

Runtime implementation (shipped)

Shipped: The external-assistant MTP decoding path is implemented in this fork. The three pieces described below were all resolved and shipped as part of the fork's Phase 2 Gemma4 external-assistant MTP work (see CHANGELOG.md).

Validated accept rate: 47.3% on Gemma4-26B-A4B-it with the external drafter (n_drafted=112 n_accept=53, --temp 0, chat-templated prompt, ROCm gfx1150, llama-speculative-simple --spec-type draft-mtp; commits ca62c0756, 6506eb25a, f2558d258). The bundled-MTP qwen3.5 path registers 63.4% accept on the same validation run. Upstream PR #23398 CUDA reference: 58.8%. Accept rate varies with draft budget, temperature, and prompt; chat templates are required for instruction-tuned targets.

The three pieces that were required:

Implementing decoding faithfully needs three pieces that do not exist in mainline llama.cpp; they were implemented as fork-divergence (+666 LoC):

  1. Cross-context K/V exposure. Reading the target context's last-layer K/V at decode time and passing those tensors as graph inputs to the drafter context. The transformers reference takes them as the shared_kv_states keyword argument of the drafter's forward.
  2. Bidirectional attention with external K/V. Each drafter layer attends over the target's existing tokens (causal-by-construction since they're already decoded), with a flipped sliding window.
  3. Centroid masked logits. A top_k + get_rows + scatter combination over the LM head. ggml's existing top_k covers the first step; the scatter requires either a dense fill plus set_rows, or a new graph op.

Upstreaming these pieces is deferred to a dedicated PR for separate review.