Enable Gemma 4 MTP (assistant_model) for faster decode

app.py

@@ -15,7 +15,12 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-12b-it")
+MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-12B-it")
+# Gemma 4's official MTP drafter — a ~0.8GB multi-token-prediction head. Passing
+# it as `assistant_model` to generate() is "all you need to enable MTP"
+# (https://ai.google.dev/gemma/docs/mtp/mtp): the target verifies several draft
+# tokens per forward pass, so we run far fewer 12B passes per reply.
+ASSISTANT_ID = os.environ.get("ASSISTANT_ID", "google/gemma-4-12B-it-assistant")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
@@ -25,6 +30,15 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+assistant_model = AutoModelForCausalLM.from_pretrained(
+    ASSISTANT_ID,
+    dtype=torch.bfloat16,
+    device_map="auto",
+)
+assistant_model.eval()
+assistant_model.generation_config.num_assistant_tokens = 4
+assistant_model.generation_config.num_assistant_tokens_schedule = "heuristic"
+
 DEFAULT_SYSTEM = (
     "You are a host on 'Small Talk', a live AI-to-AI robot podcast hosted by "
     "Reachy Mini robots. Stay fully in character. Keep every reply short, witty "
@@ -66,6 +80,7 @@ def chat(message, history, system_prompt, temperature, max_new_tokens):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     gen_kwargs = dict(
         **inputs,
+        assistant_model=assistant_model,   # ← MTP / multi-token prediction
         streamer=streamer,
         max_new_tokens=int(max_new_tokens),
         do_sample=temperature > 0,
@@ -91,8 +106,9 @@ demo = gr.ChatInterface(
     title="Small Talk · Gemma 4 12B brain",
     description=(
         "The live-banter brain for the Small Talk robot podcast — Gemma 4 12B "
-        "(QAT-trained, bf16) on ZeroGPU. Pass a persona as the system prompt to "
-        "voice a character. Callable as an API by the podcast backend."
+        "on ZeroGPU, sped up with official **MTP** (multi-token prediction). Pass "
+        "a persona as the system prompt to voice a character. Callable as an API "
+        "by the podcast backend."
     ),
 )