From c4d0a98ce989be4bc8823042b4292937c3ca8b46 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 1 Jul 2026 11:17:23 +0000
Subject: [PATCH 1/9] [None][feat] Add Gemma4 MTP assistant support

Add standalone Gemma4 assistant config, model, weight loading, and target KV-cache sharing for two-model MTP decoding. Support masked assistant logits and CUDA graph execution, and add structural coverage for assistant configuration and KV source mapping.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 examples/llm-api/quickstart_advanced.py       |   6 +-
 tensorrt_llm/_torch/configs/__init__.py       |  18 ++
 tensorrt_llm/_torch/configs/gemma4.py         |  61 ++++++
 .../checkpoints/hf/gemma4_weight_mapper.py    |   1 +
 tensorrt_llm/_torch/models/modeling_gemma4.py | 204 +++++++++++++++++-
 .../_torch/models/modeling_gemma4mm.py        |   1 +
 tensorrt_llm/_torch/pyexecutor/_util.py       |  17 ++
 .../_torch/pyexecutor/py_executor_creator.py  |   9 +-
 tensorrt_llm/_torch/speculative/utils.py      |   5 +-
 .../_torch/modeling/test_modeling_gemma4.py   |  67 ++++++
 10 files changed, 377 insertions(+), 12 deletions(-)
 create mode 100644 tensorrt_llm/_torch/configs/gemma4.py

diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py
index d2acd490d44a..081ed7138f5a 100644
--- a/examples/llm-api/quickstart_advanced.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py
index 031c2a534e2a..3bcf6fefdf5b 100644
--- a/tensorrt_llm/_torch/configs/__init__.py
+++ b/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",
diff --git a/tensorrt_llm/_torch/configs/gemma4.py b/tensorrt_llm/_torch/configs/gemma4.py
new file mode 100644
index 000000000000..85ccb7ec2347
--- /dev/null
+++ b/tensorrt_llm/_torch/configs/gemma4.py
@@ -0,0 +1,61 @@
+# 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 transformers import Gemma4TextConfig, PreTrainedConfig
+
+
+class Gemma4AssistantConfig(PreTrainedConfig):
+    """Compatibility config for Gemma4 assistant checkpoints.
+
+    Gemma4 assistant support postdates the transformers version currently
+    pinned by TensorRT-LLM. Keep the compatibility surface minimal and remove
+    this class once the pinned transformers release provides it natively.
+    """
+
+    model_type = "gemma4_assistant"
+    sub_configs = {"text_config": Gemma4TextConfig}
+
+    def __init__(
+        self,
+        text_config=None,
+        backbone_hidden_size=None,
+        use_ordered_embeddings=False,
+        num_centroids=0,
+        centroid_intermediate_top_k=0,
+        **kwargs,
+    ):
+        if text_config is None:
+            text_config = Gemma4TextConfig()
+        elif isinstance(text_config, dict):
+            text_config = Gemma4TextConfig(**text_config)
+
+        self.text_config = text_config
+        self.backbone_hidden_size = backbone_hidden_size
+        self.use_ordered_embeddings = use_ordered_embeddings
+        self.num_centroids = num_centroids
+        self.centroid_intermediate_top_k = centroid_intermediate_top_k
+        super().__init__(**kwargs)
+
+    @property
+    def hidden_size(self):
+        return self.text_config.hidden_size
+
+    @property
+    def vocab_size(self):
+        return self.text_config.vocab_size
+
+    @property
+    def num_hidden_layers(self):
+        return self.text_config.num_hidden_layers
diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
index e336ef4b894b..3f69b62c9fbb 100644
--- a/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
+++ b/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
@@ -28,6 +28,7 @@
 @register_mapper("HF", "Gemma4ForCausalLM")
 @register_mapper("HF", "Gemma4ForConditionalGeneration")
 @register_mapper("HF", "Gemma4UnifiedForConditionalGeneration")
+@register_mapper("HF", "Gemma4AssistantForCausalLM")
 class Gemma4HfWeightMapper(HfWeightMapper):
     @property
     def _is_vlm(self) -> bool:
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index bf356d5e3ba1..4d30615fa130 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -14,7 +14,9 @@
 # limitations under the License.
 """TensorRT-LLM PyTorch backend implementation for Gemma4 text model."""
 
+import dataclasses
 import math
+import weakref
 from typing import Dict, Optional, Tuple, Union
 
 import torch
@@ -55,6 +57,7 @@
 from ..modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from ..modules.rms_norm import RMSNorm
 from ..utils import ActivationType, Fp4QuantizedTensor, is_torch_compiling
+from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model
 
 _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0"
@@ -656,7 +659,7 @@ def __init__(
         # Determine if this is a KV-shared layer
         num_kv_shared = getattr(config, "num_kv_shared_layers", 0)
         first_kv_shared_layer_idx = config.num_hidden_layers - num_kv_shared
-        self.is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0
+        self.is_kv_shared_layer = num_kv_shared > 0 and layer_idx >= first_kv_shared_layer_idx
 
         # For shared layers, find the target layer to read KV cache from:
         # last non-shared layer of the same attention type (sliding/full).
@@ -1223,7 +1226,7 @@ def forward(
 # Gemma4 For Causal LM
 # ---------------------------------------------------------------------------
 @register_auto_model("Gemma4ForCausalLM")
-class Gemma4ForCausalLM(DecoderModelForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
+class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
     def __init__(
         self,
         model_config: ModelConfig[Gemma4TextConfig],
@@ -1243,12 +1246,7 @@ def __init__(
                 "moe_ep_size>1 requires a Gemma4 MoE variant (only 26B-A4B-it today)."
             )
 
-        super().__init__(
-            Gemma4TextModel(model_config),
-            config=model_config,
-            hidden_size=model_config.pretrained_config.hidden_size,
-            vocab_size=model_config.pretrained_config.vocab_size,
-        )
+        super().__init__(Gemma4TextModel(model_config), model_config)
 
     @classmethod
     def get_model_defaults(cls, llm_args) -> dict:
@@ -1426,3 +1424,193 @@ def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         # Ensure PLE nn.Linear modules match model dtype (weight loader may
         # not handle raw nn.Linear correctly, leaving them as float32).
         self.model._ensure_ple_dtype()
+
+
+class Gemma4AssistantMaskedEmbedder(nn.Module):
+    """Compute Gemma4 assistant logits for the selected centroid clusters."""
+
+    def __init__(self, model_config: ModelConfig):
+        super().__init__()
+        config = model_config.pretrained_config
+        self.hidden_size = config.hidden_size
+        self.num_centroids = config.num_centroids
+        self.centroid_intermediate_top_k = config.centroid_intermediate_top_k
+        self.vocab_size = config.vocab_size
+        self.vocab_size_per_centroid = self.vocab_size // self.num_centroids
+        self.centroids = Linear(
+            self.hidden_size,
+            self.num_centroids,
+            bias=False,
+            dtype=config.torch_dtype,
+        )
+        self.register_buffer(
+            "token_ordering",
+            torch.empty(self.vocab_size, dtype=torch.long, device="cuda"),
+        )
+
+    def forward(self, hidden_states: torch.Tensor, lm_head_weight: torch.Tensor) -> torch.Tensor:
+        centroid_logits = self.centroids(hidden_states)
+        _, top_k_indices = torch.topk(
+            centroid_logits,
+            k=self.centroid_intermediate_top_k,
+            dim=-1,
+        )
+        canonical_positions = self.token_ordering.view(
+            self.num_centroids, self.vocab_size_per_centroid
+        )[top_k_indices]
+        selected_embeddings = lm_head_weight[canonical_positions.reshape(-1)].view(
+            hidden_states.shape[0],
+            self.centroid_intermediate_top_k * self.vocab_size_per_centroid,
+            self.hidden_size,
+        )
+        selected_logits = torch.bmm(
+            hidden_states.unsqueeze(1), selected_embeddings.transpose(1, 2)
+        ).squeeze(1)
+        logits = torch.full(
+            (hidden_states.shape[0], self.vocab_size),
+            torch.finfo(hidden_states.dtype).min,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        return logits.scatter_(1, canonical_positions.flatten(1), selected_logits)
+
+
+@register_auto_model("Gemma4AssistantForCausalLM")
+class Gemma4AssistantForCausalLM(DecoderModelForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
+    """Gemma4 MTP assistant that attends directly to the target model KV cache."""
+
+    def __init__(self, model_config: ModelConfig):
+        assistant_config = model_config.pretrained_config
+        text_model_config = dataclasses.replace(
+            model_config,
+            pretrained_config=assistant_config.text_config,
+            spec_config=None,
+        )
+        super().__init__(
+            Gemma4TextModel(text_model_config),
+            config=model_config,
+            hidden_size=assistant_config.hidden_size,
+            vocab_size=assistant_config.vocab_size,
+        )
+        self.pre_projection = Linear(
+            2 * assistant_config.backbone_hidden_size,
+            assistant_config.hidden_size,
+            bias=False,
+            dtype=assistant_config.torch_dtype,
+        )
+        self.post_projection = Linear(
+            assistant_config.hidden_size,
+            assistant_config.backbone_hidden_size,
+            bias=False,
+            dtype=assistant_config.torch_dtype,
+        )
+        self.masked_embedding = (
+            Gemma4AssistantMaskedEmbedder(model_config)
+            if assistant_config.use_ordered_embeddings
+            else None
+        )
+        self._target_embed_tokens_ref = None
+
+    @classmethod
+    def get_model_defaults(cls, llm_args) -> dict:
+        return {"attn_backend": "FLASHINFER"}
+
+    def load_weights_from_target_model(self, target_model: nn.Module) -> None:
+        target_llm = target_model.llm if hasattr(target_model, "llm") else target_model
+        self._target_embed_tokens_ref = weakref.ref(target_llm.model.embed_tokens)
+
+        target_config = target_llm.config
+        num_kv_shared = getattr(target_config, "num_kv_shared_layers", 0)
+        num_source_layers = target_config.num_hidden_layers - num_kv_shared
+        source_layer_types = target_config.layer_types[:num_source_layers]
+        for layer in self.model.layers:
+            layer_type = "sliding_attention" if layer.is_sliding else "full_attention"
+            if layer_type not in source_layer_types:
+                raise ValueError(f"Target Gemma4 model has no KV source for {layer_type}")
+            source_layer_idx = (
+                len(source_layer_types) - 1 - source_layer_types[::-1].index(layer_type)
+            )
+            layer.self_attn.attn.layer_idx = source_layer_idx
+
+    def _get_target_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if self._target_embed_tokens_ref is None:
+            raise RuntimeError("Gemma4 assistant target embeddings have not been initialized")
+        target_embed_tokens = self._target_embed_tokens_ref()
+        if target_embed_tokens is None:
+            raise RuntimeError("Gemma4 assistant target embedding reference is no longer valid")
+        return target_embed_tokens(input_ids)
+
+    @staticmethod
+    def _constant_position_ids(
+        position_ids: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        positions = position_ids.squeeze(0) if position_ids.ndim == 2 else position_ids
+        seq_lens = attn_metadata.seq_lens_cuda[: attn_metadata.num_seqs]
+        last_token_indices = torch.cumsum(seq_lens, dim=0, dtype=torch.long) - 1
+        return torch.repeat_interleave(
+            positions[last_token_indices],
+            seq_lens,
+            output_size=positions.shape[0],
+        ).unsqueeze(0)
+
+    @staticmethod
+    def _last_token_states(
+        hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        last_tokens = (
+            torch.cumsum(
+                attn_metadata.seq_lens_cuda,
+                dim=0,
+                dtype=torch.long,
+            )
+            - 1
+        )
+        return hidden_states[last_tokens]
+
+    def forward(
+        self,
+        attn_metadata: AttentionMetadata,
+        input_ids: torch.IntTensor = None,
+        position_ids: Optional[torch.IntTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        return_context_logits: bool = False,
+        spec_metadata=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if input_ids is None or spec_metadata is None:
+            raise ValueError("Gemma4 assistant requires input_ids and speculative metadata")
+        target_hidden_states = spec_metadata.get_hidden_states()
+        target_embeddings = self._get_target_embeddings(input_ids)
+        assistant_inputs = self.pre_projection(
+            torch.cat([target_embeddings, target_hidden_states], dim=-1)
+        )
+        assistant_hidden_states = self.model(
+            attn_metadata=attn_metadata,
+            position_ids=self._constant_position_ids(position_ids, attn_metadata),
+            inputs_embeds=assistant_inputs,
+            spec_metadata=spec_metadata,
+        )
+
+        projected_hidden_states = self.post_projection(assistant_hidden_states)
+        spec_metadata.maybe_capture_hidden_states(
+            self.config.num_hidden_layers - 1,
+            projected_hidden_states,
+        )
+
+        if return_context_logits:
+            logits_hidden_states = assistant_hidden_states
+        else:
+            logits_hidden_states = self._last_token_states(assistant_hidden_states, attn_metadata)
+        if self.masked_embedding is not None:
+            return self.masked_embedding(logits_hidden_states, self.lm_head.weight).float()
+        return self.lm_head(logits_hidden_states).float()
+
+    def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
+        weights = weight_mapper.preprocess_weights(weights)
+        token_ordering_key = "masked_embedding.token_ordering"
+        if self.masked_embedding is not None and token_ordering_key in weights:
+            self.masked_embedding.token_ordering.copy_(weights[token_ordering_key])
+            del weights[token_ordering_key]
+        super().load_weights(weights, weight_mapper)
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 0f5a340c0a9a..bf295489f971 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -1098,6 +1098,7 @@ def forward(
             mm_token_type_ids=mm_token_type_ids,
             ple_input_ids=ple_input_ids,
             lora_params=kwargs.get("lora_params", None),
+            spec_metadata=kwargs.get("spec_metadata", None),
         )
         return logits
 
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 0149ba7ae03a..541fc752b628 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1617,6 +1617,23 @@ def build_managers(self,
                                        if draft_kv_cache_config is not None else
                                        self_kv_cache_config)
 
+        # Gemma4 assistants read the target model's KV cache directly. They
+        # still need a small draft manager for request/slot bookkeeping, but
+        # must not reserve a second full-size GPU cache during final capacity
+        # allocation.
+        if (self._draft_model_engine is not None
+                and self._draft_model_engine.kv_cache_manager_key
+                == ResourceManagerType.KV_CACHE_MANAGER):
+            draft_build_kv_cache_config = copy.deepcopy(
+                draft_build_kv_cache_config)
+            draft_build_kv_cache_config.max_gpu_total_bytes = 0
+            draft_build_kv_cache_config.free_gpu_memory_fraction = None
+            draft_build_kv_cache_config.host_cache_size = 0
+            draft_build_kv_cache_config.max_tokens = max(
+                self._max_num_tokens,
+                self._max_seq_len * self._max_batch_size,
+            )
+
         # Two-model speculative decoding: draft model has separate engine
         if self._draft_model_engine is not None:
             if self._is_kv_cache_manager_v2:
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index f82207881c8f..a53d0ec4888d 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -633,11 +633,16 @@ def drafting_loop_wrapper(model):
                 model_weights_memory_tag=model_weights_memory_tag,
                 model_weights_restore_mode=model_weights_restore_mode,
             )
-            # For DeepseekV3 MTP, we need to set the num_hidden_layers to 1 for the draft model
-            if spec_config.spec_dec_mode.is_mtp_eagle():
+            # Embedded MTP checkpoints expose a single draft layer. Standalone
+            # Gemma4 assistants keep their full four-layer text backbone.
+            if (spec_config.spec_dec_mode.is_mtp_eagle()
+                    and draft_model_engine.model.config.model_type
+                    != "gemma4_assistant"):
                 draft_model_engine.model.model_config.pretrained_config.num_hidden_layers = 1
             draft_model_engine.load_weights_from_target_model(
                 model_engine.model)
+            if draft_model_engine.model.config.model_type == "gemma4_assistant":
+                draft_model_engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
     else:
         draft_model_engine = None
 
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 62189035cafb..c5c29fe36cbb 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -130,13 +130,16 @@ def get_spec_metadata(spec_config,
             draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_mtp_eagle():
+        hidden_size = model_config.hidden_size
+        if model_config.model_type == "gemma4_assistant":
+            hidden_size = model_config.backbone_hidden_size
         return Eagle3SpecMetadata(
             max_draft_len=spec_config.max_draft_len,
             max_total_draft_tokens=spec_config.tokens_per_gen_step - 1,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
             num_layers=model_config.num_hidden_layers,
-            hidden_size=model_config.hidden_size,
+            hidden_size=hidden_size,
             max_num_tokens=max_num_tokens,
             dtype=model_config.torch_dtype,
             is_draft_model=is_draft_model,
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index acd7f2d593f5..5d8a69032020 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -28,10 +28,12 @@
 from transformers import Gemma4Config, Gemma4TextConfig
 
 from tensorrt_llm._torch.attention_backend import FlashInferAttention, FlashInferAttentionMetadata
+from tensorrt_llm._torch.configs.gemma4 import Gemma4AssistantConfig
 from tensorrt_llm._torch.metadata import KVCacheParams
 from tensorrt_llm._torch.model_config import ModelConfig
 from tensorrt_llm._torch.models.checkpoints.hf.gemma4_weight_mapper import Gemma4HfWeightMapper
 from tensorrt_llm._torch.models.modeling_gemma4 import (
+    Gemma4AssistantForCausalLM,
     Gemma4Attention,
     Gemma4DecoderLayer,
     Gemma4ForCausalLM,
@@ -104,6 +106,26 @@
     "use_double_wide_mlp": True,
 }
 
+GEMMA4_ASSISTANT_CONFIG = {
+    "text_config": {
+        **GEMMA4_SMALL_CONFIG,
+        "num_hidden_layers": 4,
+        "layer_types": [
+            "sliding_attention",
+            "sliding_attention",
+            "sliding_attention",
+            "full_attention",
+        ],
+        "num_kv_shared_layers": 4,
+    },
+    "backbone_hidden_size": 256,
+    "use_ordered_embeddings": True,
+    "num_centroids": 16,
+    "centroid_intermediate_top_k": 2,
+    "tie_word_embeddings": True,
+    "dtype": "bfloat16",
+}
+
 
 def _make_model_config(config_dict):
     """Build a ModelConfig from a raw config dict."""
@@ -112,6 +134,13 @@ def _make_model_config(config_dict):
     return ModelConfig(pretrained_config=cfg, mapping=mapping)
 
 
+def _make_assistant_model_config(config_dict=GEMMA4_ASSISTANT_CONFIG):
+    """Build a ModelConfig for a standalone Gemma4 assistant."""
+    cfg = Gemma4AssistantConfig(**deepcopy(config_dict))
+    mapping = Mapping(world_size=1, tp_size=1, rank=0)
+    return ModelConfig(pretrained_config=cfg, mapping=mapping)
+
+
 class TestGemma4Config(unittest.TestCase):
     """Tests for Gemma4 config classes."""
 
@@ -445,6 +474,44 @@ def test_reject_unrecognized_expert_weights(self):
             Gemma4HfWeightMapper()._remap_moe_keys(weights)
 
 
+class TestGemma4Assistant(unittest.TestCase):
+    """Structural tests for standalone Gemma4 MTP assistants."""
+
+    def test_assistant_config_wraps_text_config(self):
+        config = Gemma4AssistantConfig(**deepcopy(GEMMA4_ASSISTANT_CONFIG))
+
+        self.assertEqual(config.model_type, "gemma4_assistant")
+        self.assertIsInstance(config.text_config, Gemma4TextConfig)
+        self.assertEqual(config.hidden_size, 256)
+        self.assertEqual(config.vocab_size, 1024)
+        self.assertEqual(config.num_hidden_layers, 4)
+
+    def test_assistant_uses_target_kv_sources(self):
+        assistant = Gemma4AssistantForCausalLM(_make_assistant_model_config())
+        self.assertEqual(len(assistant.model.layers), 4)
+        self.assertTrue(all(layer.is_kv_shared_layer for layer in assistant.model.layers))
+
+        target_config = {
+            **GEMMA4_SMALL_CONFIG,
+            "layer_types": [
+                "sliding_attention",
+                "full_attention",
+                "sliding_attention",
+                "full_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            "num_kv_shared_layers": 2,
+        }
+        target = Gemma4ForCausalLM(_make_model_config(target_config))
+        assistant.load_weights_from_target_model(target)
+
+        expected_source_layers = [2, 2, 2, 3]
+        actual_source_layers = [layer.self_attn.attn.layer_idx for layer in assistant.model.layers]
+        self.assertEqual(actual_source_layers, expected_source_layers)
+        self.assertIs(assistant._target_embed_tokens_ref(), target.model.embed_tokens)
+
+
 # ---------------------------------------------------------------------------
 # HF reference comparison tests (sub-module + full model)
 # ---------------------------------------------------------------------------

From b53924eaa0bbcf175c7fc99f5f7ea40c31817f81 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:13:41 +0000
Subject: [PATCH 2/9] [None][docs] Document Gemma4 MTP support

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 docs/source/models/supported-models.md |  4 ++-
 examples/models/core/gemma/README.md   | 48 ++++++++++++++++++++++----
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index b03f9961f7cc..848372db456d 100644
--- a/docs/source/models/supported-models.md
+++ b/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` [^14]   | 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 [^14] | 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        |
@@ -92,6 +93,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism.
 [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
+[^14]: Gemma 4 uses two-model MTP on the PyTorch backend with a matching `Gemma4AssistantForCausalLM` checkpoint. See the [Gemma 4 example](../../../examples/models/core/gemma/README.md#mtp-speculative-decoding). Two-model MTP is deprecated and scheduled for removal in release 1.4. The Gemma 4 12B target and assistant use the `Gemma4UnifiedForConditionalGeneration` and `Gemma4UnifiedAssistantForCausalLM` architectures, which are not supported.
 
 # Multimodal Feature Support Matrix (PyTorch Backend)
 
diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md
index 1315755e85b4..c80a9ab5c9c2 100644
--- a/examples/models/core/gemma/README.md
+++ b/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:

From 1863db44aee6a0e7e69177c3b051248399ac0ba3 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:17:34 +0000
Subject: [PATCH 3/9] [None][docs] Remove Gemma4 MTP footnote

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 docs/source/models/supported-models.md | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index 848372db456d..b68d447487c9 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -20,7 +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` [^14]   | 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` |
+| `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 +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             | Yes [^14] | 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        |
@@ -93,8 +93,6 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism.
 [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
-[^14]: Gemma 4 uses two-model MTP on the PyTorch backend with a matching `Gemma4AssistantForCausalLM` checkpoint. See the [Gemma 4 example](../../../examples/models/core/gemma/README.md#mtp-speculative-decoding). Two-model MTP is deprecated and scheduled for removal in release 1.4. The Gemma 4 12B target and assistant use the `Gemma4UnifiedForConditionalGeneration` and `Gemma4UnifiedAssistantForCausalLM` architectures, which are not supported.
-
 # Multimodal Feature Support Matrix (PyTorch Backend)
 
 | Model Architecture/Feature           | Overlap Scheduler | CUDA Graph | Chunked Prefill | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Logits Post Processor | EPD Disaggregated Serving | Modality  |

From 426c58a71b724ad4a8aec6d5fb24479547e091f9 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:19:20 +0000
Subject: [PATCH 4/9] [None][docs] Restore supported models spacing

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 docs/source/models/supported-models.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index b68d447487c9..dfa283beecde 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -93,6 +93,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism.
 [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
+
 # Multimodal Feature Support Matrix (PyTorch Backend)
 
 | Model Architecture/Feature           | Overlap Scheduler | CUDA Graph | Chunked Prefill | Torch Sampler | TLLM C++ Sampler | KV Cache Reuse | Logits Post Processor | EPD Disaggregated Serving | Modality  |

From b98988996811190c4cab368714d4d6e02ce5993e Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Thu, 2 Jul 2026 15:25:30 +0000
Subject: [PATCH 5/9] [None][fix] Avoid Gemma4 Bandit false positive

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 4d30615fa130..5e131dc9bf88 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1609,8 +1609,8 @@ def forward(
 
     def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         weights = weight_mapper.preprocess_weights(weights)
-        token_ordering_key = "masked_embedding.token_ordering"
-        if self.masked_embedding is not None and token_ordering_key in weights:
-            self.masked_embedding.token_ordering.copy_(weights[token_ordering_key])
-            del weights[token_ordering_key]
+        ordering_weight_name = "masked_embedding.token_ordering"
+        if self.masked_embedding is not None and ordering_weight_name in weights:
+            self.masked_embedding.token_ordering.copy_(weights[ordering_weight_name])
+            del weights[ordering_weight_name]
         super().load_weights(weights, weight_mapper)

From 2e5f3f52ee785d7def2d937a89c375deaac357a7 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 21 Jul 2026 12:33:48 +0000
Subject: [PATCH 6/9] [None][fix] Fix Gemma4 MTP decoding correctness

Preserve Gemma4 assistant hidden-state and position semantics across
multi-token drafting, and capture target hidden states for verification.
Refresh FlashInfer paged-prefill graph state after request turnover and
add regression coverage for drafting and CUDA graph replay.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    |  47 ++++++
 tensorrt_llm/_torch/models/modeling_gemma4.py |   7 +
 .../_torch/pyexecutor/model_engine.py         |  28 +++-
 .../_torch/pyexecutor/py_executor_creator.py  |  20 ++-
 .../_torch/speculative/drafting_loops.py      |  43 ++++++
 tensorrt_llm/_torch/speculative/eagle3.py     |  31 +++-
 .../_torch/speculative/model_drafter.py       |  44 +++++-
 tensorrt_llm/_torch/speculative/utils.py      |   1 +
 .../_torch/modeling/test_modeling_gemma4.py   | 139 +++++++++++++++++
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 143 ++++++++++++++++++
 10 files changed, 486 insertions(+), 17 deletions(-)
 create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 733d477009cb..375703ce5e49 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -225,6 +225,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)
@@ -1325,6 +1327,51 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._host_paged_kv_indices = \
                 self._host_pool_indices[primary_pool_id]
 
+        # Gemma4's trtllm-gen paged-prefill graphs capture one stable block
+        # table per head dimension. Refresh those tables and KV lengths after
+        # request turnover instead of replaying with graph-warmup page IDs.
+        if (self.is_cuda_graph and self.num_contexts > 0
+                and self._vswa_layer_to_pool is not None):
+            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
+                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 = (head_dim_to_pool.get(plan_params.head_dim)
+                           if head_dim_to_pool else None)
+                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
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 5e131dc9bf88..3a19d468c302 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -56,6 +56,7 @@
 from ..modules.gemma4.fused_qkv import gemma4_fused_qkv_norm_rope_quant
 from ..modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
 from ..modules.rms_norm import RMSNorm
+from ..speculative.interface import SpecMetadata
 from ..utils import ActivationType, Fp4QuantizedTensor, is_torch_compiling
 from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model
@@ -1378,6 +1379,7 @@ def forward(
         inputs_embeds: Optional[torch.FloatTensor] = None,
         return_context_logits: bool = False,
         mm_token_type_ids: Optional[torch.Tensor] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         local_attention_mask_data = None
@@ -1404,6 +1406,11 @@ def forward(
             **kwargs,
         )
 
+        if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+            spec_metadata.maybe_capture_hidden_states(self.layer_idx, output)
+        if attn_metadata.padded_num_tokens is not None:
+            output = output[: attn_metadata.num_tokens]
+
         logits = self.logits_processor.forward(
             output,
             self.lm_head,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 50aa317e9e43..822fc6931b4c 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -281,8 +281,8 @@ def __init__(
         dist: Optional[Distributed] = None,
         spec_config: Optional[DecodingBaseConfig] = None,
         is_draft_model: bool = False,
-        drafting_loop_wrapper: Optional[Callable[[torch.nn.Module],
-                                                 torch.nn.Module]] = None,
+        drafting_loop_wrapper: Optional[Callable[
+            [torch.nn.Module], Optional[torch.nn.Module]]] = None,
         model: Optional[torch.nn.Module] = None,
         checkpoint_loader: Optional[BaseCheckpointLoader] = None,
         model_weights_memory_tag: Optional[str] = None,
@@ -431,8 +431,12 @@ def __init__(
             enable_overlap_headroom=self._enable_dsv4_overlap_headroom,
         )
         if drafting_loop_wrapper is not None:
-            self.model = drafting_loop_wrapper(self.model)
-            self.model_is_wrapped = True
+            wrapped_model = drafting_loop_wrapper(self.model)
+            if wrapped_model is not None:
+                self.model = wrapped_model
+                self.model_is_wrapped = True
+            else:
+                self.model_is_wrapped = False
         else:
             self.model_is_wrapped = False
         self.sparse_attention_config = self.model.model_config.sparse_attention_config
@@ -1469,8 +1473,12 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
         logger.info("Running autotuner warmup...")
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
-        token_num_upper_bound = min(self.max_num_tokens,
-                                    self.batch_size * (self.max_seq_len - 1))
+        if (self.is_draft_model and self.model_is_wrapped
+                and self.model.config.model_type == "gemma4_assistant"):
+            token_num_upper_bound = 1
+        else:
+            token_num_upper_bound = min(
+                self.max_num_tokens, self.batch_size * (self.max_seq_len - 1))
         curr_max_num_tokens = kv_cache_manager.get_num_available_tokens(
             token_num_upper_bound=token_num_upper_bound,
             max_num_draft_tokens=self.original_max_draft_len)
@@ -2405,10 +2413,14 @@ def _update_draft_inference_state_for_warmup(
             ResourceManagerType.SPEC_RESOURCE_MANAGER)
         if self.is_draft_model and isinstance(spec_resource_manager,
                                               Eagle3ResourceManager):
-            spec_resource_manager.is_first_draft = is_first_draft
+            is_gemma4_assistant = (self.model_is_wrapped
+                                   and self.model.config.model_type
+                                   == "gemma4_assistant")
+            spec_resource_manager.is_first_draft = (is_first_draft
+                                                    and not is_gemma4_assistant)
             if is_first_draft:
                 for req in batch.generation_requests:
-                    req.py_is_first_draft = True
+                    req.py_is_first_draft = not is_gemma4_assistant
                     req.py_draft_tokens = []
 
     def _set_up_attn_metadata(
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index a53d0ec4888d..7602aafb62dc 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -578,22 +578,34 @@ def allocation_scope(current_stage: ExecutorMemoryType):
         with allocation_scope(ExecutorMemoryType.MODEL_ENGINE_DRAFT):
             draft_spec_config = copy.copy(spec_config)
 
-            use_chain_drafter = (
+            capturable_drafter_eligible = (
                 guided_decoding_config is None
-                and draft_spec_config._allow_chain_drafter
                 and draft_spec_config._allow_greedy_draft_tokens
-                and llm_args.attn_backend == "TRTLLM"
                 and draft_spec_config.draft_len_schedule is None)
+            use_chain_drafter = (capturable_drafter_eligible
+                                 and draft_spec_config._allow_chain_drafter
+                                 and llm_args.attn_backend == "TRTLLM")
 
             logger.debug(f"USE CHAIN DRAFTER: {use_chain_drafter}")
-            if use_chain_drafter:
+            if (capturable_drafter_eligible
+                    and (use_chain_drafter
+                         or draft_spec_config.spec_dec_mode.is_mtp_eagle())):
 
                 def drafting_loop_wrapper(model):
                     from tensorrt_llm._torch.speculative.drafting_loops import (
+                        Gemma4AssistantDraftingLoopWrapper,
                         LinearDraftingLoopWrapper,
                         StaticTreeDraftingLoopWrapper)
                     from tensorrt_llm.llmapi import EagleDecodingConfig
 
+                    if model.config.model_type == "gemma4_assistant":
+                        return Gemma4AssistantDraftingLoopWrapper(
+                            spec_config.max_draft_len,
+                            spec_config.tokens_per_gen_step - 1, model)
+
+                    if not use_chain_drafter:
+                        return None
+
                     static_tree_drafter = isinstance(
                         draft_spec_config, EagleDecodingConfig
                     ) and draft_spec_config.eagle_choices is not None
diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py
index e2be40bfce5e..23673059a2a6 100644
--- a/tensorrt_llm/_torch/speculative/drafting_loops.py
+++ b/tensorrt_llm/_torch/speculative/drafting_loops.py
@@ -1,3 +1,5 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
 """
 This module contains capturable drafting loops for speculative decoding.
 
@@ -202,6 +204,47 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata,
         return new_position_ids
 
 
+class Gemma4AssistantDraftingLoopWrapper(LinearDraftingLoopWrapper):
+    """Draft tokens without advancing the target KV cache or position."""
+
+    def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
+                attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata,
+                **kwargs) -> dict[str, torch.Tensor]:
+        logits = self.draft_model.forward(input_ids=input_ids,
+                                          position_ids=position_ids,
+                                          attn_metadata=attn_metadata,
+                                          spec_metadata=spec_metadata,
+                                          return_context_logits=True)
+        logits = logits[spec_metadata.gather_ids]
+
+        new_draft_tokens = [self.sample(logits)]
+        draft_logits = [logits]
+        if self.max_draft_len > 1:
+            if not isinstance(spec_metadata, Eagle3SpecMetadata):
+                raise TypeError("Gemma4 assistant requires Eagle3 metadata")
+            batch_size = attn_metadata.num_seqs
+            spec_metadata.hidden_states_read_indices[:batch_size].copy_(
+                spec_metadata.hidden_states_write_indices[:batch_size])
+            for _ in range(self.max_draft_len - 1):
+                logits = self.draft_model.forward(
+                    input_ids=new_draft_tokens[-1],
+                    position_ids=position_ids,
+                    attn_metadata=attn_metadata,
+                    spec_metadata=spec_metadata)
+                new_draft_tokens.append(self.sample(logits))
+                draft_logits.append(logits)
+
+        return {
+            "new_draft_tokens": torch.stack(new_draft_tokens),
+            "draft_logits": torch.stack(draft_logits),
+        }
+
+    def prepare_for_generation(self, attn_metadata: AttentionMetadata,
+                               spec_metadata: SpecMetadata,
+                               position_ids: torch.Tensor) -> torch.Tensor:
+        return position_ids
+
+
 class StaticTreeDraftingLoopWrapper(BaseDraftingLoopWrapper):
 
     def __init__(self, max_draft_len: int, max_total_draft_tokens: int,
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index a1f4f82869bd..06e337177d37 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Optional, Set
 
@@ -88,6 +91,10 @@ def __init__(self,
         self.start_indices = {i: 0 for i in range(slot_size)}
         # whether the next draft forward is the first
         self.is_first_draft = True
+        # Gemma4 assistants share the target KV cache and only query the last
+        # validated position. The drafter records that position in the target
+        # model's most recent hidden-state span before preparing draft inputs.
+        self.draft_hidden_state_offsets: Dict[int, int] = {}
         self.spec_tree_manager = None
 
         if isinstance(config,
@@ -122,6 +129,7 @@ def free_resources(self, request: LlmRequest):
         slot_id = self.slot_manager.get_slot(request.request_id)
         self.seq_lens[slot_id] = 0
         self.start_indices[slot_id] = 0
+        self.draft_hidden_state_offsets.pop(request.request_id, None)
         if self.use_relaxed_acceptance_for_thinking:
             self.relaxed_delta_pool[slot_id].fill_(0)
         self.slot_manager.remove_slot(request.request_id)
@@ -207,6 +215,7 @@ class Eagle3SpecMetadata(SpecMetadata):
     is_first_draft: bool = False
     eagle3_resource_manager: Optional[Eagle3ResourceManager] = None
     is_mtp_eagle: bool = False
+    is_gemma4_assistant: bool = False
 
     eagle_choices: Optional[List[List[int]]] = None
     max_total_draft_tokens: int = 0
@@ -274,11 +283,29 @@ def prepare(self):
         for req_id, seq_len in zip(self.request_ids, self.seq_lens):
             slot_id = self.eagle3_resource_manager.slot_manager.get_slot(req_id)
             start_idx = self.eagle3_resource_manager.start_indices[slot_id]
+            # Gemma4 assistants issue one query per target iteration and reuse
+            # the target KV cache. Read the hidden state for the last validated
+            # target token, then overwrite that location with the projected
+            # assistant state for the remaining draft iterations.
+            if self.is_draft_model and self.is_gemma4_assistant:
+                assert seq_len == 1, (
+                    "Gemma4 assistant drafting expects one query token per "
+                    f"request, got {seq_len}")
+                old_seq_len = self.eagle3_resource_manager.seq_lens[slot_id]
+                hidden_state_offset = self.eagle3_resource_manager.draft_hidden_state_offsets.get(
+                    req_id, max(old_seq_len - 1, 0))
+                assert old_seq_len == 0 or 0 <= hidden_state_offset < old_seq_len, (
+                    "Gemma4 assistant hidden-state offset is outside the "
+                    f"target span: offset={hidden_state_offset}, "
+                    f"target_seq_len={old_seq_len}")
+                hidden_state_idx = start_idx + hidden_state_offset
+                hidden_states_read_indices.append(hidden_state_idx)
+                hidden_states_write_indices.append(hidden_state_idx)
             # 1) target model or (is_first_draft and is_linear_tree)
             # If this is the first draft or the target model forward, we need to
             # read/write all of the hidden states
-            if not self.is_draft_model or (is_first_draft
-                                           and spec_tree_manager is None):
+            elif not self.is_draft_model or (is_first_draft
+                                             and spec_tree_manager is None):
                 hidden_states_read_indices.extend(
                     list(range(start_idx, start_idx + seq_len)))
                 hidden_states_write_indices.extend(
diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py
index 5eae9b7e44cc..562e0a10e2da 100644
--- a/tensorrt_llm/_torch/speculative/model_drafter.py
+++ b/tensorrt_llm/_torch/speculative/model_drafter.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 from __future__ import annotations
 
 import traceback
@@ -93,6 +96,10 @@ def __init__(
         self.guided_decoder = guided_decoder
 
         self.use_static_draft_loop = draft_model_engine.model_is_wrapped
+        draft_model = (draft_model_engine.model.draft_model if
+                       self.use_static_draft_loop else draft_model_engine.model)
+        self.is_gemma4_assistant = getattr(draft_model.config, "model_type",
+                                           None) == "gemma4_assistant"
         if self.use_static_draft_loop:
             # TODO: enable sampling/guided decoding on static draft loop
             assert guided_decoder is None
@@ -163,6 +170,27 @@ def _create_generation_request(self, request: LlmRequest,
         new_request.state = LlmRequestState.GENERATION_IN_PROGRESS
         return new_request
 
+    def _create_gemma4_assistant_request(self, request: LlmRequest,
+                                         input_tokens: List[int],
+                                         is_first_draft: bool) -> LlmRequest:
+        """Create a one-token query over the target model's existing KV cache."""
+        new_request = self._create_generation_request(request, input_tokens)
+        if self.spec_resource_manager is None or not hasattr(
+                self.spec_resource_manager, "draft_hidden_state_offsets"):
+            raise RuntimeError(
+                "Gemma4 assistant requires an Eagle3 resource manager")
+        if is_first_draft:
+            slot_id = self.spec_resource_manager.slot_manager.get_slot(
+                request.py_request_id)
+            hidden_state_offset = self.spec_resource_manager.seq_lens[
+                slot_id] - 1
+        else:
+            hidden_state_offset = request.py_num_accepted_draft_tokens
+
+        self.spec_resource_manager.draft_hidden_state_offsets[
+            request.py_request_id] = hidden_state_offset
+        return new_request
+
     def _create_accepted_tokens_request(self, request: LlmRequest,
                                         input_tokens: Any,
                                         num_accepted_tokens: int) -> LlmRequest:
@@ -223,6 +251,14 @@ def _create_draft_request_for_request(
         num_draft_tokens, num_accepted_tokens = self._initialize_draft_tokens(
             request)
 
+        # First time seeing this request - context request
+        num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1
+        is_first_draft = (request.max_beam_num_tokens - 1 +
+                          num_overlap_tokens == request.py_prompt_len)
+        if self.is_gemma4_assistant:
+            return self._create_gemma4_assistant_request(
+                request, list(request.get_tokens(0)), is_first_draft)
+
         input_tokens = get_draft_model_prompt(self.spec_config.spec_dec_mode,
                                               request,
                                               self.disable_overlap_scheduler)
@@ -230,9 +266,7 @@ def _create_draft_request_for_request(
         is_eagle_style = self.spec_config.spec_dec_mode.is_eagle3(
         ) or self.spec_config.spec_dec_mode.is_mtp_eagle()
 
-        # First time seeing this request - context request
-        num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1
-        if request.max_beam_num_tokens - 1 + num_overlap_tokens == request.py_prompt_len:
+        if is_first_draft:
             # This is the first time the draft model is seeing this request.
             # Prepare a context request. We discard the first token and take
             # the newly decoded one - this is the convention for EAGLE 2 and 3.
@@ -301,6 +335,10 @@ def _prepare_draft_batch(
             for request in scheduled_requests.context_requests:
                 if request.py_disable_speculative_decoding:
                     continue
+                if self.is_gemma4_assistant:
+                    # The assistant has no private KV cache to populate during
+                    # chunked prefill. Drafting starts after target prefill.
+                    continue
                 if request.is_first_context_chunk:
                     # Ignore requests which still need to be processed by the target model.
                     continue
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index c5c29fe36cbb..b0f0382f7ff1 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -146,6 +146,7 @@ def get_spec_metadata(spec_config,
             eagle3_resource_manager=spec_resource_manager,
             layers_to_capture=None,
             is_mtp_eagle=True,
+            is_gemma4_assistant=model_config.model_type == "gemma4_assistant",
         )
     if spec_config.spec_dec_mode.is_eagle3():
         effective_dynamic_tree = _is_effective_dynamic_tree(spec_config)
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 5d8a69032020..8276807322ee 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -2804,6 +2804,145 @@ def test_cuda_graph_decode_hybrid_headdim(self):
 
         kv_cache_manager.shutdown()
 
+    @torch.no_grad()
+    @unittest.mock.patch(
+        "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
+    )
+    def test_cuda_graph_speculative_verification_hybrid_headdim(self):
+        """Speculative graph refreshes paged-prefill state after request turnover."""
+        config = Gemma4TextConfig(**deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG))
+        kv_cache_manager = self._get_kv_cache_manager(
+            config, num_blocks=32, tokens_per_block=32, batch_size=2
+        )
+        self.addCleanup(kv_cache_manager.shutdown)
+
+        capture_request_ids = [0]
+        replay_request_ids = [1]
+        capture_cached_tokens = [96]
+        replay_cached_tokens = [48]
+        verification_tokens = 6
+        capture_requests = kv_cache_manager.add_dummy_requests(
+            capture_request_ids,
+            [capture_cached_tokens[0] + verification_tokens],
+        )
+        replay_requests = kv_cache_manager.add_dummy_requests(
+            replay_request_ids,
+            [replay_cached_tokens[0] + verification_tokens],
+        )
+        self.assertIsNotNone(capture_requests)
+        self.assertIsNotNone(replay_requests)
+
+        layer_indices = [
+            config.layer_types.index("sliding_attention"),
+            config.layer_types.index("full_attention"),
+        ]
+        layers = []
+        queries = []
+        keys = []
+        values = []
+        for layer_idx in layer_indices:
+            is_sliding = config.layer_types[layer_idx] == "sliding_attention"
+            head_dim = config.head_dim if is_sliding else config.global_head_dim
+            layers.append(
+                FlashInferAttention(
+                    layer_idx=layer_idx,
+                    num_heads=config.num_attention_heads,
+                    head_dim=head_dim,
+                    num_kv_heads=config.num_key_value_heads,
+                    flashinfer_backend="trtllm-gen",
+                )
+            )
+            queries.append(
+                torch.randn(
+                    verification_tokens,
+                    config.num_attention_heads * head_dim,
+                    dtype=config.torch_dtype,
+                    device="cuda",
+                )
+            )
+            keys.append(
+                torch.randn(
+                    verification_tokens,
+                    config.num_key_value_heads * head_dim,
+                    dtype=config.torch_dtype,
+                    device="cuda",
+                )
+            )
+            values.append(torch.randn_like(keys[-1]))
+            torch.nn.init.normal_(kv_cache_manager.get_buffers(layer_idx))
+
+        def make_metadata(
+            *,
+            is_cuda_graph: bool,
+            request_ids: list[int],
+            cached_tokens: list[int],
+        ):
+            return FlashInferAttentionMetadata(
+                seq_lens=torch.tensor([verification_tokens], dtype=torch.int),
+                num_contexts=1,
+                is_cuda_graph=is_cuda_graph,
+                kv_cache_params=KVCacheParams(
+                    use_cache=True, num_cached_tokens_per_seq=cached_tokens
+                ),
+                workspace_buffer=(
+                    torch.empty(_FLASHINFER_WORKSPACE_BYTES, dtype=torch.uint8, device="cuda")
+                    if is_cuda_graph
+                    else None
+                ),
+                max_num_requests=1,
+                max_num_tokens=verification_tokens,
+                kv_cache_manager=kv_cache_manager,
+                request_ids=request_ids,
+            )
+
+        graph_metadata = make_metadata(
+            is_cuda_graph=True,
+            request_ids=capture_request_ids,
+            cached_tokens=capture_cached_tokens,
+        )
+        graph_metadata.prepare()
+        for _ in range(2):
+            for layer, query, key, value in zip(layers, queries, keys, values, strict=True):
+                layer.forward(query, key, value, graph_metadata)
+
+        graph_outputs = []
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            for layer, query, key, value in zip(layers, queries, keys, values, strict=True):
+                graph_outputs.append(layer.forward(query, key, value, graph_metadata))
+
+        graph_metadata.request_ids = replay_request_ids
+        graph_metadata.kv_cache_params = KVCacheParams(
+            use_cache=True,
+            num_cached_tokens_per_seq=replay_cached_tokens,
+        )
+        graph_metadata.prepare()
+
+        reference_metadata = make_metadata(
+            is_cuda_graph=False,
+            request_ids=replay_request_ids,
+            cached_tokens=replay_cached_tokens,
+        )
+        reference_metadata.prepare()
+        reference_outputs = [
+            layer.forward(query, key, value, reference_metadata)
+            for layer, query, key, value in zip(layers, queries, keys, values, strict=True)
+        ]
+
+        graph.replay()
+        torch.cuda.synchronize()
+
+        for layer, graph_output, reference_output in zip(
+            layers, graph_outputs, reference_outputs, strict=True
+        ):
+            torch.testing.assert_close(
+                graph_output,
+                reference_output,
+                atol=1e-2,
+                rtol=0,
+                msg=f"Layer {layer.layer_idx}: verification graph output diverges from eager",
+            )
+
     @torch.no_grad()
     @unittest.mock.patch(
         "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py b/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
new file mode 100644
index 000000000000..85ac7601fa56
--- /dev/null
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+import torch
+
+from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
+from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine
+from tensorrt_llm._torch.speculative.drafting_loops import Gemma4AssistantDraftingLoopWrapper
+from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata
+from tensorrt_llm._torch.speculative.model_drafter import ModelDrafter
+
+
+class _DummyGemma4Assistant(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.config = SimpleNamespace(model_type="gemma4_assistant")
+        self.model_config = None
+        self.model = SimpleNamespace()
+        self.calls = []
+
+    def forward(self, input_ids, position_ids, attn_metadata, **kwargs):
+        self.calls.append(
+            {
+                "input_ids": input_ids.clone(),
+                "position_ids": position_ids.clone(),
+                "kv_lens": attn_metadata.kv_lens_cuda.clone(),
+            }
+        )
+        logits = torch.zeros((input_ids.shape[0], 8))
+        logits[:, len(self.calls)] = 1
+        return logits
+
+
+def test_gemma4_drafting_loop_keeps_position_and_target_kv_length():
+    draft_model = _DummyGemma4Assistant()
+    wrapper = Gemma4AssistantDraftingLoopWrapper(
+        max_draft_len=3,
+        max_total_draft_tokens=3,
+        draft_model=draft_model,
+    )
+    wrapper.sample = lambda logits: logits.argmax(dim=-1)
+
+    attn_metadata = SimpleNamespace(
+        num_seqs=2,
+        kv_lens_cuda=torch.tensor([7, 11]),
+    )
+    spec_metadata = object.__new__(Eagle3SpecMetadata)
+    spec_metadata.gather_ids = torch.tensor([0, 1])
+    spec_metadata.hidden_states_read_indices = torch.tensor([4, 8])
+    spec_metadata.hidden_states_write_indices = torch.tensor([5, 9])
+
+    outputs = wrapper(
+        input_ids=torch.tensor([2, 3]),
+        position_ids=torch.tensor([[6, 10]]),
+        attn_metadata=attn_metadata,
+        spec_metadata=spec_metadata,
+    )
+
+    assert outputs["new_draft_tokens"].tolist() == [[1, 1], [2, 2], [3, 3]]
+    assert len(draft_model.calls) == 3
+    assert all(call["position_ids"].tolist() == [[6, 10]] for call in draft_model.calls)
+    assert all(call["kv_lens"].tolist() == [7, 11] for call in draft_model.calls)
+    assert spec_metadata.hidden_states_read_indices.tolist() == [5, 9]
+
+
+def test_gemma4_drafter_records_target_hidden_state_offset():
+    drafter = object.__new__(ModelDrafter)
+    drafter.spec_resource_manager = SimpleNamespace(
+        draft_hidden_state_offsets={},
+        seq_lens={4: 10},
+        slot_manager=SimpleNamespace(get_slot=lambda request_id: 4),
+    )
+    draft_request = SimpleNamespace()
+    drafter._create_generation_request = lambda request, tokens: draft_request
+    request = SimpleNamespace(
+        py_request_id=17,
+        py_last_context_chunk=(4, 10),
+        py_prompt_len=10,
+        py_num_accepted_draft_tokens=3,
+    )
+
+    assert (
+        drafter._create_gemma4_assistant_request(request, [1, 2], is_first_draft=True)
+        is draft_request
+    )
+    assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 9
+
+    drafter._create_gemma4_assistant_request(request, [1, 2], is_first_draft=False)
+    assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 3
+
+
+def test_gemma4_cuda_graph_warmup_uses_one_token_generation_request():
+    engine = object.__new__(PyTorchModelEngine)
+    engine.is_draft_model = True
+    engine.model_is_wrapped = True
+    engine.model = SimpleNamespace(config=SimpleNamespace(model_type="gemma4_assistant"))
+    spec_resource_manager = object.__new__(Eagle3ResourceManager)
+    spec_resource_manager.is_first_draft = True
+    resource_manager = SimpleNamespace(
+        get_resource_manager=lambda resource_type: spec_resource_manager
+    )
+    request = SimpleNamespace(py_is_first_draft=True, py_draft_tokens=[1])
+    batch = SimpleNamespace(generation_requests=[request])
+
+    engine._update_draft_inference_state_for_warmup(
+        batch, is_first_draft=True, resource_manager=resource_manager
+    )
+
+    assert not spec_resource_manager.is_first_draft
+    assert not request.py_is_first_draft
+    assert request.py_draft_tokens == []
+
+
+def test_gemma4_target_forward_captures_speculative_hidden_states():
+    model = SimpleNamespace(
+        layer_idx=-1,
+        config=SimpleNamespace(final_logit_softcapping=None),
+        model=lambda **kwargs: torch.tensor([[1.0, 2.0]]),
+        logits_processor=SimpleNamespace(forward=lambda hidden_states, *args: hidden_states),
+        lm_head=object(),
+    )
+    captured = []
+    spec_metadata = SimpleNamespace(
+        is_layer_capture=lambda layer_idx: layer_idx == -1,
+        maybe_capture_hidden_states=lambda layer_idx, hidden_states: captured.append(
+            (layer_idx, hidden_states.clone())
+        ),
+    )
+    attn_metadata = SimpleNamespace(padded_num_tokens=None)
+
+    output = Gemma4ForCausalLM.forward(
+        model,
+        attn_metadata=attn_metadata,
+        input_ids=torch.tensor([1]),
+        spec_metadata=spec_metadata,
+    )
+
+    assert torch.equal(output, torch.tensor([[1.0, 2.0]]))
+    assert len(captured) == 1
+    assert captured[0][0] == -1
+    assert torch.equal(captured[0][1], torch.tensor([[1.0, 2.0]]))

From bf69dbb3c6da8bc9a549eeaa258efb73b4f2e394 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 21 Jul 2026 15:13:33 +0000
Subject: [PATCH 7/9] [None][perf] Restore Gemma4 assistant CUDA graph replay

Normalize the Gemma4 assistant CUDA graph key across capture and runtime
lookup so the full drafting loop replays its captured graph instead of
silently falling back to eager execution. Add regression coverage for the
first-draft state transition.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/pyexecutor/cuda_graph_runner.py    | 20 +++++++++---
 .../_torch/pyexecutor/model_engine.py         |  3 ++
 .../executor/test_pytorch_model_engine.py     | 32 ++++++++++++++++++-
 tests/unittest/_torch/helpers.py              |  1 +
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index b8b67ed5ded8..15798776254c 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -103,6 +103,7 @@ class CUDAGraphRunnerConfig:
     original_max_draft_len: int
     original_max_total_draft_tokens: int
     is_draft_model: bool
+    is_gemma4_assistant: bool
     enable_attention_dp: bool
     is_encoder_decoder: bool
     batch_size: int
@@ -247,11 +248,20 @@ def get_graph_key(
 
         if self.config.is_draft_model and spec_resource_manager is not None and isinstance(
                 spec_resource_manager, Eagle3ResourceManager):
-            # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'.
-            # Because we will pad the input to 'max_draft_len' length for the first draft layer.
-            draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0
-            key = (batch_size, draft_len, spec_resource_manager.is_first_draft,
-                   short_seq_len_mode, is_all_greedy_sample)
+            if self.config.is_gemma4_assistant:
+                # Gemma4 captures the whole assistant drafting loop in one
+                # graph, but its external input always contains one query per
+                # request. The resource manager's first-draft flag is internal
+                # loop state and must not select a different graph at runtime.
+                draft_len = 0
+                is_first_draft = False
+            else:
+                # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'.
+                # Because we will pad the input to 'max_draft_len' length for the first draft layer.
+                draft_len = self.config.original_max_draft_len if spec_resource_manager.is_first_draft else 0
+                is_first_draft = spec_resource_manager.is_first_draft
+            key = (batch_size, draft_len, is_first_draft, short_seq_len_mode,
+                   is_all_greedy_sample)
         else:
             # With dynamic spec decode, the draft length may be zero even when enable_spec_decode is True,
             # so we need to get the draft length from the batch instead of using enable_spec_decode.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 822fc6931b4c..9294e5ff670c 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -771,6 +771,9 @@ def __init__(
             original_max_total_draft_tokens=self.
             original_max_total_draft_tokens,
             is_draft_model=self.is_draft_model,
+            is_gemma4_assistant=(self.is_draft_model and self.model_is_wrapped
+                                 and self.model.config.model_type
+                                 == "gemma4_assistant"),
             enable_attention_dp=self.enable_attention_dp,
             is_encoder_decoder=self._is_encoder_decoder_model(),
             batch_size=self.batch_size,
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index 51c36071274c..9b96c9722307 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -3,6 +3,7 @@
 
 import unittest
 from dataclasses import dataclass
+from types import SimpleNamespace
 from unittest.mock import Mock
 
 import torch
@@ -12,7 +13,8 @@
 from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \
     KvCacheConnectorWorker
 from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import (
-    _restore_spec_decode_capture_state, _save_spec_decode_capture_state)
+    CUDAGraphRunner, _restore_spec_decode_capture_state,
+    _save_spec_decode_capture_state)
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.model_engine import (
     PyTorchModelEngine, _build_request_multimodal_input)
@@ -28,6 +30,7 @@
 
 from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
+from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager
 from tensorrt_llm.bindings.executor import KvCacheConfig
 from tensorrt_llm.llmapi import CudaGraphConfig, SamplingParams
 from tensorrt_llm.mapping import CpType, Mapping
@@ -188,6 +191,33 @@ def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None:
 
         self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2)
 
+    def test_gemma4_assistant_graph_key_ignores_first_draft_state(self) -> None:
+        runner = object.__new__(CUDAGraphRunner)
+        runner.config = SimpleNamespace(
+            is_draft_model=True,
+            is_gemma4_assistant=True,
+            original_max_draft_len=2,
+        )
+        runner.sparse_config = None
+        runner.graphs = {}
+        runner.graph_outputs = {}
+        runner.graph_metadata = {}
+        runner.padding_dummy_requests = {}
+        runner.memory_pool = None
+        batch = SimpleNamespace(batch_size=1)
+        resource_manager = object.__new__(Eagle3ResourceManager)
+
+        resource_manager.is_first_draft = False
+        capture_key = runner.get_graph_key(
+            batch, spec_resource_manager=resource_manager)
+
+        resource_manager.is_first_draft = True
+        runtime_key = runner.get_graph_key(
+            batch, spec_resource_manager=resource_manager)
+
+        self.assertEqual(capture_key, (1, 0, False, False, True))
+        self.assertEqual(runtime_key, capture_key)
+
     def test_pad_generation_requests(self) -> None:
         model_engine, kv_cache_manager = create_model_engine_and_kvcache()
         resource_manager = ResourceManager(
diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py
index 7f167077effa..2168709de0db 100644
--- a/tests/unittest/_torch/helpers.py
+++ b/tests/unittest/_torch/helpers.py
@@ -252,6 +252,7 @@ def create_mock_cuda_graph_runner(batch_size: int, use_mrope: bool = False):
         original_max_draft_len=0,
         original_max_total_draft_tokens=0,
         is_draft_model=False,
+        is_gemma4_assistant=False,
         is_encoder_decoder=False,
         mapping=Mapping(),
         dist=None,

From 925108f900ecf3addcea19fee0d6d3a40a64852a Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 21 Jul 2026 16:33:30 +0000
Subject: [PATCH 8/9] [None][fix] Harden Gemma4 MTP runtime integration

Reuse the existing linear drafting loop through model capabilities, enforce assistant configuration invariants, and reject unsupported shared-KV combinations.\n\nFix vocab-parallel ordered logits, shared target KV cache ownership and budgeting, CUDA graph capacity, and FlashInfer KV pool identity. Add focused regressions for configuration, TP logits, graph sizing, cache budgets, and same-head-dimension pools.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    |  24 +---
 tensorrt_llm/_torch/configs/gemma4.py         |  53 ++++++-
 tensorrt_llm/_torch/models/modeling_gemma4.py |  54 ++++++--
 tensorrt_llm/_torch/pyexecutor/_util.py       |  22 +--
 .../_torch/pyexecutor/cuda_graph_runner.py    |  14 +-
 .../_torch/pyexecutor/model_engine.py         |  40 ++++--
 .../_torch/pyexecutor/py_executor_creator.py  |  24 ++--
 .../_torch/speculative/drafting_loops.py      |  76 +++++-----
 tensorrt_llm/_torch/speculative/eagle3.py     |  15 +-
 .../_torch/speculative/model_drafter.py       |  19 +--
 tensorrt_llm/_torch/speculative/utils.py      |   8 +-
 .../executor/test_kv_cache_budget_split.py    |  39 ++++++
 .../executor/test_pytorch_model_engine.py     |  30 +++-
 tests/unittest/_torch/helpers.py              |   1 -
 .../_torch/modeling/test_modeling_gemma4.py   | 131 ++++++++++++++++--
 .../hw_agnostic/test_gemma4_drafting_loop.py  |  12 +-
 16 files changed, 416 insertions(+), 146 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 375703ce5e49..d145e835f1c8 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/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.
@@ -706,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.
@@ -1327,12 +1319,10 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._host_paged_kv_indices = \
                 self._host_pool_indices[primary_pool_id]
 
-        # Gemma4's trtllm-gen paged-prefill graphs capture one stable block
-        # table per head dimension. Refresh those tables and KV lengths after
-        # request turnover instead of replaying with graph-warmup page IDs.
+        # 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):
-            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
@@ -1341,8 +1331,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                         or prefill_wrapper._backend != "trtllm-gen"):
                     continue
                 block_tables = prefill_wrapper._block_tables
-                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 block_tables is None or pool_id is None:
                     continue
                 host_pool_indices = self._host_pool_indices[pool_id]
@@ -1378,7 +1367,6 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                 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
@@ -1386,8 +1374,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                 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
@@ -1486,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)
 
diff --git a/tensorrt_llm/_torch/configs/gemma4.py b/tensorrt_llm/_torch/configs/gemma4.py
index 85ccb7ec2347..e2942239790e 100644
--- a/tensorrt_llm/_torch/configs/gemma4.py
+++ b/tensorrt_llm/_torch/configs/gemma4.py
@@ -27,20 +27,60 @@ class Gemma4AssistantConfig(PreTrainedConfig):
     model_type = "gemma4_assistant"
     sub_configs = {"text_config": Gemma4TextConfig}
 
+    # Runtime capabilities consumed by the generic two-model MTP pipeline.
+    shares_target_kv_cache = True
+    preserve_checkpoint_layer_count = True
+    freezes_draft_attention_state = True
+    cuda_graph_external_draft_len = 0
+
     def __init__(
         self,
         text_config=None,
-        backbone_hidden_size=None,
+        backbone_hidden_size=1536,
         use_ordered_embeddings=False,
-        num_centroids=0,
-        centroid_intermediate_top_k=0,
+        num_centroids=2048,
+        centroid_intermediate_top_k=32,
         **kwargs,
     ):
         if text_config is None:
-            text_config = Gemma4TextConfig()
+            text_config = Gemma4TextConfig(
+                num_hidden_layers=4,
+                num_kv_shared_layers=4,
+                hidden_size_per_layer_input=0,
+                vocab_size_per_layer_input=0,
+                enable_moe_block=False,
+                use_double_wide_mlp=False,
+            )
         elif isinstance(text_config, dict):
             text_config = Gemma4TextConfig(**text_config)
 
+        # Assistant layers are Q-only and all read the target model's KV cache.
+        # Match the native Transformers config behavior when the field is
+        # omitted, and reject partially shared variants that this architecture
+        # cannot execute correctly.
+        if not text_config.num_kv_shared_layers:
+            text_config.num_kv_shared_layers = text_config.num_hidden_layers
+        if text_config.num_kv_shared_layers != text_config.num_hidden_layers:
+            raise ValueError(
+                "All Gemma4 assistant layers must share the target KV cache: "
+                f"expected {text_config.num_hidden_layers}, got "
+                f"{text_config.num_kv_shared_layers}"
+            )
+        if text_config.hidden_size_per_layer_input != 0:
+            raise ValueError(
+                "Gemma4 assistant hidden_size_per_layer_input must be 0, "
+                f"got {text_config.hidden_size_per_layer_input}"
+            )
+        if text_config.vocab_size_per_layer_input != 0:
+            raise ValueError(
+                "Gemma4 assistant vocab_size_per_layer_input must be 0, "
+                f"got {text_config.vocab_size_per_layer_input}"
+            )
+        if text_config.enable_moe_block:
+            raise ValueError("Gemma4 assistant does not support MoE blocks")
+        if text_config.use_double_wide_mlp:
+            raise ValueError("Gemma4 assistant does not support double-wide MLPs")
+
         self.text_config = text_config
         self.backbone_hidden_size = backbone_hidden_size
         self.use_ordered_embeddings = use_ordered_embeddings
@@ -59,3 +99,8 @@ def vocab_size(self):
     @property
     def num_hidden_layers(self):
         return self.text_config.num_hidden_layers
+
+    @property
+    def speculative_hidden_size(self):
+        """Hidden-state width captured from the target model."""
+        return self.backbone_hidden_size
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 3a19d468c302..ea8705030b68 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -42,6 +42,7 @@
     PredefinedAttentionMask,
     RopeParams,
 )
+from ..distributed import AllReduce
 from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
 from ..model_config import ModelConfig
 from ..modules.decoder_layer import DecoderLayer
@@ -1450,12 +1451,37 @@ def __init__(self, model_config: ModelConfig):
             bias=False,
             dtype=config.torch_dtype,
         )
+        self.vocab_all_reduce = AllReduce(
+            mapping=model_config.mapping,
+            dtype=config.torch_dtype,
+        )
         self.register_buffer(
             "token_ordering",
             torch.empty(self.vocab_size, dtype=torch.long, device="cuda"),
         )
 
-    def forward(self, hidden_states: torch.Tensor, lm_head_weight: torch.Tensor) -> torch.Tensor:
+    @staticmethod
+    def _selected_logits_for_vocab_shard(
+        hidden_states: torch.Tensor,
+        lm_head_weight: torch.Tensor,
+        canonical_positions: torch.Tensor,
+        vocab_start_index: int,
+    ) -> torch.Tensor:
+        """Compute selected logits owned by one vocab-parallel shard."""
+        local_positions = canonical_positions - vocab_start_index
+        is_local = (local_positions >= 0) & (local_positions < lm_head_weight.shape[0])
+        safe_positions = local_positions.clamp(0, lm_head_weight.shape[0] - 1)
+        selected_embeddings = lm_head_weight[safe_positions.reshape(-1)].view(
+            hidden_states.shape[0],
+            canonical_positions.shape[1],
+            hidden_states.shape[1],
+        )
+        selected_logits = torch.bmm(
+            hidden_states.unsqueeze(1), selected_embeddings.transpose(1, 2)
+        ).squeeze(1)
+        return selected_logits.masked_fill(~is_local, 0)
+
+    def forward(self, hidden_states: torch.Tensor, lm_head: nn.Module) -> torch.Tensor:
         centroid_logits = self.centroids(hidden_states)
         _, top_k_indices = torch.topk(
             centroid_logits,
@@ -1464,22 +1490,28 @@ def forward(self, hidden_states: torch.Tensor, lm_head_weight: torch.Tensor) ->
         )
         canonical_positions = self.token_ordering.view(
             self.num_centroids, self.vocab_size_per_centroid
-        )[top_k_indices]
-        selected_embeddings = lm_head_weight[canonical_positions.reshape(-1)].view(
-            hidden_states.shape[0],
-            self.centroid_intermediate_top_k * self.vocab_size_per_centroid,
-            self.hidden_size,
+        )[top_k_indices].flatten(1)
+
+        vocab_start_index = 0
+        is_vocab_sharded = lm_head.tp_mode == TensorParallelMode.COLUMN and lm_head.tp_size > 1
+        if is_vocab_sharded:
+            vocab_start_index = lm_head.tp_rank * lm_head.out_features
+        selected_logits = self._selected_logits_for_vocab_shard(
+            hidden_states,
+            lm_head.weight,
+            canonical_positions,
+            vocab_start_index,
         )
-        selected_logits = torch.bmm(
-            hidden_states.unsqueeze(1), selected_embeddings.transpose(1, 2)
-        ).squeeze(1)
+        if is_vocab_sharded:
+            selected_logits = self.vocab_all_reduce(selected_logits)
+
         logits = torch.full(
             (hidden_states.shape[0], self.vocab_size),
             torch.finfo(hidden_states.dtype).min,
             dtype=hidden_states.dtype,
             device=hidden_states.device,
         )
-        return logits.scatter_(1, canonical_positions.flatten(1), selected_logits)
+        return logits.scatter_(1, canonical_positions, selected_logits)
 
 
 @register_auto_model("Gemma4AssistantForCausalLM")
@@ -1611,7 +1643,7 @@ def forward(
         else:
             logits_hidden_states = self._last_token_states(assistant_hidden_states, attn_metadata)
         if self.masked_embedding is not None:
-            return self.masked_embedding(logits_hidden_states, self.lm_head.weight).float()
+            return self.masked_embedding(logits_hidden_states, self.lm_head).float()
         return self.lm_head(logits_hidden_states).float()
 
     def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 541fc752b628..743673dc6076 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1580,11 +1580,16 @@ def build_managers(self,
         # Split combined KV cache budgets before creating managers. Skip during
         # estimation — estimation uses max_tokens-based logic and must not
         # mutate the config.
-        has_draft = (
-            self._draft_model_engine is not None  # two-model
+        draft_shares_target_kv_cache = (
+            self._draft_model_engine is not None
+            and self._draft_model_engine.kv_cache_manager_key
+            == ResourceManagerType.KV_CACHE_MANAGER)
+        has_independent_draft_cache = (
+            (self._draft_model_engine is not None
+             and not draft_shares_target_kv_cache)  # two-model
             or self._should_create_separate_draft_kv_cache())  # one-model
         draft_kv_cache_config = None
-        if not estimating_kv_cache and has_draft:
+        if not estimating_kv_cache and has_independent_draft_cache:
             # Used when each manager sizes pools from max_gpu_total_bytes (V2
             # and V1 VSWA). V1 non-VSWA GPU uses shared max_tokens instead.
             if self._needs_gpu_kv_cache_budget_split(self_kv_cache_config):
@@ -1617,13 +1622,10 @@ def build_managers(self,
                                        if draft_kv_cache_config is not None else
                                        self_kv_cache_config)
 
-        # Gemma4 assistants read the target model's KV cache directly. They
-        # still need a small draft manager for request/slot bookkeeping, but
-        # must not reserve a second full-size GPU cache during final capacity
-        # allocation.
-        if (self._draft_model_engine is not None
-                and self._draft_model_engine.kv_cache_manager_key
-                == ResourceManagerType.KV_CACHE_MANAGER):
+        # Shared-target-KV drafters still need a small draft manager for
+        # request/slot bookkeeping, but they neither reserve a second full-size
+        # cache nor reduce the target manager's GPU/host cache budget.
+        if draft_shares_target_kv_cache:
             draft_build_kv_cache_config = copy.deepcopy(
                 draft_build_kv_cache_config)
             draft_build_kv_cache_config.max_gpu_total_bytes = 0
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 15798776254c..7b4bebdd7f9b 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -103,7 +103,6 @@ class CUDAGraphRunnerConfig:
     original_max_draft_len: int
     original_max_total_draft_tokens: int
     is_draft_model: bool
-    is_gemma4_assistant: bool
     enable_attention_dp: bool
     is_encoder_decoder: bool
     batch_size: int
@@ -112,6 +111,7 @@ class CUDAGraphRunnerConfig:
     kv_cache_manager_key: Any
     dynamic_draft_len_mapping: Optional[Dict[int, int]] = None
     sparse_attention_config: Optional[BaseSparseAttentionConfig] = None
+    draft_model_external_draft_len: Optional[int] = None
 
 
 class CUDAGraphRunner:
@@ -159,6 +159,8 @@ def __init__(self, config: CUDAGraphRunnerConfig):
     def _create_shared_static_tensors(self):
         """Allocates static tensors sized for the largest possible batch."""
         runtime_draft_token_buffer_width = (
+            self.config.draft_model_external_draft_len
+            if self.config.draft_model_external_draft_len is not None else
             self.config.original_max_total_draft_tokens
             if self.config.spec_config is not None else 0)
         token_per_request = runtime_draft_token_buffer_width + 1
@@ -248,12 +250,10 @@ def get_graph_key(
 
         if self.config.is_draft_model and spec_resource_manager is not None and isinstance(
                 spec_resource_manager, Eagle3ResourceManager):
-            if self.config.is_gemma4_assistant:
-                # Gemma4 captures the whole assistant drafting loop in one
-                # graph, but its external input always contains one query per
-                # request. The resource manager's first-draft flag is internal
-                # loop state and must not select a different graph at runtime.
-                draft_len = 0
+            if self.config.draft_model_external_draft_len is not None:
+                # Capturable draft models may execute the whole drafting loop
+                # internally while exposing a smaller fixed input shape.
+                draft_len = self.config.draft_model_external_draft_len
                 is_first_draft = False
             else:
                 # If 'is_first_draft' is True, even with tree decoding, the length of draft_len will only be 'max_draft_len', not 'max_total_draft_token'.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 9294e5ff670c..e17ea52f2c90 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -66,7 +66,8 @@
                            prepare_attn_metadata_for_draft_replay,
                            restore_attn_metadata_after_draft_replay,
                            update_spec_config_from_loaded_model)
-from ..speculative.drafting_loops import BaseDraftingLoopWrapper
+from ..speculative.drafting_loops import (BaseDraftingLoopWrapper,
+                                          get_draft_model_capability)
 from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata
 from ..speculative.spec_sampler_base import SampleStateTensorsSpec
 from ..utils import (get_model_extra_attrs,
@@ -439,6 +440,11 @@ def __init__(
                 self.model_is_wrapped = False
         else:
             self.model_is_wrapped = False
+        self._shares_target_kv_cache = bool(
+            is_draft_model and get_draft_model_capability(
+                self.model, "shares_target_kv_cache", False))
+        self._cuda_graph_external_draft_len = get_draft_model_capability(
+            self.model, "cuda_graph_external_draft_len", None)
         self.sparse_attention_config = self.model.model_config.sparse_attention_config
         # In case that some tests use stub models and override `_load_model`.
         if not hasattr(self.model, 'extra_attrs'):
@@ -658,9 +664,13 @@ def __init__(
 
         self._cuda_graph_padding_enabled = cuda_graph_padding_enabled
 
+        cuda_graph_max_total_draft_tokens = (
+            self._cuda_graph_external_draft_len
+            if self._cuda_graph_external_draft_len is not None else
+            self.original_max_total_draft_tokens)
         self._cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes(
             cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens,
-            self.original_max_total_draft_tokens,
+            cuda_graph_max_total_draft_tokens,
             self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else []
 
         self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if
@@ -750,7 +760,10 @@ def __init__(
         # We look up this key in resource_manager during forward to find the
         # kv cache manager. Can be changed to support multiple model engines
         # with different KV cache managers.
-        self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER
+        self.kv_cache_manager_key = (ResourceManagerType.DRAFT_KV_CACHE_MANAGER
+                                     if is_draft_model
+                                     and not self._shares_target_kv_cache else
+                                     ResourceManagerType.KV_CACHE_MANAGER)
         self.lora_model_config: Optional[LoraModelConfig] = None
         self._trtllm_gen_jit_warmup = False
 
@@ -771,15 +784,13 @@ def __init__(
             original_max_total_draft_tokens=self.
             original_max_total_draft_tokens,
             is_draft_model=self.is_draft_model,
-            is_gemma4_assistant=(self.is_draft_model and self.model_is_wrapped
-                                 and self.model.config.model_type
-                                 == "gemma4_assistant"),
             enable_attention_dp=self.enable_attention_dp,
             is_encoder_decoder=self._is_encoder_decoder_model(),
             batch_size=self.batch_size,
             mapping=self.mapping,
             dist=self.dist,
             kv_cache_manager_key=self.kv_cache_manager_key,
+            draft_model_external_draft_len=self._cuda_graph_external_draft_len,
             sparse_attention_config=self.sparse_attention_config,
         )
         self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config)
@@ -1476,8 +1487,8 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager):
         logger.info("Running autotuner warmup...")
         kv_cache_manager = resource_manager.get_resource_manager(
             self.kv_cache_manager_key)
-        if (self.is_draft_model and self.model_is_wrapped
-                and self.model.config.model_type == "gemma4_assistant"):
+        if (self.is_draft_model
+                and self._cuda_graph_external_draft_len is not None):
             token_num_upper_bound = 1
         else:
             token_num_upper_bound = min(
@@ -2416,14 +2427,15 @@ def _update_draft_inference_state_for_warmup(
             ResourceManagerType.SPEC_RESOURCE_MANAGER)
         if self.is_draft_model and isinstance(spec_resource_manager,
                                               Eagle3ResourceManager):
-            is_gemma4_assistant = (self.model_is_wrapped
-                                   and self.model.config.model_type
-                                   == "gemma4_assistant")
-            spec_resource_manager.is_first_draft = (is_first_draft
-                                                    and not is_gemma4_assistant)
+            freezes_draft_attention_state = bool(
+                get_draft_model_capability(self.model,
+                                           "freezes_draft_attention_state",
+                                           False))
+            spec_resource_manager.is_first_draft = (
+                is_first_draft and not freezes_draft_attention_state)
             if is_first_draft:
                 for req in batch.generation_requests:
-                    req.py_is_first_draft = not is_gemma4_assistant
+                    req.py_is_first_draft = not freezes_draft_attention_state
                     req.py_draft_tokens = []
 
     def _set_up_attn_metadata(
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 7602aafb62dc..5023afd51103 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -587,9 +587,8 @@ def allocation_scope(current_stage: ExecutorMemoryType):
                                  and llm_args.attn_backend == "TRTLLM")
 
             logger.debug(f"USE CHAIN DRAFTER: {use_chain_drafter}")
-            if (capturable_drafter_eligible
-                    and (use_chain_drafter
-                         or draft_spec_config.spec_dec_mode.is_mtp_eagle())):
+            if (use_chain_drafter
+                    or draft_spec_config.spec_dec_mode.is_mtp_eagle()):
 
                 def drafting_loop_wrapper(model):
                     from tensorrt_llm._torch.speculative.drafting_loops import (
@@ -598,7 +597,18 @@ def drafting_loop_wrapper(model):
                         StaticTreeDraftingLoopWrapper)
                     from tensorrt_llm.llmapi import EagleDecodingConfig
 
-                    if model.config.model_type == "gemma4_assistant":
+                    shares_target_kv_cache = bool(
+                        getattr(model.config, "shares_target_kv_cache", False))
+                    if shares_target_kv_cache:
+                        if guided_decoding_config is not None:
+                            raise ValueError(
+                                "Guided decoding is not supported with draft "
+                                "models that share the target KV cache")
+                        if not capturable_drafter_eligible:
+                            raise ValueError(
+                                "Draft models that share the target KV cache "
+                                "require the capturable greedy drafting loop "
+                                "without a draft length schedule")
                         return Gemma4AssistantDraftingLoopWrapper(
                             spec_config.max_draft_len,
                             spec_config.tokens_per_gen_step - 1, model)
@@ -648,13 +658,11 @@ def drafting_loop_wrapper(model):
             # Embedded MTP checkpoints expose a single draft layer. Standalone
             # Gemma4 assistants keep their full four-layer text backbone.
             if (spec_config.spec_dec_mode.is_mtp_eagle()
-                    and draft_model_engine.model.config.model_type
-                    != "gemma4_assistant"):
+                    and not getattr(draft_model_engine.model.config,
+                                    "preserve_checkpoint_layer_count", False)):
                 draft_model_engine.model.model_config.pretrained_config.num_hidden_layers = 1
             draft_model_engine.load_weights_from_target_model(
                 model_engine.model)
-            if draft_model_engine.model.config.model_type == "gemma4_assistant":
-                draft_model_engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
     else:
         draft_model_engine = None
 
diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py
index 23673059a2a6..8cbd6736a589 100644
--- a/tensorrt_llm/_torch/speculative/drafting_loops.py
+++ b/tensorrt_llm/_torch/speculative/drafting_loops.py
@@ -11,7 +11,7 @@
 """
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from typing import Optional, final
 
 import torch
@@ -59,6 +59,12 @@ def load_weights_from_target_model(self, target_model) -> None:
             self.draft_model.load_weights_from_target_model(target_model)
 
 
+def get_draft_model_capability(model: torch.nn.Module, name: str, default=None):
+    """Read a capability from a wrapped or unwrapped draft model."""
+    draft_model = getattr(model, "draft_model", model)
+    return getattr(draft_model.config, name, default)
+
+
 @contextmanager
 def save_metadata_state(attn_metadata: AttentionMetadata,
                         spec_metadata: SpecMetadata) -> None:
@@ -126,11 +132,13 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
         draft_logits = [logits]
         if self.max_draft_len > 1:
             is_eagle3 = isinstance(spec_metadata, Eagle3SpecMetadata)
-            with save_metadata_state(attn_metadata, spec_metadata):
+            with self.drafting_metadata_context(attn_metadata, spec_metadata):
                 batch_size = attn_metadata.num_seqs
 
                 new_position_ids = self.prepare_for_generation(
                     attn_metadata, spec_metadata, position_ids)
+                self.prepare_hidden_states_for_generation(
+                    spec_metadata, batch_size)
                 for i in range(self.max_draft_len - 1):
                     logits = self.draft_model.forward(
                         input_ids=new_draft_tokens[-1],
@@ -139,8 +147,8 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
                         spec_metadata=spec_metadata)
                     new_draft_tokens.append(self.sample(logits))
                     draft_logits.append(logits)
-                    new_position_ids += 1
-                    attn_metadata.kv_lens_cuda[:batch_size] += 1
+                    self.advance_generation_state(new_position_ids,
+                                                  attn_metadata, batch_size)
                     if i == 0 and is_eagle3:
                         spec_metadata.hidden_states_read_indices[:batch_size].copy_(
                             spec_metadata.
@@ -151,6 +159,20 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
             "draft_logits": torch.stack(draft_logits)
         }
 
+    def drafting_metadata_context(self, attn_metadata: AttentionMetadata,
+                                  spec_metadata: SpecMetadata):
+        return save_metadata_state(attn_metadata, spec_metadata)
+
+    def prepare_hidden_states_for_generation(self, spec_metadata: SpecMetadata,
+                                             batch_size: int) -> None:
+        pass
+
+    def advance_generation_state(self, position_ids: torch.Tensor,
+                                 attn_metadata: AttentionMetadata,
+                                 batch_size: int) -> None:
+        position_ids += 1
+        attn_metadata.kv_lens_cuda[:batch_size] += 1
+
     def sample(self, logits: torch.Tensor) -> torch.Tensor:
         # TODO: inject the sampler here so we can support non-greedy
         tokens, _ = greedy_search_sampling_batch(logits, return_probs=False)
@@ -207,37 +229,21 @@ def prepare_for_generation(self, attn_metadata: AttentionMetadata,
 class Gemma4AssistantDraftingLoopWrapper(LinearDraftingLoopWrapper):
     """Draft tokens without advancing the target KV cache or position."""
 
-    def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor,
-                attn_metadata: AttentionMetadata, spec_metadata: SpecMetadata,
-                **kwargs) -> dict[str, torch.Tensor]:
-        logits = self.draft_model.forward(input_ids=input_ids,
-                                          position_ids=position_ids,
-                                          attn_metadata=attn_metadata,
-                                          spec_metadata=spec_metadata,
-                                          return_context_logits=True)
-        logits = logits[spec_metadata.gather_ids]
-
-        new_draft_tokens = [self.sample(logits)]
-        draft_logits = [logits]
-        if self.max_draft_len > 1:
-            if not isinstance(spec_metadata, Eagle3SpecMetadata):
-                raise TypeError("Gemma4 assistant requires Eagle3 metadata")
-            batch_size = attn_metadata.num_seqs
-            spec_metadata.hidden_states_read_indices[:batch_size].copy_(
-                spec_metadata.hidden_states_write_indices[:batch_size])
-            for _ in range(self.max_draft_len - 1):
-                logits = self.draft_model.forward(
-                    input_ids=new_draft_tokens[-1],
-                    position_ids=position_ids,
-                    attn_metadata=attn_metadata,
-                    spec_metadata=spec_metadata)
-                new_draft_tokens.append(self.sample(logits))
-                draft_logits.append(logits)
-
-        return {
-            "new_draft_tokens": torch.stack(new_draft_tokens),
-            "draft_logits": torch.stack(draft_logits),
-        }
+    def drafting_metadata_context(self, attn_metadata: AttentionMetadata,
+                                  spec_metadata: SpecMetadata):
+        return nullcontext()
+
+    def prepare_hidden_states_for_generation(self, spec_metadata: SpecMetadata,
+                                             batch_size: int) -> None:
+        if not isinstance(spec_metadata, Eagle3SpecMetadata):
+            raise TypeError("Gemma4 assistant requires Eagle3 metadata")
+        spec_metadata.hidden_states_read_indices[:batch_size].copy_(
+            spec_metadata.hidden_states_write_indices[:batch_size])
+
+    def advance_generation_state(self, position_ids: torch.Tensor,
+                                 attn_metadata: AttentionMetadata,
+                                 batch_size: int) -> None:
+        pass
 
     def prepare_for_generation(self, attn_metadata: AttentionMetadata,
                                spec_metadata: SpecMetadata,
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 06e337177d37..5cdf92b76df0 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -215,7 +215,7 @@ class Eagle3SpecMetadata(SpecMetadata):
     is_first_draft: bool = False
     eagle3_resource_manager: Optional[Eagle3ResourceManager] = None
     is_mtp_eagle: bool = False
-    is_gemma4_assistant: bool = False
+    shares_target_kv_cache: bool = False
 
     eagle_choices: Optional[List[List[int]]] = None
     max_total_draft_tokens: int = 0
@@ -283,19 +283,18 @@ def prepare(self):
         for req_id, seq_len in zip(self.request_ids, self.seq_lens):
             slot_id = self.eagle3_resource_manager.slot_manager.get_slot(req_id)
             start_idx = self.eagle3_resource_manager.start_indices[slot_id]
-            # Gemma4 assistants issue one query per target iteration and reuse
-            # the target KV cache. Read the hidden state for the last validated
-            # target token, then overwrite that location with the projected
-            # assistant state for the remaining draft iterations.
-            if self.is_draft_model and self.is_gemma4_assistant:
+            # Shared-target-KV drafters issue one query per target iteration.
+            # Read the hidden state for the last validated target token, then
+            # overwrite that location for the remaining draft iterations.
+            if self.is_draft_model and self.shares_target_kv_cache:
                 assert seq_len == 1, (
-                    "Gemma4 assistant drafting expects one query token per "
+                    "Shared-target-KV drafting expects one query token per "
                     f"request, got {seq_len}")
                 old_seq_len = self.eagle3_resource_manager.seq_lens[slot_id]
                 hidden_state_offset = self.eagle3_resource_manager.draft_hidden_state_offsets.get(
                     req_id, max(old_seq_len - 1, 0))
                 assert old_seq_len == 0 or 0 <= hidden_state_offset < old_seq_len, (
-                    "Gemma4 assistant hidden-state offset is outside the "
+                    "Shared-target-KV hidden-state offset is outside the "
                     f"target span: offset={hidden_state_offset}, "
                     f"target_seq_len={old_seq_len}")
                 hidden_state_idx = start_idx + hidden_state_offset
diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py
index 562e0a10e2da..bffdbafd5833 100644
--- a/tensorrt_llm/_torch/speculative/model_drafter.py
+++ b/tensorrt_llm/_torch/speculative/model_drafter.py
@@ -21,6 +21,7 @@
 from ..pyexecutor.scheduler import ScheduledRequests
 from ..pyexecutor.seq_slot_manager import SeqSlotManager
 from .drafter import Drafter
+from .drafting_loops import get_draft_model_capability
 from .spec_sampler_base import SampleStateTensorsSpec
 
 if TYPE_CHECKING:
@@ -96,10 +97,9 @@ def __init__(
         self.guided_decoder = guided_decoder
 
         self.use_static_draft_loop = draft_model_engine.model_is_wrapped
-        draft_model = (draft_model_engine.model.draft_model if
-                       self.use_static_draft_loop else draft_model_engine.model)
-        self.is_gemma4_assistant = getattr(draft_model.config, "model_type",
-                                           None) == "gemma4_assistant"
+        self.shares_target_kv_cache = bool(
+            get_draft_model_capability(draft_model_engine.model,
+                                       "shares_target_kv_cache", False))
         if self.use_static_draft_loop:
             # TODO: enable sampling/guided decoding on static draft loop
             assert guided_decoder is None
@@ -170,7 +170,7 @@ def _create_generation_request(self, request: LlmRequest,
         new_request.state = LlmRequestState.GENERATION_IN_PROGRESS
         return new_request
 
-    def _create_gemma4_assistant_request(self, request: LlmRequest,
+    def _create_shared_target_kv_request(self, request: LlmRequest,
                                          input_tokens: List[int],
                                          is_first_draft: bool) -> LlmRequest:
         """Create a one-token query over the target model's existing KV cache."""
@@ -178,7 +178,8 @@ def _create_gemma4_assistant_request(self, request: LlmRequest,
         if self.spec_resource_manager is None or not hasattr(
                 self.spec_resource_manager, "draft_hidden_state_offsets"):
             raise RuntimeError(
-                "Gemma4 assistant requires an Eagle3 resource manager")
+                "A shared-target-KV drafter requires an Eagle3 resource manager"
+            )
         if is_first_draft:
             slot_id = self.spec_resource_manager.slot_manager.get_slot(
                 request.py_request_id)
@@ -255,8 +256,8 @@ def _create_draft_request_for_request(
         num_overlap_tokens = 0 if self.disable_overlap_scheduler else 1
         is_first_draft = (request.max_beam_num_tokens - 1 +
                           num_overlap_tokens == request.py_prompt_len)
-        if self.is_gemma4_assistant:
-            return self._create_gemma4_assistant_request(
+        if self.shares_target_kv_cache:
+            return self._create_shared_target_kv_request(
                 request, list(request.get_tokens(0)), is_first_draft)
 
         input_tokens = get_draft_model_prompt(self.spec_config.spec_dec_mode,
@@ -335,7 +336,7 @@ def _prepare_draft_batch(
             for request in scheduled_requests.context_requests:
                 if request.py_disable_speculative_decoding:
                     continue
-                if self.is_gemma4_assistant:
+                if self.shares_target_kv_cache:
                     # The assistant has no private KV cache to populate during
                     # chunked prefill. Drafting starts after target prefill.
                     continue
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index b0f0382f7ff1..e9da38e4f38f 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -130,9 +130,8 @@ def get_spec_metadata(spec_config,
             draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_mtp_eagle():
-        hidden_size = model_config.hidden_size
-        if model_config.model_type == "gemma4_assistant":
-            hidden_size = model_config.backbone_hidden_size
+        hidden_size = getattr(model_config, "speculative_hidden_size",
+                              model_config.hidden_size)
         return Eagle3SpecMetadata(
             max_draft_len=spec_config.max_draft_len,
             max_total_draft_tokens=spec_config.tokens_per_gen_step - 1,
@@ -146,7 +145,8 @@ def get_spec_metadata(spec_config,
             eagle3_resource_manager=spec_resource_manager,
             layers_to_capture=None,
             is_mtp_eagle=True,
-            is_gemma4_assistant=model_config.model_type == "gemma4_assistant",
+            shares_target_kv_cache=getattr(model_config,
+                                           "shares_target_kv_cache", False),
         )
     if spec_config.spec_dec_mode.is_eagle3():
         effective_dynamic_tree = _is_effective_dynamic_tree(spec_config)
diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
index 109460a5a57b..1ffdd83c7d0d 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
@@ -14,11 +14,13 @@
 # limitations under the License.
 """Tests for KV cache budget splitting between target and draft managers."""
 
+from types import SimpleNamespace
 from unittest.mock import Mock
 
 import pytest
 
 from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig
 
 GB = 1 << 30
@@ -65,6 +67,43 @@ def _make_creator(
 
 
 class TestSplitGpuBudgetForDraft:
+    def test_shared_target_cache_keeps_full_target_budget(self):
+        total_gpu = 10 * GB
+        total_host = 20 * GB
+        c = _make_creator(
+            max_gpu_total_bytes=total_gpu,
+            host_cache_size=total_host,
+        )
+        c._draft_model_engine = SimpleNamespace(
+            kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER
+        )
+        c._skip_est = False
+        c._is_encoder_decoder = Mock(return_value=False)
+        c._is_kv_cache_manager_v2 = True
+        c._kv_connector_manager = None
+        c._max_num_tokens = 128
+        c._should_create_separate_draft_kv_cache = Mock(return_value=False)
+        c._split_kv_cache_budget_for_draft = Mock()
+        c._create_kv_cache_manager = Mock(side_effect=["target", "draft"])
+
+        resources = {}
+        c.build_managers(resources, estimating_kv_cache=False)
+
+        c._split_kv_cache_budget_for_draft.assert_not_called()
+        target_config = c._create_kv_cache_manager.call_args_list[0].kwargs[
+            "kv_cache_config_override"
+        ]
+        draft_config = c._create_kv_cache_manager.call_args_list[1].kwargs[
+            "kv_cache_config_override"
+        ]
+        assert target_config.max_gpu_total_bytes == total_gpu
+        assert target_config.host_cache_size == total_host
+        assert draft_config.max_gpu_total_bytes == 0
+        assert draft_config.host_cache_size == 0
+        assert draft_config.max_tokens == c._max_seq_len * c._max_batch_size
+        assert resources[ResourceManagerType.KV_CACHE_MANAGER] == "target"
+        assert resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] == "draft"
+
     def test_gpu_budget_split_proportionally(self):
         total_gpu = 10 * GB
         c = _make_creator(
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index 9b96c9722307..9b25578da6d3 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -17,7 +17,8 @@
     _save_spec_decode_capture_state)
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.model_engine import (
-    PyTorchModelEngine, _build_request_multimodal_input)
+    PyTorchModelEngine, _build_request_multimodal_input,
+    _filter_cuda_graph_batch_sizes)
 from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
 
 # isort: off
@@ -191,11 +192,12 @@ def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None:
 
         self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2)
 
-    def test_gemma4_assistant_graph_key_ignores_first_draft_state(self) -> None:
+    def test_external_draft_len_graph_key_ignores_first_draft_state(
+            self) -> None:
         runner = object.__new__(CUDAGraphRunner)
         runner.config = SimpleNamespace(
             is_draft_model=True,
-            is_gemma4_assistant=True,
+            draft_model_external_draft_len=0,
             original_max_draft_len=2,
         )
         runner.sparse_config = None
@@ -218,6 +220,28 @@ def test_gemma4_assistant_graph_key_ignores_first_draft_state(self) -> None:
         self.assertEqual(capture_key, (1, 0, False, False, True))
         self.assertEqual(runtime_key, capture_key)
 
+    def test_external_draft_len_preserves_cuda_graph_batch_capacity(
+            self) -> None:
+        batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
+
+        regular_draft_sizes = _filter_cuda_graph_batch_sizes(
+            batch_sizes,
+            max_batch_size=128,
+            max_num_tokens=128,
+            max_total_draft_tokens=5,
+            enable_padding=False,
+        )
+        external_draft_sizes = _filter_cuda_graph_batch_sizes(
+            batch_sizes,
+            max_batch_size=128,
+            max_num_tokens=128,
+            max_total_draft_tokens=0,
+            enable_padding=False,
+        )
+
+        self.assertEqual(regular_draft_sizes, [1, 2, 4, 8, 16])
+        self.assertEqual(external_draft_sizes, batch_sizes)
+
     def test_pad_generation_requests(self) -> None:
         model_engine, kv_cache_manager = create_model_engine_and_kvcache()
         resource_manager = ResourceManager(
diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py
index 2168709de0db..7f167077effa 100644
--- a/tests/unittest/_torch/helpers.py
+++ b/tests/unittest/_torch/helpers.py
@@ -252,7 +252,6 @@ def create_mock_cuda_graph_runner(batch_size: int, use_mrope: bool = False):
         original_max_draft_len=0,
         original_max_total_draft_tokens=0,
         is_draft_model=False,
-        is_gemma4_assistant=False,
         is_encoder_decoder=False,
         mapping=Mapping(),
         dist=None,
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 8276807322ee..ca2e6cfb90cb 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -19,13 +19,14 @@
 """
 
 import math
+import tempfile
 import unittest
 import unittest.mock
 from copy import deepcopy
 from typing import TYPE_CHECKING
 
 import torch
-from transformers import Gemma4Config, Gemma4TextConfig
+from transformers import AutoConfig, Gemma4Config, Gemma4TextConfig
 
 from tensorrt_llm._torch.attention_backend import FlashInferAttention, FlashInferAttentionMetadata
 from tensorrt_llm._torch.configs.gemma4 import Gemma4AssistantConfig
@@ -34,6 +35,7 @@
 from tensorrt_llm._torch.models.checkpoints.hf.gemma4_weight_mapper import Gemma4HfWeightMapper
 from tensorrt_llm._torch.models.modeling_gemma4 import (
     Gemma4AssistantForCausalLM,
+    Gemma4AssistantMaskedEmbedder,
     Gemma4Attention,
     Gemma4DecoderLayer,
     Gemma4ForCausalLM,
@@ -117,6 +119,7 @@
             "full_attention",
         ],
         "num_kv_shared_layers": 4,
+        "vocab_size_per_layer_input": 0,
     },
     "backbone_hidden_size": 256,
     "use_ordered_embeddings": True,
@@ -486,6 +489,64 @@ def test_assistant_config_wraps_text_config(self):
         self.assertEqual(config.vocab_size, 1024)
         self.assertEqual(config.num_hidden_layers, 4)
 
+    def test_assistant_config_default_constructor_is_serializable(self):
+        config = Gemma4AssistantConfig()
+
+        self.assertEqual(config.text_config.num_kv_shared_layers, 4)
+        self.assertIn("text_config", config.to_dict())
+
+    def test_assistant_config_auto_config_round_trip(self):
+        config = Gemma4AssistantConfig(**deepcopy(GEMMA4_ASSISTANT_CONFIG))
+
+        with tempfile.TemporaryDirectory() as directory:
+            config.save_pretrained(directory)
+            restored = AutoConfig.from_pretrained(directory)
+
+        self.assertIsInstance(restored, Gemma4AssistantConfig)
+        self.assertEqual(restored.backbone_hidden_size, 256)
+        self.assertEqual(restored.text_config.num_kv_shared_layers, 4)
+
+    def test_assistant_config_defaults_to_sharing_all_target_kv_layers(self):
+        config_dict = deepcopy(GEMMA4_ASSISTANT_CONFIG)
+        config_dict["text_config"].pop("num_kv_shared_layers")
+
+        config = Gemma4AssistantConfig(**config_dict)
+
+        self.assertEqual(
+            config.text_config.num_kv_shared_layers,
+            config.text_config.num_hidden_layers,
+        )
+
+    def test_assistant_config_rejects_partially_shared_target_kv(self):
+        config_dict = deepcopy(GEMMA4_ASSISTANT_CONFIG)
+        config_dict["text_config"]["num_kv_shared_layers"] = 2
+
+        with self.assertRaisesRegex(ValueError, "must share the target KV cache"):
+            Gemma4AssistantConfig(**config_dict)
+
+    def test_ordered_embedding_combines_vocab_parallel_shards(self):
+        hidden_states = torch.tensor([[1.0, 2.0, 3.0, 4.0], [4.0, 3.0, 2.0, 1.0]])
+        lm_head_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
+        canonical_positions = torch.tensor([[0, 5, 7], [3, 4, 6]])
+
+        shard_logits = []
+        for vocab_start_index, shard in ((0, lm_head_weight[:4]), (4, lm_head_weight[4:])):
+            shard_logits.append(
+                Gemma4AssistantMaskedEmbedder._selected_logits_for_vocab_shard(
+                    hidden_states,
+                    shard,
+                    canonical_positions,
+                    vocab_start_index,
+                )
+            )
+        actual = sum(shard_logits)
+        selected_weights = lm_head_weight[canonical_positions]
+        expected = torch.bmm(hidden_states.unsqueeze(1), selected_weights.transpose(1, 2)).squeeze(
+            1
+        )
+
+        torch.testing.assert_close(actual, expected)
+
     def test_assistant_uses_target_kv_sources(self):
         assistant = Gemma4AssistantForCausalLM(_make_assistant_model_config())
         self.assertEqual(len(assistant.model.layers), 4)
@@ -739,7 +800,13 @@ def test_assistant_uses_target_kv_sources(self):
 }
 
 
-def _build_gemma4_kv_cache_manager(config, num_blocks=4, tokens_per_block=32, batch_size=1):
+def _build_gemma4_kv_cache_manager(
+    config,
+    num_blocks=4,
+    tokens_per_block=32,
+    batch_size=1,
+    force_vswa=False,
+):
     """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads.
 
     Mirrors ``Gemma4Attention``'s layout (global kv heads only for K=V layers)
@@ -794,7 +861,7 @@ def _build_gemma4_kv_cache_manager(config, num_blocks=4, tokens_per_block=32, ba
     # exceeds sliding_window.
     sliding_window = getattr(config, "sliding_window", None)
     max_attn_window = None
-    needs_vswa = isinstance(head_dim, list) and len(set(head_dim)) > 1
+    needs_vswa = force_vswa or (isinstance(head_dim, list) and len(set(head_dim)) > 1)
     if not needs_vswa:
         needs_vswa = isinstance(num_kv_heads, list) and len(set(num_kv_heads)) > 1
     if needs_vswa and sliding_window:
@@ -2305,9 +2372,11 @@ def _make_trtllm_gen_decode_case(
         self,
         initial_page_counts: list[int],
         *,
+        config_dict: dict | None = None,
         reserved_page_counts: list[int] | None = None,
         max_pages: int = 64,
         manager_batch_size: int | None = None,
+        force_vswa: bool = False,
     ) -> tuple[
         "KVCacheManagerV2",
         list["FlashInferAttention"],
@@ -2323,12 +2392,13 @@ def _make_trtllm_gen_decode_case(
         if manager_batch_size is None:
             manager_batch_size = batch_size
 
-        config = Gemma4TextConfig(**deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG))
+        config = Gemma4TextConfig(**deepcopy(config_dict or GEMMA4_E2B_REAL_DIMS_CONFIG))
         kv_cache_manager = self._get_kv_cache_manager(
             config,
             num_blocks=max_pages,
             tokens_per_block=_TRTLLM_GEN_TOKENS_PER_BLOCK,
             batch_size=manager_batch_size,
+            force_vswa=force_vswa,
         )
         self.addCleanup(kv_cache_manager.shutdown)
         self.assertTrue(kv_cache_manager.is_vswa, "Expected VSWA manager")
@@ -2447,14 +2517,13 @@ def _prepare_decode_page_counts(
     def _expected_decode_block_table(
         self,
         metadata: "FlashInferAttentionMetadata",
-        head_dim: int,
+        pool_id: int,
         page_counts: list[int],
         *,
         rows: int,
         width: int,
     ) -> torch.Tensor:
         """Build the expected compact table from one VSWA pool's host indices."""
-        pool_id = metadata._vswa_head_dim_to_pool[head_dim]
         pool_indices = metadata._host_pool_indices[pool_id].numpy()
         expected = torch.zeros((rows, width), dtype=torch.int32)
         source_offset = metadata.num_context_blocks
@@ -2489,7 +2558,7 @@ def test_cuda_graph_trtllm_gen_block_table_transitions(self) -> None:
                 block_tables = wrappers.decode_wrapper._block_tables
                 expected = self._expected_decode_block_table(
                     metadata,
-                    plan_params.head_dim,
+                    plan_params.kv_pool_id,
                     new_page_counts,
                     rows=len(initial_page_counts),
                     width=max(initial_page_counts),
@@ -2518,11 +2587,12 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N
 
         initial_state = {}
         for plan_params, wrappers in metadata._plan_params_to_wrappers.items():
+            self.assertIsNotNone(plan_params.kv_pool_id)
             block_tables = wrappers.decode_wrapper._block_tables
             self.assertGreaterEqual(block_tables.size(1), 65)
             self.assertEqual(block_tables.size(1), metadata.kv_cache_manager.max_blocks_per_seq)
             self.assertEqual(wrappers.host_decode_block_tables.size(1), 64)
-            initial_state[plan_params.head_dim] = (
+            initial_state[plan_params.kv_pool_id] = (
                 block_tables.data_ptr(),
                 wrappers.host_decode_block_tables.data_ptr(),
             )
@@ -2534,13 +2604,13 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N
         for plan_params, wrappers in metadata._plan_params_to_wrappers.items():
             with self.subTest(head_dim=plan_params.head_dim):
                 block_tables = wrappers.decode_wrapper._block_tables
-                old_device_ptr, old_host_ptr = initial_state[plan_params.head_dim]
+                old_device_ptr, old_host_ptr = initial_state[plan_params.kv_pool_id]
                 self.assertEqual(block_tables.data_ptr(), old_device_ptr)
                 self.assertNotEqual(wrappers.host_decode_block_tables.data_ptr(), old_host_ptr)
                 self.assertGreaterEqual(wrappers.host_decode_block_tables.size(1), 65)
                 expected = self._expected_decode_block_table(
                     metadata,
-                    plan_params.head_dim,
+                    plan_params.kv_pool_id,
                     new_page_counts,
                     rows=len(new_page_counts),
                     width=max(new_page_counts),
@@ -2552,6 +2622,47 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N
                     rtol=0,
                 )
 
+    @torch.no_grad()
+    @unittest.mock.patch(
+        "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
+    )
+    def test_cuda_graph_trtllm_gen_distinguishes_same_head_dim_pools(self) -> None:
+        """Plan keys retain the KV pool when sliding and full head dims match."""
+        config_dict = deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG)
+        config_dict["global_head_dim"] = config_dict["head_dim"]
+        initial_page_counts = [5, 3]
+        _, _, metadata, _, _, _ = self._make_trtllm_gen_decode_case(
+            initial_page_counts,
+            config_dict=config_dict,
+            force_vswa=True,
+        )
+
+        plan_params = list(metadata._plan_params_to_wrappers)
+        self.assertEqual({params.head_dim for params in plan_params}, {256})
+        self.assertEqual(len({params.kv_pool_id for params in plan_params}), 2)
+
+        new_page_counts = [2, 1]
+        self._prepare_decode_page_counts(metadata, [0, 1], new_page_counts)
+        torch.cuda.synchronize()
+
+        for params, wrappers in metadata._plan_params_to_wrappers.items():
+            with self.subTest(pool_id=params.kv_pool_id):
+                expected = self._expected_decode_block_table(
+                    metadata,
+                    params.kv_pool_id,
+                    new_page_counts,
+                    rows=len(new_page_counts),
+                    width=max(new_page_counts),
+                )
+                torch.testing.assert_close(
+                    wrappers.decode_wrapper._block_tables[
+                        : len(new_page_counts), : max(new_page_counts)
+                    ].cpu(),
+                    expected,
+                    atol=0,
+                    rtol=0,
+                )
+
     @torch.no_grad()
     @unittest.mock.patch(
         "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py b/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
index 85ac7601fa56..b6eb66fcb6b1 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
@@ -15,7 +15,11 @@
 class _DummyGemma4Assistant(torch.nn.Module):
     def __init__(self) -> None:
         super().__init__()
-        self.config = SimpleNamespace(model_type="gemma4_assistant")
+        self.config = SimpleNamespace(
+            model_type="gemma4_assistant",
+            shares_target_kv_cache=True,
+            freezes_draft_attention_state=True,
+        )
         self.model_config = None
         self.model = SimpleNamespace()
         self.calls = []
@@ -82,12 +86,12 @@ def test_gemma4_drafter_records_target_hidden_state_offset():
     )
 
     assert (
-        drafter._create_gemma4_assistant_request(request, [1, 2], is_first_draft=True)
+        drafter._create_shared_target_kv_request(request, [1, 2], is_first_draft=True)
         is draft_request
     )
     assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 9
 
-    drafter._create_gemma4_assistant_request(request, [1, 2], is_first_draft=False)
+    drafter._create_shared_target_kv_request(request, [1, 2], is_first_draft=False)
     assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 3
 
 
@@ -95,7 +99,7 @@ def test_gemma4_cuda_graph_warmup_uses_one_token_generation_request():
     engine = object.__new__(PyTorchModelEngine)
     engine.is_draft_model = True
     engine.model_is_wrapped = True
-    engine.model = SimpleNamespace(config=SimpleNamespace(model_type="gemma4_assistant"))
+    engine.model = SimpleNamespace(config=SimpleNamespace(freezes_draft_attention_state=True))
     spec_resource_manager = object.__new__(Eagle3ResourceManager)
     spec_resource_manager.is_first_draft = True
     resource_manager = SimpleNamespace(

From 3b9c14bb1bc17fa8930a79a1849205fa7d4a25b0 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Fri, 24 Jul 2026 05:52:52 +0000
Subject: [PATCH 9/9] [None][fix] Avoid redundant Gemma4 assistant KV cache

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 20 ++++---------------
 .../executor/test_kv_cache_budget_split.py    | 17 +++++-----------
 2 files changed, 9 insertions(+), 28 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 743673dc6076..0078f6c46d49 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1622,22 +1622,10 @@ def build_managers(self,
                                        if draft_kv_cache_config is not None else
                                        self_kv_cache_config)
 
-        # Shared-target-KV drafters still need a small draft manager for
-        # request/slot bookkeeping, but they neither reserve a second full-size
-        # cache nor reduce the target manager's GPU/host cache budget.
-        if draft_shares_target_kv_cache:
-            draft_build_kv_cache_config = copy.deepcopy(
-                draft_build_kv_cache_config)
-            draft_build_kv_cache_config.max_gpu_total_bytes = 0
-            draft_build_kv_cache_config.free_gpu_memory_fraction = None
-            draft_build_kv_cache_config.host_cache_size = 0
-            draft_build_kv_cache_config.max_tokens = max(
-                self._max_num_tokens,
-                self._max_seq_len * self._max_batch_size,
-            )
-
-        # Two-model speculative decoding: draft model has separate engine
-        if self._draft_model_engine is not None:
+        # Two-model speculative decoding with an independent draft KV cache.
+        # Shared-target-KV draft engines use the primary manager instead.
+        if (self._draft_model_engine is not None
+                and not draft_shares_target_kv_cache):
             if self._is_kv_cache_manager_v2:
                 assert draft_kv_cache_config is None, (
                     "KVCacheManagerV2 does not support two-model speculative "
diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
index 1ffdd83c7d0d..610676911cc0 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
@@ -67,7 +67,7 @@ def _make_creator(
 
 
 class TestSplitGpuBudgetForDraft:
-    def test_shared_target_cache_keeps_full_target_budget(self):
+    def test_shared_target_cache_skips_draft_manager(self):
         total_gpu = 10 * GB
         total_host = 20 * GB
         c = _make_creator(
@@ -84,25 +84,18 @@ def test_shared_target_cache_keeps_full_target_budget(self):
         c._max_num_tokens = 128
         c._should_create_separate_draft_kv_cache = Mock(return_value=False)
         c._split_kv_cache_budget_for_draft = Mock()
-        c._create_kv_cache_manager = Mock(side_effect=["target", "draft"])
+        c._create_kv_cache_manager = Mock(return_value="target")
 
         resources = {}
         c.build_managers(resources, estimating_kv_cache=False)
 
         c._split_kv_cache_budget_for_draft.assert_not_called()
-        target_config = c._create_kv_cache_manager.call_args_list[0].kwargs[
-            "kv_cache_config_override"
-        ]
-        draft_config = c._create_kv_cache_manager.call_args_list[1].kwargs[
-            "kv_cache_config_override"
-        ]
+        c._create_kv_cache_manager.assert_called_once()
+        target_config = c._create_kv_cache_manager.call_args.kwargs["kv_cache_config_override"]
         assert target_config.max_gpu_total_bytes == total_gpu
         assert target_config.host_cache_size == total_host
-        assert draft_config.max_gpu_total_bytes == 0
-        assert draft_config.host_cache_size == 0
-        assert draft_config.max_tokens == c._max_seq_len * c._max_batch_size
         assert resources[ResourceManagerType.KV_CACHE_MANAGER] == "target"
-        assert resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] == "draft"
+        assert resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] is None
 
     def test_gpu_budget_split_proportionally(self):
         total_gpu = 10 * GB