melcheikh/gemma-4-31B-it-qat-assistant-NVFP4-Blackwell

Gemma 4 31B IT QAT Assistant NVFP4 Blackwell

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:

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:

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

   # Change this:
   self.mlp = Gemma4MLP(
       hidden_size=self.hidden_size,
       intermediate_size=text_config.intermediate_size,
       hidden_activation=text_config.hidden_activation,
       quant_config=None,
       prefix=f"{prefix}.mlp",
   )
   
   # To this:
   self.mlp = Gemma4MLP(
       hidden_size=self.hidden_size,
       intermediate_size=text_config.intermediate_size,
       hidden_activation=text_config.hidden_activation,
       quant_config=quant_config,
       prefix=f"{prefix}.mlp",
   )
   

2. 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.

   # Change this:
   self.q_proj = ColumnParallelLinear(..., quant_config=None, ...)
   self.o_proj = RowParallelLinear(..., quant_config=None, ...)
   
   # To this:
   self.q_proj = ColumnParallelLinear(..., quant_config=quant_config, ...)
   self.o_proj = RowParallelLinear(..., quant_config=quant_config, ...)
   

3. 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:

   vllm serve melcheikh/gemma-4-31B-it-qat-NVFP4-Blackwell \
       --quantization modelopt_fp4 \
       --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