[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`                              |
@@ -75,7 +76,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

@@ -303,6 +303,9 @@ def setup_llm(args, **kwargs):
     if spec_decode_algo == 'MTP':
         if not args.use_one_model:
             print("Running MTP eagle with two model style.")
+            if args.draft_model_dir is None:
+                raise ValueError(
+                    "--draft_model_dir is required for two-model MTP")
         spec_config = MTPDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             use_relaxed_acceptance_for_thinking=args.
@@ -313,7 +316,8 @@ 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=args.model_dir
+            if args.use_one_model else args.draft_model_dir)
     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,42 @@ 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 two-model Multi-Token Prediction (MTP) speculative decoding on the PyTorch backend. Each target checkpoint must use the matching assistant checkpoint listed above. 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: false
+  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 reads the target model's KV cache directly, 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. This path uses two-model MTP, which is deprecated and scheduled for removal in release 1.4.
+
+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 \
+    --no-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:

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -165,6 +165,7 @@ class PlanParams:
     multi_item_params: Optional[FlashInferMultiItemParams] = None
     sm_scale: Optional[float] = None
     window_left: Optional[int] = None
+    kv_pool_id: Optional[int] = None
 
 
 # NB: Some features (multi-item scoring) are only supported with the paged KV-cache wrapper.
@@ -225,6 +226,8 @@ class FlashInferWrappers:
     # and columns before narrowing future updates.
     decode_block_table_active_rows: int = field(default=0, repr=False)
     decode_block_table_active_width: int = field(default=0, repr=False)
+    host_prefill_block_tables: Optional[torch.Tensor] = field(default=None,
+                                                              repr=False)
 
 
 @dataclass(kw_only=True)
@@ -704,15 +707,6 @@ def _post_init_with_buffers(self, buffers) -> None:
                     self._vswa_layer_to_pool[layer_idx] = pool_id
                     if pool_id not in self._vswa_pool_to_rep_layer:
                         self._vswa_pool_to_rep_layer[pool_id] = layer_idx
-                # Build head_dim → pool_id mapping using V2 per-layer head_dim
-                self._vswa_head_dim_to_pool: Dict[int, int] = {}
-                if hasattr(mgr, 'head_dim_per_layer'):
-                    for layer_idx, pool_id in self._vswa_layer_to_pool.items():
-                        hd = mgr.head_dim_per_layer[
-                            mgr.layer_offsets[layer_idx]]
-                        if hd not in self._vswa_head_dim_to_pool:
-                            self._vswa_head_dim_to_pool[hd] = pool_id
-
                 # Pre-allocate VSWA pool cache buffers.  These must be
                 # stable (never reallocated) so that CUDA-graph-recorded
                 # copies reference valid addresses across replays.
@@ -1325,22 +1319,62 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._host_paged_kv_indices = \
                 self._host_pool_indices[primary_pool_id]
 
+        # trtllm-gen paged-prefill graphs capture one stable block table per KV
+        # pool. Refresh those tables and KV lengths after request turnover.
+        if (self.is_cuda_graph and self.num_contexts > 0
+                and self._vswa_layer_to_pool is not None):
+            for plan_params, wrappers in self._plan_params_to_wrappers.items():
+                if plan_params.attention_mask_data is not None:
+                    continue
+                prefill_wrapper = wrappers.prefill_wrapper
+                if (prefill_wrapper is None
+                        or prefill_wrapper._backend != "trtllm-gen"):
+                    continue
+                block_tables = prefill_wrapper._block_tables
+                pool_id = plan_params.kv_pool_id
+                if block_tables is None or pool_id is None:
+                    continue
+                host_pool_indices = self._host_pool_indices[pool_id]
+                host_block_tables = wrappers.host_prefill_block_tables
+                if (host_block_tables is None
+                        or host_block_tables.shape != block_tables.shape):
+                    host_block_tables = torch.zeros(
+                        block_tables.shape,
+                        dtype=torch.int32,
+                        pin_memory=prefer_pinned(),
+                    )
+                    wrappers.host_prefill_block_tables = host_block_tables
+                else:
+                    host_block_tables.zero_()
+                source_offset = 0
+                for row, num_blocks_for_row in enumerate(
+                        num_blocks[:self.num_contexts]):
+                    copy_width = min(int(num_blocks_for_row),
+                                     block_tables.size(1))
+                    host_block_tables[row, :copy_width].copy_(
+                        host_pool_indices[source_offset:source_offset +
+                                          copy_width])
+                    source_offset += int(num_blocks_for_row)
+                block_tables.copy_(host_block_tables, non_blocking=True)
+                prefill_wrapper._kv_lens_buffer[:self.num_contexts].copy_(
+                    _to_int32_tensor(kv_lens_host[:self.num_contexts]),
+                    non_blocking=True,
+                )
+
         # CUDA graph + trtllm-gen: update _block_tables and _kv_lens_buffer
         # so the trtllm-gen decode kernel uses current page indices.
         if (self.is_cuda_graph and self._vswa_layer_to_pool is not None
                 and self._vswa_pool_indices_cache is not None
                 and self.num_generations > 0):
             decode_blocks = num_blocks[self.num_contexts:]
-            head_dim_to_pool = getattr(self, '_vswa_head_dim_to_pool', None)
             for plan_params, wrappers in self._plan_params_to_wrappers.items():
                 if plan_params.attention_mask_data is not None:
                     continue
                 decode_wrapper = wrappers.decode_wrapper
                 block_tables = getattr(decode_wrapper, '_block_tables', None)
                 if block_tables is None:
                     continue
-                pool_id = (head_dim_to_pool.get(plan_params.head_dim)
-                           if head_dim_to_pool else None)
+                pool_id = plan_params.kv_pool_id
                 if pool_id is None:
                     continue
                 batch_size, table_width = block_tables.shape
@@ -1439,6 +1473,7 @@ def plan(self,
             attention_mask_type=AttentionMaskType(attention_mask_type),
             attention_mask_data=attention_mask_data,
             multi_item_params=self._multi_item_params,
+            kv_pool_id=getattr(self, "_vswa_active_pool_id", None),
         )
         return self._plan_with_params(plan_params, flashinfer_backend)
 

tensorrt_llm/_torch/configs/__init__.py

@@ -1,6 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from tensorrt_llm._torch.configs.cosmos3 import Cosmos3Config
 from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config
 from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config
+from tensorrt_llm._torch.configs.gemma4 import Gemma4AssistantConfig
 from tensorrt_llm._torch.configs.gemma4_unified import (
     Gemma4UnifiedAudioConfig,
     Gemma4UnifiedConfig,
@@ -36,6 +52,7 @@ def _register_custom_configs_with_transformers() -> None:
         "deepseek_v32": DeepseekV3Config,
         "kimi_k2": DeepseekV3Config,
         "deepseek_v4": DeepseekV4Config,
+        "gemma4_assistant": Gemma4AssistantConfig,
         "laguna": LagunaConfig,
         "gemma4_unified": Gemma4UnifiedConfig,
         "gemma4_unified_text": Gemma4UnifiedTextConfig,
@@ -59,6 +76,7 @@ def _register_custom_configs_with_transformers() -> None:
     "Cosmos3Config",
     "DeepseekV3Config",
     "DeepseekV4Config",
+    "Gemma4AssistantConfig",
     "Gemma4UnifiedAudioConfig",
     "Gemma4UnifiedConfig",
     "Gemma4UnifiedTextConfig",