catplusplus/gemma-4-12B-it-assistant-fp8

Google DeepMind Gemma 4 Collection | GitHub | Launch Blog | Documentation
License: Apache 2.0 | Quantization & Calibration: catplusplus

---

Gemma 4 12B IT Assistant — MSE-Calibrated Dynamic FP8 Speculative Drafter

This repository provides an MSE-calibrated, dynamic mixed-precision FP8 Multi-Token Prediction (MTP) assistant model for Google Gemma 4 12B Instruct (google/gemma-4-12B-it and NVFP4 checkpoints such as xdavxd/gemma-4-12B-it-heretic-v2-NVFP4).

It delivers up to 2.5x–3.0x decoding acceleration in speculative decoding pipelines under vLLM and SGLang while cutting linear compute memory bandwidth traffic by 50%, maintaining reference-grade numerical fidelity (SNR > 44.0–45.8 dB, Cosine Similarity > 0.9999), and retaining the exact draft acceptance rate of the uncompressed BF16 drafter (45%–51%).

---

🌟 Key Highlights & Innovations

1. Dynamic Mixed Precision (Unsloth-Profile Sensitivity Protection):

   • FP8 E4M3 (torch.float8_e4m3fn): Applied to all 16 compute-heavy feed-forward and query projections (gate_proj, up_proj, down_proj, and q_proj). Reduces linear tensor memory traffic by 50% on GPU memory buses.
   • Protected Native BF16: Sensitive cross-model recurrent bridge layers (pre_projection $[1024 \times 7680]$ and post_projection $[3840 \times 1024]$), attention output projections (o_proj, which lack pre-norm damping), and token embeddings are strictly preserved in BF16. This eliminates compounding recurrent drift across multi-step draft sequences.

2. Data-Driven MSE Line-Search Calibration:

   • Rather than relying on simple maximum-absolute scaling, each quantized tensor underwent an optimal grid search minimizing Mean Squared Reconstruction Error ($\min_s |W - \hat{W}|_F^2$) against multi-domain coding and systems reasoning distributions.
   • Achieves reference-grade Signal-to-Noise Ratio (44.02 dB – 45.85 dB) and near-perfect directional fidelity (Cosine Similarity > 0.9999).

3. 100% Zero-Patch Engine Compatibility:

   • Cleansed phantom KV projection scale artifacts (k_proj.weight_scale, v_proj.weight_scale). Because Gemma 4's MTP assistant decoder shares its Key and Value projections with backbone layers 46 & 47, removing these phantom keys allows vLLM's AutoWeightsLoader and SGLang to load 100% of weights out of the box in 1.18 seconds without manual code modifications.

4. Consumer GPU Friendly (16 GB VRAM Budgets):

   • Reclaims ~150–200 MB of VRAM compared to the uncompressed drafter (691 MB safetensors vs 845 MB BF16) and halves linear activation workspace buffers.
   • On a single 16 GB GPU (e.g. RTX 5060 Ti, RTX 4070 Ti Super, RTX 4080 Mobile), this reclaimed margin provides an extra ~4,000–6,000 tokens of FP8 KV cache, preventing out-of-memory errors during long-context reasoning.

---

📊 Empirical Benchmarks & Compute Stats

Hardware Testbed

• GPUs: Dual NVIDIA GeForce RTX 5060 Ti (16 GB VRAM each, Ada/Blackwell SM120).
• Target Backbone: xdavxd/gemma-4-12B-it-heretic-v2-NVFP4 (12B Instruct, native CUTLASS NVFP4 SM120 tensor cores executing OMMA.SF.16864.F32.E2M1.E2M1.UE4M3.4X PTX atoms).
• Serving Engine: vLLM 0.19.0+ / ThinkMAX (CutlassFP8ScaledMMLinearKernel for drafter).

Decoding Speed & Speculative Throughput

Mode	Backbone Model	Drafter Model	Speculative Acceptance	Mean ITL (ms)	Throughput (tok/s)
Direct Autoregressive	Gemma-4-12B NVFP4	None (Spec Off)	N/A	27.9 ms	35.8 tok/s
Deep Reasoning (MTP BF16)	Gemma-4-12B NVFP4	Gemma-4 Assistant BF16	51.2%	33.5 ms	78.4 tok/s
Deep Reasoning (MTP FP8 Calibrated)	Gemma-4-12B NVFP4	Gemma-4 Assistant FP8	50.8%	32.6 ms	85.0 – 92.7 tok/s (Peak 103.1 tok/s)
Batch Concurrency ($C=4$)	Gemma-4-12B NVFP4	Gemma-4 Assistant FP8	45.4%	29.9 ms	309.8 tok/s aggregate

Signal-to-Noise Ratio & Reconstruction Quality

Layer Category	Precision	SNR (dB)	Cosine Similarity	Relative Size Reduction
gate_proj.weight	FP8 E4M3 + Scale	44.91 dB	0.999968	-50.0%
up_proj.weight	FP8 E4M3 + Scale	45.85 dB	0.999974	-50.0%
down_proj.weight	FP8 E4M3 + Scale	44.02 dB	0.999912	-50.0%
q_proj.weight	FP8 E4M3 + Scale	45.10 dB	0.999981	-50.0%
o_proj.weight	BF16 (Protected)	$\infty$ (Exact)	1.000000	0.0% (Stability Guard)
pre_projection	BF16 (Protected)	$\infty$ (Exact)	1.000000	0.0% (Recurrent Guard)
post_projection	BF16 (Protected)	$\infty$ (Exact)	1.000000	0.0% (Recurrent Guard)

Code Execution Gauntlet (100% Pass Rate)

Tested against a rigorous multi-step systems programming suite:

• Thread-Safe TTL LRU Cache: Concurrency, mutex locking, and eviction mechanics — PASSED
• Resilient Async Event Dispatcher: AsyncIO error isolation and task dispatch — PASSED
• Interval Tree with Overlap Search: Recursive binary tree search and node balancing — PASSED
• API Schema Dataclasses & JSON Validation: Strict type serialization and coercion — PASSED
• OpenCode Autonomous Agent: Verified autonomously building packages, writing tests, and running test runners with zero syntax or runtime errors.

---

🚀 Quickstart & Usage

1. Running with vLLM

Launch vLLM with MTP speculative decoding enabled:

vllm serve xdavxd/gemma-4-12B-it-heretic-v2-NVFP4 \
  --spec-method mtp \
  --spec-model catplusplus/gemma-4-12B-it-assistant-fp8 \
  --host 0.0.0.0 \
  --port 8000 \
  --gpu-memory-utilization 0.94 \
  --max-model-len 131072 \
  --kv-cache-dtype fp8 \
  --enable-prefix-caching \
  --enable-chunked-prefill

2. Running with SGLang

Launch SGLang with Frozen-KV speculative draft decoding:

python3 -m sglang.launch_server \
  --model-path xdavxd/gemma-4-12B-it-heretic-v2-NVFP4 \
  --speculative-draft-model-path catplusplus/gemma-4-12B-it-assistant-fp8 \
  --speculative-algorithm Frozen-KV \
  --speculative-num-steps 3 \
  --host 0.0.0.0 \
  --port 8000 \
  --mem-fraction-static 0.90 \
  --cuda-graph-backend-prefill tc_piecewise

3. ThinkMAX Configuration Example (~/.thinkmax/config.yaml)

workers:
  summarizer-gemma4:
    listen: "0.0.0.0:9005"
    backend: "127.0.0.1:19005"
    command:
      executable: "/home/olegk/venv/vllm/bin/python"
      args:
        - "/home/olegk/venv/vllm/bin/vllm"
        - "serve"
        - "xdavxd/gemma-4-12B-it-heretic-v2-NVFP4"
        - "--port"
        - "19005"
        - "--spec-method"
        - "mtp"
        - "--spec-model"
        - "catplusplus/gemma-4-12B-it-assistant-fp8"
        - "--kv-cache-dtype"
        - "fp8"

---

📜 Architectural Details (Gemma 4 MTP Drafter)

Multi-Token Prediction (MTP) in Gemma 4 works by extending the base 12B transformer with a lightweight 4-layer recurrent assistant decoder:

• Hidden size: $1024$
• Intermediate size: $8192$
• Number of attention heads: 16 (Head dimension 256)
• Shared KV Layers: 4 (shares Key and Value representations directly with backbone layers 46 & 47, eliminating separate KV projections in the drafter)
• Cross-model bridges: pre_projection $[1024 \times 7680]$ routes representations from the base model into the drafter; post_projection $[3840 \times 1024]$ routes drafted representations back to the verification head.

---

⚖️ License & Attribution

• Base Weights & Architecture: Developed by Google DeepMind. Released under the Gemma Terms of Use and Apache 2.0.
• Quantization, Calibration & Packaging: Oleg Kibirev (@catplusplus).