[None][feat] Add Gemma4 MTP assistant support

docs/source/models/supported-models.md

@@ -20,6 +20,7 @@ The following is a table of supported models for the PyTorch backend:
 | `Gemma3nForConditionalGeneration` [^7]| Gemma 3n                           | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it` |
 | `Gemma4ForConditionalGeneration`     | Gemma 4                            | `google/gemma-4-E2B-it`, `google/gemma-4-E4B-it`, `google/gemma-4-26B-A4B-it` [^6], `google/gemma-4-31B-it` [^6] |
 | `Gemma4UnifiedForConditionalGeneration` | Gemma 4 12B Unified (encoder-free) | `google/gemma-4-12B`, `google/gemma-4-12B-it`              |
+| `Gemma4AssistantForCausalLM`         | Gemma 4 MTP assistant              | `google/gemma-4-E2B-it-assistant`, `google/gemma-4-E4B-it-assistant`, `google/gemma-4-26B-A4B-it-assistant`, `google/gemma-4-31B-it-assistant` |
 | `Glm4MoeForCausalLM`                 | GLM-4.5, GLM-4.6, GLM-4.7          | `THUDM/GLM-4-100B-A10B`                      |
 | `Glm4MoeLiteForCausalLM` [^5]        | GLM-4.7-Flash                      | `zai-org/GLM-4.7-Flash`                      |
 | `GlmMoeDsaForCausalLM`               | GLM-5                              | `zai-org/GLM-5`                              |
@@ -76,7 +77,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `GptOssForCausalLM`              | Yes               | Yes        | Yes                        | Yes                   | Yes             | No  | Yes              | No                | Yes    | Yes           | Yes              | Yes            | N/A                      | Yes                   | Yes             |
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
-| `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
+| `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | Yes | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | Yes | No               | No                | No     | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
 | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |

examples/llm-api/quickstart_advanced.py

@@ -308,6 +308,8 @@ def setup_llm(args, **kwargs):
     if spec_decode_algo == 'MTP':
         if not args.use_one_model:
             print("Running MTP eagle with two model style.")
+        speculative_model = (args.draft_model_dir if args.draft_model_dir
+                             is not None else args.model_dir)
         spec_config = MTPDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             use_relaxed_acceptance_for_thinking=args.
@@ -318,7 +320,7 @@ def setup_llm(args, **kwargs):
             use_dynamic_tree=args.use_dynamic_tree,
             dynamic_tree_max_topK=args.dynamic_tree_max_topK,
             max_total_draft_tokens=args.max_total_draft_tokens,
-            speculative_model=args.model_dir)
+            speculative_model=speculative_model)
     elif spec_decode_algo == "EAGLE3":
         spec_config = Eagle3DecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,

examples/models/core/gemma/README.md

@@ -8,12 +8,12 @@ loaded directly. The legacy TensorRT engine flow (`convert_checkpoint.py` /
 
 Gemma 4 runs on the **PyTorch backend** — HuggingFace checkpoints are loaded directly. The legacy TensorRT engine flow (`convert_checkpoint.py` / `trtllm-build`) is not required and is not covered here.
 
-| HuggingFace checkpoint        | Modalities                       | Notes                                  |
-|-------------------------------|----------------------------------|----------------------------------------|
-| `google/gemma-4-E2B-it`       | text + image + video + audio     | Single-GPU friendly                    |
-| `google/gemma-4-E4B-it`       | text + image + video + audio     | Single-GPU friendly                    |
-| `google/gemma-4-26B-A4B-it`   | text + image + video (MoE)       | Multi-GPU recommended; no audio tower  |
-| `google/gemma-4-31B-it`       | text + image + video             | Multi-GPU recommended; no audio tower  |
+| HuggingFace checkpoint      | Modalities                   | Matching MTP assistant                         | Notes                                 |
+| --------------------------- | ---------------------------- | ---------------------------------------------- | ------------------------------------- |
+| `google/gemma-4-E2B-it`     | text + image + video + audio | `google/gemma-4-E2B-it-assistant`              | Single-GPU friendly                   |
+| `google/gemma-4-E4B-it`     | text + image + video + audio | `google/gemma-4-E4B-it-assistant`              | Single-GPU friendly                   |
+| `google/gemma-4-26B-A4B-it` | text + image + video (MoE)   | `google/gemma-4-26B-A4B-it-assistant`          | Multi-GPU recommended; no audio tower |
+| `google/gemma-4-31B-it`     | text + image + video         | `google/gemma-4-31B-it-assistant`              | Multi-GPU recommended; no audio tower |
 
 All four variants ship the vision tower (image + video). The audio tower is only present on `E2B` / `E4B`. The examples below use `google/gemma-4-E4B-it` (small, full multimodal) — swap the model name for the other variants and bump `--tp_size` (e.g. `4` or `8`) for the larger checkpoints.
 
@@ -47,6 +47,49 @@ curl http://localhost:8000/v1/chat/completions \
 
 The `/v1/chat/completions` endpoint applies the Gemma 4 chat template automatically.
 
+### MTP speculative decoding
+
+Gemma 4 supports Multi-Token Prediction (MTP) speculative decoding through the
+PyTorch execution path. The target loads its matching assistant checkpoint,
+and the Q-only assistant reads the target model's KV cache. Create a server
+configuration for the target/assistant pair:
+
+```bash
+cat > gemma4_mtp.yaml <<'EOF'
+speculative_config:
+  decoding_type: MTP
+  max_draft_len: 3
+  mtp_eagle_one_model: true
+  speculative_model: google/gemma-4-E4B-it-assistant
+kv_cache_config:
+  enable_block_reuse: false
+EOF
+
+trtllm-serve google/gemma-4-E4B-it \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --config gemma4_mtp.yaml
+```
+
+The assistant shares the target model's KV cache, so TensorRT-LLM does not
+allocate a second full-size GPU KV cache for it. The current implementation
+supports the `Gemma4AssistantForCausalLM` assistants for E2B, E4B, 26B-A4B,
+and 31B. Gemma 4 12B uses the `Gemma4UnifiedForConditionalGeneration` and
+`Gemma4UnifiedAssistantForCausalLM` architectures, which are not supported.
+
+For offline inference, pass both checkpoints to the advanced LLM API example:
+
+```bash
+python3 examples/llm-api/quickstart_advanced.py \
+    --model_dir google/gemma-4-E4B-it \
+    --draft_model_dir google/gemma-4-E4B-it-assistant \
+    --spec_decode_algo MTP \
+    --spec_decode_max_draft_len 3 \
+    --use_one_model \
+    --disable_kv_cache_reuse \
+    --apply_chat_template
+```
+
 ### Accuracy evaluation with `trtllm-eval`
 
 `trtllm-eval` is the canonical entry point for accuracy benchmarks. Two tasks relevant to Gemma 4: