gemma-4-31B-it-assistant-FP8-block

Model Overview

• Model Architecture: Gemma4AssistantForCausalLM (Multi-Token Prediction drafter)
  • Input: Text
  • Output: Text (draft tokens)
• Model Optimizations:
  • Weight quantization: FP8 (block-wise, 128×128)
  • Activation quantization: FP8 (dynamic, per group of 128)
• Release Date: 2026-08-24
• Version: 1.0
• Model Developer: BarraHome

This is a quantized version of google/gemma-4-31B-it-assistant, the Multi-Token Prediction (MTP) drafter for google/gemma-4-31B-it.

This checkpoint is not a standalone model. It is a draft model: it predicts several tokens ahead so the 31B target model can verify them in parallel, which speeds up decoding while the target model's output distribution — and therefore its quality — is preserved. It must be loaded alongside google/gemma-4-31B-it.

Model Optimizations

The weights and activations of the linear operators inside the drafter's transformer blocks were quantized to FP8 with LLM Compressor, using the data-free FP8_BLOCK scheme: weights use block-wise scaling over 128×128 blocks, activations are quantized dynamically per group of 128. The vision-related modules, lm_head, and embed_tokens are left in their original precision.

22 tensors were converted to float8_e4m3fn (the 4 decoder blocks' q_proj, o_proj, gate_proj, up_proj, down_proj, plus pre_projection and post_projection). Everything else — layer norms, layer scalars, and the embedding table — stays in BF16.

Size

The reduction is ~21%, not the ~50% typical of FP8 conversions. The reason is that model.embed_tokens.weight (262144 × 1024) is 512 MiB on its own — about 73% of this checkpoint — and it is excluded from quantization because it is tied to the output head. The quantized linear layers themselves shrink close to 2×; the embedding table dominates what is left.

Deployment

Use with Transformers

Load the target model and pass this drafter as assistant_model:

from transformers import AutoProcessor, AutoModelForCausalLM

TARGET_MODEL_ID = "google/gemma-4-31B-it"
ASSISTANT_MODEL_ID = "BarraHome/gemma-4-31B-it-assistant-FP8-block"

processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)
target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID,
    dtype="auto",
    device_map="auto",
)

# Assistant model (the drafter)
assistant_model = AutoModelForCausalLM.from_pretrained(
    ASSISTANT_MODEL_ID,
    dtype="auto",
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a short joke about saving RAM."},
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
    enable_thinking=False,
).to(target_model.device)
input_len = inputs["input_ids"].shape[-1]

outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)
print(processor.parse_response(response))

Reading the FP8 weights requires compressed-tensors:

pip install -U transformers torch accelerate compressed-tensors

Use with vLLM

Serve the 31B target model and point speculative decoding at this drafter:

vllm serve google/gemma-4-31B-it \
  --speculative-config '{"method": "mtp", "model": "BarraHome/gemma-4-31B-it-assistant-FP8-block", "num_speculative_tokens": 3}'

Note: this snippet has not been verified against a running server. MTP drafters attach to the target model's hidden states (backbone_hidden_size: 5376 in config.json), so it needs a vLLM build that supports the gemma4_assistant architecture. Check the vLLM speculative decoding docs for the current syntax and for whether your version handles this architecture.

Creation

Produced with data-free FP8 block quantization via LLM Compressor:

from llmcompressor import model_free_ptq

MODEL_ID = "google/gemma-4-31B-it-assistant"
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-block"

model_free_ptq(
    model_stub=MODEL_ID,
    save_directory=SAVE_DIR,
    scheme="FP8_BLOCK",
    ignore=["re:.*vision.*", "lm_head", "re:.*embed_tokens.*"],
    max_workers=8,
    device="cuda:0",
)

Evaluation

This model has not been evaluated. No benchmarks were run and no acceptance-rate or draft-quality measurements were taken against the BF16 drafter.

Because the drafter's proposals are verified by the target model, quantizing it should not change the target model's output quality — a rejected draft token is simply discarded. What quantization can affect is the acceptance rate, and therefore the actual speedup. If you depend on the speedup, measure tokens/second and acceptance rate against the BF16 drafter on your own workload before deploying.

License

Apache 2.0, inherited from the base model. See the Gemma 4 license.