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

Model Overview

• Model Architecture: Gemma4AssistantForCausalLM (Multi-Token Prediction drafter)
  • Input: Text
  • Output: Text (draft tokens)
• Model Optimizations:
  • Weight quantization: FP8, static, per output channel
  • Activation quantization: FP8, dynamic, per token
• 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_DYNAMIC scheme. Weight scales are computed offline, one per output channel; activation scales are computed at runtime, one per token. No calibration data is used. 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.

Relationship to the FP8-block variant

A sibling checkpoint, BarraHome/gemma-4-31B-it-assistant-FP8-block, was produced from the same base model with the FP8_BLOCK scheme. The two differ only in scale granularity:

This variant is 217,672 bytes larger. Per-channel weight scales carry more elements than a 128×128 block grid does — for example gate_proj stores a [8192, 1] scale here versus [64, 8] in the block variant.

Neither variant has been benchmarked, so this repo makes no claim that one is more accurate or faster than the other. Pick based on what your serving stack supports.

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-dynamic"

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-dynamic", "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 dynamic 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-dynamic"

model_free_ptq(
    model_stub=MODEL_ID,
    save_directory=SAVE_DIR,
    scheme="FP8_DYNAMIC",
    ignore=["re:.*vision.*", "lm_head", "re:.*embed_tokens.*"],
)

Produced with compressed-tensors version 0.18.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.