This repository contains the Gemma 4 31B Instruction-Tuned Assistant draft model quantized to native FP4 (NVFP4) precision for high-efficiency inference on NVIDIA Blackwell GPU architectures (e.g., B200, B100).

This assistant is designed to be used as a draft model for Speculative Decoding alongside the main gemma-4-31B-it-qat-NVFP4-Blackwell model.

Model Details

Base Model: google/gemma-4-31B-it-qat-q4_0-unquantized-assistant
Quantization Format: NVFP4 (native FP4 on Blackwell)
Group Size: 16 (per-group scaling)
Excluded Modules: lm_head
Producer: NVIDIA Model Optimizer (modelopt)

Speculative Decoding Quickstart

Run vLLM on Blackwell GPUs:

bash

vllm serve melcheikh/gemma-4-31B-it-qat-NVFP4-Blackwell \
    --quantization modelopt \
    --speculative-model melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell \
    --num-speculative-tokens 3 \
    --speculative-draft-limit 4

License

Subject to the Gemma Terms of Use.

⚠️ Troubleshooting: vLLM Dimension Mismatch Error

If you attempt to run speculative decoding using this quantized NVFP4 assistant model in vLLM, you may encounter the following initialization error:

RuntimeError: start (0) + length (8192) exceeds dimension size (4096).

Why this happens

By default, vLLM's MTP (Multi-Token Prediction) loader (gemma4_mtp.py) assumes that the assistant model is not quantized. Because of this, it initializes the MLP (Gemma4MLP) and attention projection (Gemma4MTPAttention) layers with quant_config=None, expecting full-sized weights (e.g., shape [1024, 8192]) instead of the ModelOpt quantized/packed weights (shape [1024, 4096]).

The Patch

To fix this, edit the vllm package file vllm/model_executor/models/gemma4_mtp.py in your environment:

Enable quant_config on the MLP layers: Locate self.mlp = Gemma4MLP(...) around line 300 and change quant_config=None to quant_config=quant_config.

python
1# Change this: 2self.mlp = Gemma4MLP( 3 hidden_size=self.hidden_size, 4 intermediate_size=text_config.intermediate_size, 5 hidden_activation=text_config.hidden_activation, 6 quant_config=None, 7 prefix=f"{prefix}.mlp", 8) 9 10# To this: 11self.mlp = Gemma4MLP( 12 hidden_size=self.hidden_size, 13 intermediate_size=text_config.intermediate_size, 14 hidden_activation=text_config.hidden_activation, 15 quant_config=quant_config, 16 prefix=f"{prefix}.mlp", 17)
Enable quant_config on the attention projection layers: Locate self.q_proj and self.o_proj inside Gemma4MTPAttention.__init__ around line 185 and change quant_config=None to quant_config=quant_config.

python
1# Change this: 2self.q_proj = ColumnParallelLinear(..., quant_config=None, ...) 3self.o_proj = RowParallelLinear(..., quant_config=None, ...) 4 5# To this: 6self.q_proj = ColumnParallelLinear(..., quant_config=quant_config, ...) 7self.o_proj = RowParallelLinear(..., quant_config=quant_config, ...)
Running vLLM with maximized context: Because speculative decoding requires additional KV cache memory, limit --max-num-seqs to 1 to avoid running out of VRAM (CUDA OOM) at longer contexts:

bash
1vllm serve melcheikh/gemma-4-31B-it-qat-NVFP4-Blackwell \ 2 --quantization modelopt_fp4 \ 3 --kv-cache-dtype fp8 --max-model-len 32000 --gpu-memory-utilization 0.97 --max-num-seqs 1 --spec-model melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell --spec-tokens 3