From 28c0a430a63528f4a55678777465b89bf61e36f0 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 01/26] [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 +++++++++++++++++-
 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 ++++++
 9 files changed, 376 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 bb06009bdc33..265303c95a09 100644
--- a/examples/llm-api/quickstart_advanced.py
+++ b/examples/llm-api/quickstart_advanced.py
@@ -308,6 +308,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.
@@ -318,7 +321,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 708e50e0ac5f..64eedab07343 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,
@@ -37,6 +53,7 @@ def _register_custom_configs_with_transformers() -> None:
         "deepseek_v32": DeepseekV3Config,
         "kimi_k2": DeepseekV3Config,
         "deepseek_v4": DeepseekV4Config,
+        "gemma4_assistant": Gemma4AssistantConfig,
         "laguna": LagunaConfig,
         # minicpmv4_6 is only registered in transformers>=5.7.0; register our
         # own composite config so AutoTokenizer.from_pretrained works on older
@@ -64,6 +81,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 22cde4599c43..00698a44b63f 100644
--- a/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
+++ b/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py
@@ -32,6 +32,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/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 9c0023cac331..53b4c3ce8b51 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1736,6 +1736,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 6149ba5ec0c3..d2fb79512930 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -640,11 +640,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 bffa8833058c..8b01b5a4245b 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -131,13 +131,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 b66eea54cc95..36f0d6c830b6 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -29,10 +29,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,
@@ -105,6 +107,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."""
@@ -113,6 +135,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."""
 
@@ -531,6 +560,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 80522a4b2e586fa3b5822757ce65cec4af307cca 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 02/26] [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 4f4db5525b60..cab451c990bd 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` [^15]   | Gemma 4 MTP assistant              | `google/gemma-4-E2B-it-assistant`, `google/gemma-4-E4B-it-assistant`, `google/gemma-4-26B-A4B-it-assistant`, `google/gemma-4-31B-it-assistant` |
 | `Glm4MoeForCausalLM`                 | GLM-4.5, GLM-4.6, GLM-4.7          | `THUDM/GLM-4-100B-A10B`                      |
 | `Glm4MoeLiteForCausalLM` [^5]        | GLM-4.7-Flash                      | `zai-org/GLM-4.7-Flash`                      |
 | `GlmMoeDsaForCausalLM`               | GLM-5                              | `zai-org/GLM-5`                              |
@@ -76,7 +77,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `GptOssForCausalLM`              | Yes               | Yes        | Yes                        | Yes                   | Yes             | No  | Yes              | No                | Yes    | Yes           | Yes              | Yes            | N/A                      | Yes                   | Yes             |
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
-| `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
+| `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | Yes [^15] | 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        |
@@ -94,6 +95,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^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]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`).
+[^15]: 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 66faefd015297d00209d727a3388b3a3ae401be4 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 03/26] [None][docs] Remove Gemma4 MTP footnote

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

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index cab451c990bd..b1ff9698cc8d 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` [^15]   | 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`                              |
@@ -77,7 +77,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `GptOssForCausalLM`              | Yes               | Yes        | Yes                        | Yes                   | Yes             | No  | Yes              | No                | Yes    | Yes           | Yes              | Yes            | N/A                      | Yes                   | Yes             |
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
-| `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | Yes [^15] | 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        |
@@ -95,7 +95,6 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^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]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`).
-[^15]: 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)
 

From b075c1c144e40e32ab939ef28be78b793d29b2d6 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 04/26] [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 2276ccda73457563ce511156e20a5d1bc84ee205 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 05/26] [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 cfb0360d1c80..acf595cfa74e 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -311,8 +311,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,
@@ -461,8 +461,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
@@ -1550,8 +1554,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)
@@ -2513,10 +2521,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 d2fb79512930..48b0cf454ab3 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -585,22 +585,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 133814bb2d22..e54e23cb6bee 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 09ba663c744d..bcf506763768 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 8b01b5a4245b..dfc23d13d3b8 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -147,6 +147,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 36f0d6c830b6..dde38f32996a 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -2890,6 +2890,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 ff036a453e4556087536da3e261932edbb244696 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 06/26] [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     | 28 +++++++++++++++++++
 tests/unittest/_torch/helpers.py              |  1 +
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index 05343d0deddd..6cc5c9af44d9 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
@@ -270,11 +271,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 acf595cfa74e..fffbd324019c 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -816,6 +816,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 884fb8e2eebb..1b377b292a35 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -36,6 +36,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._torch.speculative.spec_sampler_base import \
     SampleStateTensorsSpec
 from tensorrt_llm.bindings.executor import KvCacheConfig
@@ -1115,6 +1116,33 @@ def test_promoted_context_precedes_speculative_overlap_generation(
             [generation.py_seq_slot], 0)
         kv_cache_manager.shutdown()
 
+    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 2c91e6445c1a3b950d64ae1c8802f0963cfd68c3 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 07/26] [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     |  29 +++-
 tests/unittest/_torch/helpers.py              |   1 -
 .../_torch/modeling/test_modeling_gemma4.py   | 131 ++++++++++++++++--
 .../hw_agnostic/test_gemma4_drafting_loop.py  |  12 +-
 16 files changed, 415 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 53b4c3ce8b51..dc861b191b25 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1699,11 +1699,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):
@@ -1736,13 +1741,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 6cc5c9af44d9..c64b49646862 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
@@ -271,12 +273,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 fffbd324019c..5c247f3233b8 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -67,7 +67,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,
@@ -469,6 +470,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'):
@@ -703,9 +709,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
@@ -795,7 +805,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
 
@@ -816,15 +829,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)
@@ -1557,8 +1568,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(
@@ -2524,14 +2535,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 48b0cf454ab3..43f21df70741 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -594,9 +594,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 (
@@ -605,7 +604,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)
@@ -655,13 +665,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 e54e23cb6bee..0254235aa3d5 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 bcf506763768..1713958a1f95 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 dfc23d13d3b8..79aafa0859d4 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -131,9 +131,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,
@@ -147,7 +146,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 1b377b292a35..b49ec0b4ba0d 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -21,7 +21,7 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.model_engine import (
     PyTorchModelEngine, _build_request_multimodal_input,
-    _make_single_token_context_graph_batch)
+    _filter_cuda_graph_batch_sizes, _make_single_token_context_graph_batch)
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           SeqLenAwareSparseAttentionConfig,
                                           TorchLlmArgs)
@@ -1116,11 +1116,12 @@ def test_promoted_context_precedes_speculative_overlap_generation(
             [generation.py_seq_slot], 0)
         kv_cache_manager.shutdown()
 
-    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
@@ -1143,6 +1144,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 dde38f32996a..86984ec1e3d6 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -19,6 +19,7 @@
 """
 
 import math
+import tempfile
 import unittest
 import unittest.mock
 from copy import deepcopy
@@ -26,7 +27,7 @@
 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
@@ -35,6 +36,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,
@@ -118,6 +120,7 @@
             "full_attention",
         ],
         "num_kv_shared_layers": 4,
+        "vocab_size_per_layer_input": 0,
     },
     "backbone_hidden_size": 256,
     "use_ordered_embeddings": True,
@@ -572,6 +575,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)
@@ -825,7 +886,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)
@@ -880,7 +947,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:
@@ -2391,9 +2458,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"],
@@ -2409,12 +2478,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")
@@ -2533,14 +2603,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
@@ -2575,7 +2644,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),
@@ -2604,11 +2673,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(),
             )
@@ -2620,13 +2690,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),
@@ -2638,6 +2708,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 a0470c28506d0af62ee4242d9433eae082357f4b 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 08/26] [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 dc861b191b25..b6373ca907b1 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1741,22 +1741,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

From 02fd45e56d89e05b6d1ecf80fce07ef4c8d41b00 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Fri, 24 Jul 2026 09:21:57 +0000
Subject: [PATCH 09/26] [None][test] Clarify Gemma4 MTP graph pool coverage

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/flashinfer.py    | 6 ++++--
 tests/unittest/_torch/modeling/test_modeling_gemma4.py | 1 +
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index d145e835f1c8..4a8df996b199 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -1319,8 +1319,10 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._host_paged_kv_indices = \
                 self._host_pool_indices[primary_pool_id]
 
-        # trtllm-gen paged-prefill graphs capture one stable block table per KV
-        # pool. Refresh those tables and KV lengths after request turnover.
+        # Decoder graph batches have no scheduled context requests, but linear
+        # speculative verification reclassifies draft-token extensions as
+        # contexts. Refresh each trtllm-gen paged-prefill graph's stable block
+        # table and KV lengths after request turnover.
         if (self.is_cuda_graph and self.num_contexts > 0
                 and self._vswa_layer_to_pool is not None):
             for plan_params, wrappers in self._plan_params_to_wrappers.items():
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 86984ec1e3d6..c52258c08cae 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -2724,6 +2724,7 @@ def test_cuda_graph_trtllm_gen_distinguishes_same_head_dim_pools(self) -> None:
         )
 
         plan_params = list(metadata._plan_params_to_wrappers)
+        self.assertEqual(len(plan_params), 2)
         self.assertEqual({params.head_dim for params in plan_params}, {256})
         self.assertEqual(len({params.kv_pool_id for params in plan_params}), 2)
 

From 8ac60c41ec572499a39eecda06a6f6fd42eed595 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Fri, 24 Jul 2026 12:17:51 +0000
Subject: [PATCH 10/26] feat: move Gemma4 MTP to one-model shared KV

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 examples/llm-api/quickstart_advanced.py       |   5 +-
 examples/models/core/gemma/README.md          |  15 +-
 .../_torch/attention_backend/flashinfer.py    | 152 +++++++++++-
 tensorrt_llm/_torch/configs/gemma4.py         |   9 +-
 tensorrt_llm/_torch/models/modeling_gemma4.py |  85 +++++--
 .../_torch/models/modeling_gemma4mm.py        |  32 +++
 .../_torch/models/modeling_speculative.py     |  54 ++++-
 tensorrt_llm/_torch/models/modeling_utils.py  |   4 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |  25 +-
 .../_torch/pyexecutor/cuda_graph_runner.py    |  20 +-
 .../_torch/pyexecutor/model_engine.py         |  83 +++----
 .../_torch/pyexecutor/model_loader.py         |  14 +-
 .../_torch/pyexecutor/py_executor_creator.py  |  55 ++---
 .../_torch/pyexecutor/resource_manager.py     |   4 +-
 tensorrt_llm/_torch/speculative/__init__.py   |  19 +-
 .../_torch/speculative/drafting_loops.py      |  57 +----
 tensorrt_llm/_torch/speculative/eagle3.py     | 200 +++++++++++++--
 tensorrt_llm/_torch/speculative/interface.py  |  77 +++++-
 .../_torch/speculative/model_drafter.py       |  45 +---
 tensorrt_llm/_torch/speculative/utils.py      |  17 +-
 tensorrt_llm/llmapi/llm_args.py               |   3 +
 .../executor/test_kv_cache_budget_split.py    |  32 ---
 .../executor/test_pytorch_model_engine.py     |  53 +---
 .../_torch/modeling/test_gemma4_multimodal.py |  16 ++
 .../_torch/modeling/test_modeling_gemma4.py   |  52 +++-
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 229 +++++++++++-------
 26 files changed, 890 insertions(+), 467 deletions(-)

diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py
index 265303c95a09..07d04206bde5 100644
--- a/examples/llm-api/quickstart_advanced.py
+++ b/examples/llm-api/quickstart_advanced.py
@@ -311,6 +311,8 @@ def setup_llm(args, **kwargs):
             if args.draft_model_dir is None:
                 raise ValueError(
                     "--draft_model_dir is required for two-model MTP")
+        speculative_model = (args.draft_model_dir if args.draft_model_dir
+                             is not None else args.model_dir)
         spec_config = MTPDecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
             use_relaxed_acceptance_for_thinking=args.
@@ -321,8 +323,7 @@ def setup_llm(args, **kwargs):
             use_dynamic_tree=args.use_dynamic_tree,
             dynamic_tree_max_topK=args.dynamic_tree_max_topK,
             max_total_draft_tokens=args.max_total_draft_tokens,
-            speculative_model=args.model_dir
-            if args.use_one_model else args.draft_model_dir)
+            speculative_model=speculative_model)
     elif spec_decode_algo == "EAGLE3":
         spec_config = Eagle3DecodingConfig(
             max_draft_len=args.spec_decode_max_draft_len,
diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md
index c80a9ab5c9c2..da4b166b76cb 100644
--- a/examples/models/core/gemma/README.md
+++ b/examples/models/core/gemma/README.md
@@ -49,14 +49,18 @@ The `/v1/chat/completions` endpoint applies the Gemma 4 chat template automatica
 
 ### 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:
+Gemma 4 supports Multi-Token Prediction (MTP) speculative decoding through the
+one-model PyTorch execution path. The target and its matching assistant
+checkpoint are loaded into one engine, and the Q-only assistant reads the
+target model's KV cache. Create a server configuration for the
+target/assistant pair:
 
 ```bash
 cat > gemma4_mtp.yaml <<'EOF'
 speculative_config:
   decoding_type: MTP
   max_draft_len: 3
-  mtp_eagle_one_model: false
+  mtp_eagle_one_model: true
   speculative_model: google/gemma-4-E4B-it-assistant
 kv_cache_config:
   enable_block_reuse: false
@@ -68,7 +72,11 @@ trtllm-serve google/gemma-4-E4B-it \
     --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.
+The assistant shares the target model's KV cache, so TensorRT-LLM does not
+allocate a second full-size GPU KV cache for it. The current implementation
+supports the `Gemma4AssistantForCausalLM` assistants for E2B, E4B, 26B-A4B,
+and 31B. Gemma 4 12B uses the `Gemma4UnifiedForConditionalGeneration` and
+`Gemma4UnifiedAssistantForCausalLM` architectures, which are not supported.
 
 For offline inference, pass both checkpoints to the advanced LLM API example:
 
@@ -78,7 +86,6 @@ python3 examples/llm-api/quickstart_advanced.py \
     --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
 ```
diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 4a8df996b199..8c222b98bce7 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import functools
 import math
 import os
@@ -270,6 +271,12 @@ class FlashInferAttentionMetadata(AttentionMetadata):
 
     _multi_item_params: Optional[FlashInferMultiItemParams] = field(
         init=False, default=None)
+    _shared_kv_draft_metadata: Optional["FlashInferAttentionMetadata"] = field(
+        init=False, default=None, repr=False)
+    _shared_kv_runtime_lens: torch.Tensor = field(init=False, repr=False)
+    _is_shared_kv_draft_view: bool = field(init=False,
+                                           default=False,
+                                           repr=False)
 
     def needs_plan(self, plan_params: PlanParams) -> bool:
         if plan_params not in self._plan_params_to_wrappers:
@@ -291,6 +298,14 @@ def get_decode_wrapper(
     ) -> flashinfer.BatchDecodeWithPagedKVCacheWrapper:
         assert plan_params in self._plan_params_to_wrappers, "Plan params not found, make sure to call plan()"
         result = self._plan_params_to_wrappers[plan_params].decode_wrapper
+        if self._is_shared_kv_draft_view:
+            if result._backend != "trtllm-gen":
+                raise ValueError(
+                    "The shared-target-KV draft metadata view requires the "
+                    "FlashInfer trtllm-gen decode backend.")
+            num_seqs = self.num_seqs
+            result._kv_lens_buffer[:num_seqs].copy_(
+                self._shared_kv_runtime_lens[:num_seqs])
         return result
 
     def get_ragged_prefill_wrapper(
@@ -581,6 +596,130 @@ def batch_indices(self) -> torch.Tensor:
     def positions(self) -> torch.Tensor:
         return self._positions[:self.num_tokens]
 
+    def get_shared_kv_draft_metadata(self) -> "FlashInferAttentionMetadata":
+        """Return a one-query decode view over this metadata's target KV."""
+        if self._shared_kv_draft_metadata is None:
+            draft_metadata = copy.copy(self)
+            draft_metadata._is_shared_kv_draft_view = True
+            draft_metadata._shared_kv_draft_metadata = None
+            draft_metadata.workspace_buffer = self.workspace_buffer
+            draft_metadata.cuda_graph_buffers = None
+            draft_metadata.cross = None
+            draft_metadata.max_num_tokens = self.max_num_requests
+            draft_metadata._seq_lens = None
+            draft_metadata._seq_lens_cuda = None
+            draft_metadata._seq_lens_kv = None
+            draft_metadata._seq_lens_kv_cuda = None
+            draft_metadata._saved_tensors = {}
+            draft_metadata.seq_lens = torch.ones((self.max_num_requests, ),
+                                                 dtype=torch.int)
+            draft_metadata.seq_lens_kv = None
+            draft_metadata.num_contexts = 0
+            draft_metadata.__post_init__()
+            self._shared_kv_draft_metadata = draft_metadata
+            draft_metadata._sync_shared_kv_draft_view(self)
+        return self._shared_kv_draft_metadata
+
+    def _sync_shared_kv_draft_view(
+            self, target: "FlashInferAttentionMetadata") -> None:
+        """Refresh host-planned page tables before target graph replay."""
+        if not self._is_shared_kv_draft_view:
+            raise RuntimeError("Only a shared-KV draft metadata view can sync")
+
+        num_seqs = target.num_seqs
+        self.request_ids = target.request_ids
+        self.prompt_lens = target.prompt_lens
+        self.kv_cache_params = target.kv_cache_params
+        self.kv_cache_manager = target.kv_cache_manager
+        self.all_rank_num_tokens = target.all_rank_num_tokens
+        self.seq_lens = torch.ones((num_seqs, ), dtype=torch.int)
+        self.seq_lens_kv = None
+        self.num_contexts = 0
+        self.num_blocks = list(target.num_blocks)
+        self.num_context_blocks = 0
+        self.num_generation_blocks = sum(self.num_blocks)
+        self.num_ctx_cached_tokens = 0
+        self._multi_item_params = None
+
+        total_blocks = self.num_generation_blocks
+        self._paged_kv_indices[:total_blocks].copy_(
+            target._paged_kv_indices[:total_blocks], non_blocking=True)
+        self._host_paged_kv_indices = target._host_paged_kv_indices
+
+        if (target._vswa_layer_to_pool is not None
+                and target._vswa_pool_indices_cache is not None):
+            self._vswa_pool_indices_cache = {}
+            self._host_pool_indices = dict(target._host_pool_indices)
+            for pool_id, source in target._vswa_pool_indices_cache.items():
+                destination = getattr(self, f"_vswa_pool_buf_{pool_id}")
+                destination[:total_blocks].copy_(source[:total_blocks],
+                                                 non_blocking=True)
+                self._vswa_pool_indices_cache[pool_id] = destination
+            primary_pool_id = self._vswa_layer_to_pool.get(0, 0)
+            self._vswa_active_pool_id = primary_pool_id
+            self._host_paged_kv_indices = self._host_pool_indices[
+                primary_pool_id]
+            self._paged_kv_indices[:total_blocks].copy_(
+                self._vswa_pool_indices_cache[primary_pool_id][:total_blocks],
+                non_blocking=True)
+        else:
+            self._vswa_pool_indices_cache = None
+            self._host_pool_indices = {}
+            self._vswa_active_pool_id = None
+
+        host_indptr = maybe_pin_memory(
+            torch.from_numpy(
+                np.concatenate([[0], np.cumsum(self.num_blocks)
+                                ]).astype(np.int32, copy=False)))
+        self.paged_kv_indptr_decode[:num_seqs + 1].copy_(host_indptr,
+                                                         non_blocking=True)
+        self._host_paged_kv_indptr_decode = host_indptr
+        self.paged_kv_indptr_prefill[0].zero_()
+        self.paged_kv_indptr = self.paged_kv_indptr_decode[:num_seqs + 1]
+        self._paged_kv_last_page_len[:num_seqs].copy_(
+            target._paged_kv_last_page_len[:num_seqs], non_blocking=True)
+        self._cached_token_lens[:num_seqs].copy_(
+            target._cached_token_lens[:num_seqs], non_blocking=True)
+
+        host_qo_indptr = torch.arange(num_seqs + 1,
+                                      dtype=torch.int32,
+                                      pin_memory=prefer_pinned())
+        self._qo_indptr[:num_seqs + 1].copy_(host_qo_indptr, non_blocking=True)
+        full_kv_lens = (target._cached_token_lens[:num_seqs] +
+                        target.seq_lens_kv_cuda[:num_seqs])
+        self._shared_kv_runtime_lens[:num_seqs].copy_(full_kv_lens)
+
+        # Re-plan every known assistant wrapper outside graph capture. The
+        # assistant forces trtllm-gen, whose decode plan does not retain
+        # workspace state shared with another wrapper.
+        self._clean_cached_plans(defer_plan=False)
+
+    def update_shared_kv_draft_lengths(
+        self,
+        target: "FlashInferAttentionMetadata",
+        num_accepted_tokens: torch.Tensor,
+        num_contexts: int,
+    ) -> None:
+        """Publish the accepted target prefix to assistant decode wrappers."""
+        if not self._is_shared_kv_draft_view:
+            raise RuntimeError(
+                "Draft KV lengths can only be set on the shared-KV view")
+        num_seqs = target.num_seqs
+        cached_lens = target._cached_token_lens[:num_seqs]
+        runtime_lens = self._shared_kv_runtime_lens[:num_seqs]
+        if num_contexts > 0:
+            runtime_lens[:num_contexts].copy_(
+                cached_lens[:num_contexts] +
+                target.seq_lens_kv_cuda[:num_contexts])
+        runtime_lens[num_contexts:num_seqs].copy_(
+            cached_lens[num_contexts:num_seqs] +
+            num_accepted_tokens[num_contexts:num_seqs])
+        self._cached_token_lens[:num_seqs].copy_(runtime_lens)
+        torch.remainder(runtime_lens - 1,
+                        self.page_size,
+                        out=self._paged_kv_last_page_len[:num_seqs])
+        self._paged_kv_last_page_len[:num_seqs].add_(1)
+
     def __post_init__(self) -> None:
         super().__post_init__()
         self._post_init_with_buffers(self.cuda_graph_buffers)
@@ -635,6 +774,9 @@ def _post_init_with_buffers(self, buffers) -> None:
         self._cached_token_lens = torch.empty((self.max_num_requests, ),
                                               dtype=torch.int,
                                               device='cuda')
+        self._shared_kv_runtime_lens = torch.empty((self.max_num_requests, ),
+                                                   dtype=torch.int,
+                                                   device='cuda')
         self._batch_indices = torch.empty((self.max_num_tokens, ),
                                           dtype=torch.int,
                                           device='cuda')
@@ -759,6 +901,9 @@ def create_cuda_graph_metadata(self,
                                                       encode_only)
         metadata.max_num_requests = max_batch_size
         metadata.max_num_tokens = max_batch_size * (1 + max_draft_tokens)
+        # The graph owns distinct assistant wrappers and stable metadata
+        # buffers; never inherit the eager view through the shallow copy.
+        metadata._shared_kv_draft_metadata = None
         # Post init again to make sure all tensors are allocated
         metadata.__post_init__()
         return metadata
@@ -993,7 +1138,8 @@ def _build_decode_block_tables(
         num_gens = self.num_generations
         if num_gens == 0:
             return None
-        host_paged_kv_indices = self._host_paged_kv_indices
+        host_paged_kv_indices = self._host_pool_indices.get(
+            plan_params.kv_pool_id, self._host_paged_kv_indices)
         if host_paged_kv_indices is None:
             return None
         gen_num_blocks = np.asarray(self.num_blocks[self.num_contexts:],
@@ -1444,6 +1590,10 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                                                              non_blocking=True)
                     if self.num_generations < batch_size:
                         kv_lens_buf[self.num_generations:batch_size].zero_()
+        if (not self._is_shared_kv_draft_view
+                and self._shared_kv_draft_metadata is not None):
+            self._shared_kv_draft_metadata._sync_shared_kv_draft_view(self)
+
         if self.cross is not None and self.cross is not self:
             self.cross.prepare()
 
diff --git a/tensorrt_llm/_torch/configs/gemma4.py b/tensorrt_llm/_torch/configs/gemma4.py
index e2942239790e..7609ae3c7fa6 100644
--- a/tensorrt_llm/_torch/configs/gemma4.py
+++ b/tensorrt_llm/_torch/configs/gemma4.py
@@ -27,11 +27,14 @@ class Gemma4AssistantConfig(PreTrainedConfig):
     model_type = "gemma4_assistant"
     sub_configs = {"text_config": Gemma4TextConfig}
 
-    # Runtime capabilities consumed by the generic two-model MTP pipeline.
+    # Runtime ownership contract for the one-model speculative pipeline.
+    loads_external_weights = True
+    num_draft_modules = 1
+    owns_independent_kv_cache = False
+    num_draft_kv_layers = 0
     shares_target_kv_cache = True
-    preserve_checkpoint_layer_count = True
     freezes_draft_attention_state = True
-    cuda_graph_external_draft_len = 0
+    requires_external_draft_metadata_view = True
 
     def __init__(
         self,
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index ea8705030b68..4d46246dbb32 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -59,7 +59,7 @@
 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_speculative import SpecDecOneEngineForCausalLM, _slice_spec_position_ids
 from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model
 
 _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0"
@@ -1384,6 +1384,8 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
         local_attention_mask_data = None
+        resource_manager = kwargs.pop("resource_manager", None)
+        orig_input_ids = kwargs.pop("orig_input_ids", None)
         # Only build bidirectional masks when use_bidirectional_attention is
         # set to "vision" (26B, 31B).  E2B/E4B have this as None and should
         # use standard causal attention even for multimodal tokens. Gemma4
@@ -1412,6 +1414,38 @@ def forward(
         if attn_metadata.padded_num_tokens is not None:
             output = output[: attn_metadata.num_tokens]
 
+        if self.spec_worker is not None:
+            logits = self.logits_processor.forward(
+                output[spec_metadata.gather_ids],
+                self.lm_head,
+                attn_metadata,
+                True,
+            )
+            if self.config.final_logit_softcapping is not None:
+                cap = self.config.final_logit_softcapping
+                logits = torch.tanh(logits / cap) * cap
+
+            spec_input_ids = input_ids if input_ids is not None else orig_input_ids
+            spec_position_ids = position_ids
+            if attn_metadata.padded_num_tokens is not None:
+                if spec_input_ids is not None:
+                    spec_input_ids = spec_input_ids[: attn_metadata.num_tokens]
+                if position_ids is not None:
+                    spec_position_ids = _slice_spec_position_ids(
+                        position_ids, attn_metadata.num_tokens
+                    )
+
+            return self.spec_worker(
+                input_ids=spec_input_ids,
+                position_ids=spec_position_ids,
+                hidden_states=output,
+                logits=logits,
+                attn_metadata=attn_metadata,
+                spec_metadata=spec_metadata,
+                draft_model=self.draft_model,
+                resource_manager=resource_manager,
+            )
+
         logits = self.logits_processor.forward(
             output,
             self.lm_head,
@@ -1608,6 +1642,32 @@ def _last_token_states(
         )
         return hidden_states[last_tokens]
 
+    def forward_draft_step(
+        self,
+        input_ids: torch.IntTensor,
+        position_ids: torch.IntTensor,
+        recurrent_hidden_states: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        spec_metadata=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Run one Q-only assistant step over a frozen target KV prefix."""
+        target_embeddings = self._get_target_embeddings(input_ids)
+        assistant_inputs = self.pre_projection(
+            torch.cat([target_embeddings, recurrent_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)
+        if self.masked_embedding is not None:
+            logits = self.masked_embedding(assistant_hidden_states, self.lm_head).float()
+        else:
+            logits = self.lm_head(assistant_hidden_states).float()
+        return logits, projected_hidden_states
+
     def forward(
         self,
         attn_metadata: AttentionMetadata,
@@ -1620,31 +1680,22 @@ def forward(
     ) -> 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(
+        logits, projected_hidden_states = self.forward_draft_step(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            recurrent_hidden_states=spec_metadata.get_hidden_states(),
             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).float()
-        return self.lm_head(logits_hidden_states).float()
+            return logits
+        last_token_indices = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
+        return logits[last_token_indices]
 
     def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         weights = weight_mapper.preprocess_weights(weights)
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 67622709c417..a344a800d984 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -732,6 +732,33 @@ def post_config(self):
         self.config = self.llm.config
         self.model_config.pretrained_config = self.llm.config
 
+    @property
+    def model(self):
+        return self.llm.model
+
+    @property
+    def lm_head(self):
+        return self.llm.lm_head
+
+    @property
+    def epilogue(self):
+        return self.llm.epilogue
+
+    @property
+    def spec_worker(self):
+        return self.llm.spec_worker
+
+    @property
+    def draft_config(self):
+        return self.llm.draft_config
+
+    @property
+    def draft_model(self):
+        return self.llm.draft_model
+
+    def load_draft_weights(self, weights, weight_mapper=None):
+        self.llm.load_draft_weights(weights, weight_mapper)
+
     @property
     def language_model(self) -> torch.nn.Module:
         return self.llm
@@ -743,6 +770,8 @@ def get_language_model_extra_forward_kwargs(
         position_ids: Optional[torch.Tensor],
         mm_inputs: PreparedLlmInputs,
         lora_params=None,
+        spec_metadata=None,
+        resource_manager=None,
         **forward_kwargs,
     ) -> Dict:
         """Build Gemma4-specific language-model forward arguments."""
@@ -770,6 +799,9 @@ def get_language_model_extra_forward_kwargs(
             "mm_token_type_ids": mm_token_type_ids,
             "ple_input_ids": ple_input_ids,
             "lora_params": lora_params,
+            "spec_metadata": spec_metadata,
+            "resource_manager": resource_manager,
+            "orig_input_ids": raw_input_ids,
         }
 
     @property
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index a0964b63dece..bca7d4e1cf53 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -30,8 +30,8 @@
 except ImportError:
     _flashinfer_rope = None
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
-from ..speculative import (SpecMetadata, get_spec_worker,
-                           should_use_separate_draft_kv_cache)
+from ..speculative import (DraftModelCapabilities, SpecMetadata,
+                           get_spec_worker, should_use_separate_draft_kv_cache)
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_auto import AutoModelForCausalLM
@@ -1889,6 +1889,11 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
+    elif (spec_dec_mode.is_mtp_eagle_one_model() and draft_config is not None
+          and getattr(model_config.spec_config, "_draft_model_capabilities",
+                      None) is not None and model_config.spec_config.
+          _draft_model_capabilities.loads_external_weights):
+        return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
         return MTPForCausalLM(model_config,
                               model_config.pretrained_config.num_hidden_layers,
@@ -1992,6 +1997,51 @@ def __init__(self,
                     model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
 
+                elif (spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                      and spec_config.speculative_model is not None
+                      and model_config.pretrained_config.model_type
+                      in ("gemma4", "gemma4_text")):
+                    self.draft_config = ModelConfig.from_pretrained(
+                        spec_config.speculative_model,
+                        trust_remote_code=True,
+                        attn_backend=model_config.attn_backend,
+                        moe_backend=model_config.moe_backend,
+                        mapping=model_config.mapping,
+                        spec_config=None,
+                        max_num_tokens=model_config.max_num_tokens,
+                        moe_max_num_tokens=model_config.moe_max_num_tokens)
+                    self.draft_config.quant_config.kv_cache_quant_algo = \
+                        model_config.quant_config.kv_cache_quant_algo
+                    self.draft_config.extra_attrs = model_config.extra_attrs
+                    draft_pretrained_config = (
+                        self.draft_config.pretrained_config)
+                    draft_architectures = getattr(draft_pretrained_config,
+                                                  "architectures", None) or []
+                    if "Gemma4AssistantForCausalLM" in draft_architectures:
+                        # Newer Transformers releases provide a native
+                        # Gemma4AssistantConfig without TRT-LLM runtime
+                        # ownership attributes. Its architecture has the same
+                        # all-Q-only, shared-target-KV contract.
+                        capabilities = (
+                            DraftModelCapabilities.external_shared_target_kv())
+                    else:
+                        capabilities = DraftModelCapabilities.from_config(
+                            draft_pretrained_config)
+                    if not (capabilities.loads_external_weights
+                            and capabilities.shares_target_kv_cache
+                            and not capabilities.owns_independent_kv_cache
+                            and capabilities.num_draft_modules == 1
+                            and capabilities.num_draft_kv_layers == 0
+                            and capabilities.freezes_draft_attention_state and
+                            capabilities.requires_external_draft_metadata_view):
+                        raise ValueError(
+                            "External one-model MTP assistants must load "
+                            "external weights, expose one logical draft module, "
+                            "own zero KV layers, share the target KV cache, "
+                            "freeze draft attention state, and require an "
+                            "external draft metadata view.")
+                    spec_config._draft_model_capabilities = capabilities
+
                 elif spec_config.spec_dec_mode.is_external_drafter():
                     self.draft_config = ModelConfig.from_pretrained(
                         model_config.spec_config.speculative_model,
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index eb91a6456209..51be81a4cf4c 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -312,9 +312,9 @@ def __pp_init__(self):
         total_num_layers = num_hidden_layers
         spec_config = getattr(self.model_config, "spec_config", None)
         if spec_config is not None:
-            from ..speculative.utils import get_num_spec_layers
+            from ..speculative.utils import get_num_draft_kv_layers
 
-            num_spec_layers = get_num_spec_layers(spec_config) or 0
+            num_spec_layers = get_num_draft_kv_layers(spec_config) or 0
             total_num_layers += num_spec_layers
             if num_spec_layers > 0 and mapping.is_last_pp_rank():
                 pp_layer_list.extend(
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index b6373ca907b1..d5bc614a0862 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -44,7 +44,7 @@
 from ..hostfunc import set_low_latency_dispatch
 from ..model_config import ModelConfig
 from ..models.modeling_multimodal_mixin import MultimodalModelMixin
-from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers,
+from ..speculative import (get_num_draft_kv_layers, get_num_extra_kv_tokens,
                            get_spec_decoder, should_use_separate_draft_kv_cache)
 from ..utils import is_gdn_replay_enabled
 from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params,
@@ -1207,7 +1207,7 @@ def _get_num_draft_layers(self) -> int:
         """
         if self._speculative_config.spec_dec_mode.is_external_drafter():
             return self._draft_config.pretrained_config.num_hidden_layers
-        return get_num_spec_layers(self._speculative_config)
+        return get_num_draft_kv_layers(self._speculative_config)
 
     def _create_one_model_draft_kv_cache_manager(
         self,
@@ -1699,16 +1699,11 @@ 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.
-        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
+        has_draft = (
+            self._draft_model_engine is not None  # two-model
             or self._should_create_separate_draft_kv_cache())  # one-model
         draft_kv_cache_config = None
-        if not estimating_kv_cache and has_independent_draft_cache:
+        if not estimating_kv_cache and has_draft:
             # 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):
@@ -1741,10 +1736,8 @@ def build_managers(self,
                                        if draft_kv_cache_config is not None else
                                        self_kv_cache_config)
 
-        # 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):
+        # Two-model speculative decoding: draft model has separate engine
+        if self._draft_model_engine is not None:
             if self._is_kv_cache_manager_v2:
                 assert draft_kv_cache_config is None, (
                     "KVCacheManagerV2 does not support two-model speculative "
@@ -1804,7 +1797,7 @@ def _build_per_layer_num_kv_heads(
     if spec_config is None or draft_config is None:
         return num_key_value_heads
 
-    from ..speculative.utils import get_num_spec_layers
+    from ..speculative.utils import get_num_draft_kv_layers
     draft_pretrained = draft_config.pretrained_config
     draft_num_kv_heads = getattr(
         draft_pretrained, 'num_key_value_heads',
@@ -1813,7 +1806,7 @@ def _build_per_layer_num_kv_heads(
     if draft_num_kv_heads is None or draft_num_kv_heads == num_key_value_heads:
         return num_key_value_heads
 
-    num_spec_layers = get_num_spec_layers(spec_config)
+    num_spec_layers = get_num_draft_kv_layers(spec_config)
     logger.info(f"Per-layer KV heads for speculative decoding: "
                 f"target={num_key_value_heads} x {num_hidden_layers} layers, "
                 f"draft={draft_num_kv_heads} x {num_spec_layers} layers, "
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index c64b49646862..05343d0deddd 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -111,7 +111,6 @@ 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,8 +158,6 @@ 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
@@ -273,18 +270,11 @@ 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.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'.
-                # 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)
+            # 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)
         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 5c247f3233b8..06e55f838931 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -66,9 +66,9 @@
                            get_num_extra_kv_tokens, get_spec_metadata,
                            prepare_attn_metadata_for_draft_replay,
                            restore_attn_metadata_after_draft_replay,
+                           should_extend_context,
                            update_spec_config_from_loaded_model)
-from ..speculative.drafting_loops import (BaseDraftingLoopWrapper,
-                                          get_draft_model_capability)
+from ..speculative.drafting_loops import BaseDraftingLoopWrapper
 from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata
 from ..speculative.spec_sampler_base import SampleStateTensorsSpec
 from ..utils import (get_model_extra_attrs,
@@ -312,8 +312,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], Optional[torch.nn.Module]]] = None,
+        drafting_loop_wrapper: Optional[Callable[[torch.nn.Module],
+                                                 torch.nn.Module]] = None,
         model: Optional[torch.nn.Module] = None,
         checkpoint_loader: Optional[BaseCheckpointLoader] = None,
         model_weights_memory_tag: Optional[str] = None,
@@ -462,19 +462,10 @@ def __init__(
             enable_overlap_headroom=self._enable_dsv4_overlap_headroom,
         )
         if drafting_loop_wrapper is not None:
-            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
+            self.model = drafting_loop_wrapper(self.model)
+            self.model_is_wrapped = True
         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'):
@@ -709,13 +700,9 @@ 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,
-            cuda_graph_max_total_draft_tokens,
+            self.original_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
@@ -805,10 +792,7 @@ 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
-                                     and not self._shares_target_kv_cache else
-                                     ResourceManagerType.KV_CACHE_MANAGER)
+        self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER
         self.lora_model_config: Optional[LoraModelConfig] = None
         self._trtllm_gen_jit_warmup = False
 
@@ -835,7 +819,6 @@ def __init__(
             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)
@@ -1568,12 +1551,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._cuda_graph_external_draft_len is not None):
-            token_num_upper_bound = 1
-        else:
-            token_num_upper_bound = min(
-                self.max_num_tokens, self.batch_size * (self.max_seq_len - 1))
+        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)
@@ -2535,15 +2514,10 @@ def _update_draft_inference_state_for_warmup(
             ResourceManagerType.SPEC_RESOURCE_MANAGER)
         if self.is_draft_model and isinstance(spec_resource_manager,
                                               Eagle3ResourceManager):
-            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)
+            spec_resource_manager.is_first_draft = is_first_draft
             if is_first_draft:
                 for req in batch.generation_requests:
-                    req.py_is_first_draft = not freezes_draft_attention_state
+                    req.py_is_first_draft = True
                     req.py_draft_tokens = []
 
     def _set_up_attn_metadata(
@@ -3418,8 +3392,9 @@ def _prepare_incremental_update_metadata(
         attn_metadata.beam_width = 1
         attn_metadata.prompt_lens = prompt_lengths
         attn_metadata.num_contexts = num_extend_ctx_requests if (
-            enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
-                self.attn_backend) and spec_config.is_linear_tree) else 0
+            enable_spec_decode
+            and should_extend_context(spec_config, self.attn_backend)
+            and spec_config.is_linear_tree) else 0
         attn_metadata.num_chunked_ctx_requests = attn_metadata.num_contexts
 
         # Create KV cache params and prepare metadata
@@ -3698,9 +3673,8 @@ def _apply_incremental_update_target(
         num_extend_dummy_requests = 0
         num_previous_batch = 0
 
-        use_extend_ctx = (self.enable_spec_decode
-                          and spec_config.spec_dec_mode.extend_ctx(
-                              self.attn_backend) and spec_config.is_linear_tree)
+        use_extend_ctx = (self.enable_spec_decode and should_extend_context(
+            spec_config, self.attn_backend) and spec_config.is_linear_tree)
 
         for idx, request in enumerate(extend_requests):
             request_accepted_path[request.py_request_id] = \
@@ -3773,8 +3747,8 @@ def _apply_incremental_update_target(
 
         # Determine if we're using extend_ctx mode for linear tree decoding
         num_extend_ctx_requests = 0
-        if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
-                self.attn_backend) and spec_config.is_linear_tree:
+        if self.enable_spec_decode and should_extend_context(
+                spec_config, self.attn_backend) and spec_config.is_linear_tree:
             num_extend_ctx_requests = num_extend_requests
 
         virtual_num_tokens = num_generation_tokens
@@ -4357,7 +4331,8 @@ def append_cross_attention_state(request: LlmRequest,
                                        if is_promoted_context else
                                        request.max_beam_num_tokens - 1)
                 draft_lens.append(num_draft_tokens)
-                if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                if self.enable_spec_decode and should_extend_context(
+                        spec_config,
                         self.attn_backend) and spec_config.is_linear_tree:
                     # We're treating the prompt lengths as context requests here, so
                     # the the prompt lens should not include the cached tokens.
@@ -4409,7 +4384,8 @@ def append_cross_attention_state(request: LlmRequest,
                     request.py_num_compressed_tokens)
                 request.cached_tokens = (past_seen_token_num +
                                          runtime_tokens_per_gen_step)
-                if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                if self.enable_spec_decode and should_extend_context(
+                        spec_config,
                         self.attn_backend) and spec_config.is_linear_tree:
                     prompt_lengths.append(runtime_tokens_per_gen_step)
                 else:
@@ -5007,8 +4983,8 @@ def previous_seq_slots_device():
         # Use num_chunked_ctx_requests to record the number of extend context requests,
         # so that we can update the kv_lens_cuda correctly in _preprocess_inputs.
         attn_metadata.num_chunked_ctx_requests = 0
-        if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
-                self.attn_backend) and spec_config.is_linear_tree:
+        if self.enable_spec_decode and should_extend_context(
+                spec_config, self.attn_backend) and spec_config.is_linear_tree:
             # For the tree decoding, we want to use XQA to process the draft tokens for the target model.
             # Therefore, we do not treat them as the chunked context requests.
             attn_metadata.num_contexts += len(extend_requests)
@@ -5653,8 +5629,8 @@ def _get_lora_params_from_requests(
             tokens_per_seq = 1
             if (self.enable_spec_decode and self.runtime_draft_len > 0
                     and self.spec_config.is_linear_tree
-                    and not self.spec_config.spec_dec_mode.extend_ctx(
-                        self.attn_backend)):
+                    and not should_extend_context(self.spec_config,
+                                                  self.attn_backend)):
                 tokens_per_seq = self.runtime_draft_len + 1
             return self.cuda_graph_lora_manager.prepare_cuda_graph_lora_params(
                 scheduled_requests, attn_metadata, peft_cache_manager,
@@ -5773,8 +5749,9 @@ def _get_eager_lora_params_from_requests(
             # count so the kernel correctly expands LoRA weights for all tokens.
             if (self.enable_spec_decode and self.runtime_draft_len > 0
                     and self.spec_config.is_linear_tree
-                    and not self.spec_config.spec_dec_mode.extend_ctx(
-                        self.attn_backend) and num_generations > 0):
+                    and not should_extend_context(self.spec_config,
+                                                  self.attn_backend)
+                    and num_generations > 0):
                 tokens_per_req = self.runtime_draft_len + 1
                 host_request_types = host_request_types.clone()
                 host_request_types[num_contexts:num_seqs].fill_(0)  # kCONTEXT
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 5d7bd983c8df..238053adcaf9 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -41,6 +41,7 @@
                                      timing)
 from ..modules.fused_moe.moe_load_balancer import (
     MoeLoadBalancer, maybe_create_moe_load_balancer)
+from ..speculative import needs_external_draft_weights
 from ..virtual_memory import RestoreMode
 from ..virtual_memory import scope as virtual_memory_scope
 from .config_utils import (is_hybrid_linear, resolve_hf_torch_dtype,
@@ -551,9 +552,7 @@ def load(
                 model = AutoModelForCausalLM.from_config(config)
                 is_meta_init = False
 
-            loads_draft_weights = (
-                self.spec_config is not None
-                and self.spec_config.spec_dec_mode.need_load_draft_weights())
+            loads_draft_weights = needs_external_draft_weights(self.spec_config)
             speculative_mode = self._speculative_mode_name(self.spec_config)
             post_transform_qualification = self._qualify_post_transform_profile(
                 model,
@@ -704,8 +703,7 @@ def init_meta_tensor(t: torch.Tensor):
                     self._call_load_weights(model.load_weights, weights,
                                             self.weight_mapper)
 
-                if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
-                ):
+                if needs_external_draft_weights(self.spec_config):
                     weights = checkpoint_loader.load_weights(
                         self.spec_config.speculative_model,
                         mapping=self.mapping)
@@ -840,8 +838,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                                     "commit an unpopulated model to the GMS "
                                     "pool.")
 
-                            if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
-                            ):
+                            if needs_external_draft_weights(self.spec_config):
                                 draft_weights = checkpoint_loader.load_weights(
                                     self.spec_config.speculative_model,
                                     mapping=self.mapping)
@@ -968,8 +965,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                 self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
                     model, config)
                 initialize_dummy_weights(model)
-                if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
-                ):
+                if needs_external_draft_weights(self.spec_config):
                     model.draft_model.load_weights_from_target_model(model)
 
             elif load_format == LoadFormat.VISION_ONLY:
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 43f21df70741..35733e8bc03b 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -40,7 +40,8 @@
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
                     validate_feature_combination)
-from .config_utils import is_hybrid_linear, is_minimax_m3
+from .config_utils import (is_hybrid_linear, is_minimax_m3,
+                           load_pretrained_config)
 from .connectors.kv_cache_connector import KvCacheConnectorManager
 from .dwdp import DwdpManager
 from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder
@@ -451,8 +452,19 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
-        # Check FLASHINFER compatibility with one-engine speculative decoding
-        if llm_args.attn_backend == "FLASHINFER":
+        # External MTP assistants with shared target KV use a dedicated
+        # FlashInfer decode metadata view. Other one-engine modes still rely
+        # on the one-query-per-sequence decode contract.
+        supports_shared_kv_flashinfer = False
+        if (spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                and spec_config.speculative_model is not None
+                and checkpoint_dir is not None):
+            target_config = load_pretrained_config(
+                checkpoint_dir, trust_remote_code=llm_args.trust_remote_code)
+            supports_shared_kv_flashinfer = target_config.model_type in (
+                "gemma4", "gemma4_text")
+        if (llm_args.attn_backend == "FLASHINFER"
+                and not supports_shared_kv_flashinfer):
             raise ValueError(
                 f"FLASHINFER attention backend is not supported with one-engine speculative "
                 f"decoding mode '{spec_config.spec_dec_mode.name}'. The FLASHINFER backend's "
@@ -585,44 +597,22 @@ def allocation_scope(current_stage: ExecutorMemoryType):
         with allocation_scope(ExecutorMemoryType.MODEL_ENGINE_DRAFT):
             draft_spec_config = copy.copy(spec_config)
 
-            capturable_drafter_eligible = (
+            use_chain_drafter = (
                 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
-                    or draft_spec_config.spec_dec_mode.is_mtp_eagle()):
+            if use_chain_drafter:
 
                 def drafting_loop_wrapper(model):
                     from tensorrt_llm._torch.speculative.drafting_loops import (
-                        Gemma4AssistantDraftingLoopWrapper,
                         LinearDraftingLoopWrapper,
                         StaticTreeDraftingLoopWrapper)
                     from tensorrt_llm.llmapi import EagleDecodingConfig
 
-                    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)
-
-                    if not use_chain_drafter:
-                        return None
-
                     static_tree_drafter = isinstance(
                         draft_spec_config, EagleDecodingConfig
                     ) and draft_spec_config.eagle_choices is not None
@@ -662,11 +652,8 @@ def drafting_loop_wrapper(model):
                 model_weights_memory_tag=model_weights_memory_tag,
                 model_weights_restore_mode=model_weights_restore_mode,
             )
-            # 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 not getattr(draft_model_engine.model.config,
-                                    "preserve_checkpoint_layer_count", False)):
+            # 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():
                 draft_model_engine.model.model_config.pretrained_config.num_hidden_layers = 1
             draft_model_engine.load_weights_from_target_model(
                 model_engine.model)
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index b4faf969d0da..4a4abe503583 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -168,7 +168,7 @@ def get_pp_layers(
     spec_config: Optional["DecodingBaseConfig"] = None,
     layer_mask: Optional[List[bool]] = None,
 ) -> Tuple[List[int], int]:
-    from ..speculative.utils import get_num_spec_layers
+    from ..speculative.utils import get_num_draft_kv_layers
 
     total_num_layers = num_layers
     if layer_mask is not None:
@@ -193,7 +193,7 @@ def get_pp_layers(
     # When layer_mask is provided, the caller explicitly controls which layers
     # to include, so we should not add extra layers automatically.
     if spec_config is not None and layer_mask is None:
-        num_spec_layers = get_num_spec_layers(spec_config)
+        num_spec_layers = get_num_draft_kv_layers(spec_config)
         total_num_layers += num_spec_layers
         if mapping.is_last_pp_rank():
             pp_layers.extend(
diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py
index 347f3aae5dc2..07f2127c4e00 100644
--- a/tensorrt_llm/_torch/speculative/__init__.py
+++ b/tensorrt_llm/_torch/speculative/__init__.py
@@ -3,9 +3,12 @@
 from .draft_target import (DraftTargetOneModelSpecMetadata,
                            DraftTargetOneModelWorker)
 from .eagle3 import Eagle3SpecMetadata, MTPEagleWorker
-from .interface import (SpecMetadata, SpecWorkerBase,
+from .interface import (DraftModelCapabilities, SpecMetadata, SpecWorkerBase,
+                        get_draft_model_capabilities,
+                        needs_external_draft_weights,
                         prepare_attn_metadata_for_draft_replay,
                         restore_attn_metadata_after_draft_replay,
+                        should_extend_context,
                         should_use_separate_draft_kv_cache)
 from .mtp import MTPSampler, MTPSpecMetadata, MTPWorker
 from .ngram import NGramDrafter, NGramPoolManager
@@ -18,16 +21,18 @@
                                 SpecSamplerBase)
 from .spec_tree_manager import SpecTreeManager
 from .suffix_automaton import SuffixAutomatonManager
-from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens,
-                    get_num_spec_layers, get_spec_decoder, get_spec_drafter,
-                    get_spec_metadata, get_spec_resource_manager,
-                    get_spec_worker, update_spec_config_from_draft_model_config,
+from .utils import (get_draft_kv_cache_manager, get_num_draft_kv_layers,
+                    get_num_extra_kv_tokens, get_num_spec_layers,
+                    get_spec_decoder, get_spec_drafter, get_spec_metadata,
+                    get_spec_resource_manager, get_spec_worker,
+                    update_spec_config_from_draft_model_config,
                     update_spec_config_from_loaded_model,
                     update_spec_config_from_model_config)
 
 __all__ = [
     "DFlashSpecMetadata",
     "DFlashWorker",
+    "DraftModelCapabilities",
     "DraftTargetOneModelSpecMetadata",
     "DraftTargetOneModelWorker",
     "Eagle3SpecMetadata",
@@ -52,6 +57,8 @@
     "SpecSamplerBase",
     "SpecWorkerBase",
     "get_draft_kv_cache_manager",
+    "get_draft_model_capabilities",
+    "get_num_draft_kv_layers",
     "get_num_extra_kv_tokens",
     "get_num_spec_layers",
     "get_spec_decoder",
@@ -60,7 +67,9 @@
     "get_spec_resource_manager",
     "get_spec_worker",
     "prepare_attn_metadata_for_draft_replay",
+    "needs_external_draft_weights",
     "restore_attn_metadata_after_draft_replay",
+    "should_extend_context",
     "should_use_separate_draft_kv_cache",
     "update_spec_config_from_draft_model_config",
     "update_spec_config_from_loaded_model",
diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py
index 0254235aa3d5..133814bb2d22 100644
--- a/tensorrt_llm/_torch/speculative/drafting_loops.py
+++ b/tensorrt_llm/_torch/speculative/drafting_loops.py
@@ -1,5 +1,3 @@
-# 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.
 
@@ -11,7 +9,7 @@
 """
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from typing import Optional, final
 
 import torch
@@ -59,12 +57,6 @@ 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:
@@ -132,13 +124,11 @@ 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 self.drafting_metadata_context(attn_metadata, spec_metadata):
+            with save_metadata_state(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],
@@ -147,8 +137,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)
-                    self.advance_generation_state(new_position_ids,
-                                                  attn_metadata, batch_size)
+                    new_position_ids += 1
+                    attn_metadata.kv_lens_cuda[:batch_size] += 1
                     if i == 0 and is_eagle3:
                         spec_metadata.hidden_states_read_indices[:batch_size].copy_(
                             spec_metadata.
@@ -159,20 +149,6 @@ 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)
@@ -226,31 +202,6 @@ 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 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,
-                               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 1713958a1f95..9ac7afe65836 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -13,6 +13,7 @@
 from tensorrt_llm.mapping import Mapping
 
 from ..attention_backend import AttentionMetadata
+from ..attention_backend.flashinfer import FlashInferAttentionMetadata
 from ..model_config import ModelConfig
 from ..pyexecutor.llm_request import LlmRequest
 from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager
@@ -91,10 +92,6 @@ 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,
@@ -129,7 +126,6 @@ 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)
@@ -215,7 +211,6 @@ class Eagle3SpecMetadata(SpecMetadata):
     is_first_draft: bool = False
     eagle3_resource_manager: Optional[Eagle3ResourceManager] = None
     is_mtp_eagle: bool = False
-    shares_target_kv_cache: bool = False
 
     eagle_choices: Optional[List[List[int]]] = None
     max_total_draft_tokens: int = 0
@@ -283,28 +278,11 @@ 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]
-            # 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, (
-                    "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, (
-                    "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
-                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
-            elif not self.is_draft_model or (is_first_draft
-                                             and spec_tree_manager is None):
+            if 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(
@@ -679,6 +657,41 @@ def __init__(self,
         self._saved_position_offsets = None
         self._saved_position_offsets_cpp = None
         self._saved_generation_lengths = None
+        self._uses_external_shared_target_kv = False
+
+    def set_draft_model(self, draft_model) -> None:
+        super().set_draft_model(draft_model)
+        capabilities = getattr(self.spec_config, "_draft_model_capabilities",
+                               None)
+        self._uses_external_shared_target_kv = bool(
+            capabilities is not None and capabilities.loads_external_weights
+            and capabilities.shares_target_kv_cache
+            and not capabilities.owns_independent_kv_cache)
+        if not self._uses_external_shared_target_kv:
+            return
+        if self.use_dynamic_tree:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP supports only the "
+                "linear draft path.")
+        if self.spec_config.draft_len_schedule is not None:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP does not support a "
+                "draft length schedule.")
+        if self.sa_enhancer is not None:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP does not support the "
+                "suffix automaton enhancer.")
+        if self.spec_config.use_rejection_sampling:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP does not support "
+                "rejection sampling.")
+
+    def set_guided_decoder(self, guided_decoder) -> bool:
+        if self._uses_external_shared_target_kv:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP does not support "
+                "guided decoding.")
+        return super().set_guided_decoder(guided_decoder)
 
     @property
     def max_draft_len(self) -> int:
@@ -746,6 +759,11 @@ def _forward_impl(self,
                       resource_manager=None):
 
         runtime_draft_len = spec_metadata.runtime_draft_len
+        if self._uses_external_shared_target_kv:
+            return self._forward_external_shared_target_kv(
+                input_ids, position_ids, hidden_states, logits, attn_metadata,
+                spec_metadata, draft_model)
+
         # skip the draft forward if the runtime draft length is 0
         if runtime_draft_len == 0:
             return self.skip_drafting(input_ids, position_ids, hidden_states,
@@ -840,6 +858,134 @@ def _forward_impl(self,
             'next_new_tokens': next_new_tokens,
         }
 
+    def _forward_external_shared_target_kv(
+        self,
+        input_ids,
+        position_ids,
+        hidden_states,
+        logits,
+        attn_metadata,
+        spec_metadata,
+        draft_model,
+    ):
+        """Draft with an external Q-only assistant over accepted target KV."""
+        if not isinstance(attn_metadata, FlashInferAttentionMetadata):
+            raise TypeError(
+                "Gemma4 shared-target-KV one-model MTP currently requires "
+                "FlashInfer attention metadata.")
+        if self.guided_decoder is not None:
+            raise ValueError(
+                "Gemma4 shared-target-KV one-model MTP does not support "
+                "guided decoding.")
+
+        batch_size = attn_metadata.num_seqs
+        num_contexts = batch_size - spec_metadata.num_generations
+        runtime_draft_len = spec_metadata.runtime_draft_len
+        if runtime_draft_len == 0:
+            target_tokens = self._sample_tokens_for_batch(
+                logits, spec_metadata, num_contexts, batch_size)
+            accepted_tokens = target_tokens.unsqueeze(1)
+            num_accepted_tokens = torch.ones(batch_size,
+                                             dtype=torch.int,
+                                             device=logits.device)
+            next_draft_tokens = torch.empty((batch_size, 0),
+                                            dtype=torch.int32,
+                                            device=logits.device)
+            return {
+                "logits": logits,
+                "new_tokens": accepted_tokens,
+                "new_tokens_lens": num_accepted_tokens,
+                "next_draft_tokens": next_draft_tokens,
+                "next_new_tokens": accepted_tokens,
+            }
+
+        raw_logits = logits
+        accepted_tokens, num_accepted_tokens = (
+            self.sample_and_accept_draft_tokens(input_ids, logits,
+                                                attn_metadata, spec_metadata))
+
+        (
+            draft_input_ids,
+            recurrent_hidden_states,
+            draft_position_ids,
+        ) = self._prepare_external_shared_target_kv_draft_inputs(
+            accepted_tokens=accepted_tokens,
+            num_accepted_tokens=num_accepted_tokens,
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            sequence_lengths=attn_metadata.seq_lens_cuda[:batch_size],
+            num_contexts=num_contexts,
+            batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
+        )
+
+        draft_metadata = attn_metadata.get_shared_kv_draft_metadata()
+        draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
+                                                      num_accepted_tokens,
+                                                      num_contexts)
+        draft_metadata.use_spec_decoding = False
+        draft_metadata.padded_num_tokens = None
+        draft_metadata.all_rank_num_tokens = (
+            spec_metadata.subseq_all_rank_num_tokens)
+
+        next_draft_tokens = []
+        for draft_step in range(runtime_draft_len):
+            draft_logits, recurrent_hidden_states = (
+                draft_model.forward_draft_step(
+                    input_ids=draft_input_ids,
+                    position_ids=draft_position_ids,
+                    recurrent_hidden_states=recurrent_hidden_states,
+                    attn_metadata=draft_metadata,
+                    spec_metadata=spec_metadata,
+                ))
+            draft_input_ids = self.sample_draft_tokens(
+                draft_logits,
+                spec_metadata,
+                batch_size,
+                draft_step=draft_step,
+            )
+            next_draft_tokens.append(draft_input_ids)
+
+        next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
+        batch_indices = spec_metadata.batch_indices_cuda[:batch_size]
+        next_new_tokens = self._prepare_next_new_tokens(accepted_tokens,
+                                                        next_draft_tokens,
+                                                        batch_indices,
+                                                        batch_size,
+                                                        num_accepted_tokens)
+        attn_metadata.use_spec_decoding = True
+        return {
+            "logits": raw_logits,
+            "new_tokens": accepted_tokens,
+            "new_tokens_lens": num_accepted_tokens,
+            "next_draft_tokens": next_draft_tokens,
+            "next_new_tokens": next_new_tokens,
+        }
+
+    @staticmethod
+    def _prepare_external_shared_target_kv_draft_inputs(
+        *,
+        accepted_tokens: torch.Tensor,
+        num_accepted_tokens: torch.Tensor,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        num_contexts: int,
+        batch_indices: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Select the accepted token, hidden row, and fixed draft position."""
+        sequence_starts = torch.cumsum(
+            sequence_lengths, dim=0, dtype=torch.long) - sequence_lengths
+        recurrent_indices = sequence_starts + sequence_lengths - 1
+        recurrent_indices[num_contexts:].copy_(
+            sequence_starts[num_contexts:] +
+            num_accepted_tokens[num_contexts:] - 1)
+        draft_input_ids = accepted_tokens[batch_indices,
+                                          num_accepted_tokens - 1]
+        recurrent_hidden_states = hidden_states[recurrent_indices]
+        draft_position_ids = (
+            _select_mtp_position_ids(position_ids, recurrent_indices) + 1)
+        return draft_input_ids, recurrent_hidden_states, draft_position_ids
+
     def _forward_draft_loop(self, inputs, attn_metadata, spec_metadata,
                             draft_model, draft_kv_cache_manager, num_contexts,
                             num_gens, batch_size, num_accepted_tokens,
@@ -1168,7 +1314,9 @@ def sample_and_accept_draft_tokens(
         acceptance is enabled (both Eagle3 and MTP Eagle); ignored otherwise.
         """
         batch_size = attn_metadata.num_seqs
-        num_contexts = attn_metadata.num_contexts
+        num_contexts = (batch_size - spec_metadata.num_generations
+                        if self._uses_external_shared_target_kv else
+                        attn_metadata.num_contexts)
         num_gens = batch_size - num_contexts
 
         runtime_draft_len = spec_metadata.runtime_draft_len
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 721b3942ed04..683414763511 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum, auto
-from typing import TYPE_CHECKING, List, Optional, Type
+from typing import TYPE_CHECKING, Any, List, Optional, Type
 
 import torch
 from packaging.version import Version
@@ -107,6 +107,78 @@ def rejection_sampling_one_model(
 _FORCE_ACCEPT_RNG_SLOT_STRIDE = 1009
 
 
+@dataclass(frozen=True)
+class DraftModelCapabilities:
+    """Runtime ownership contract for an externally loaded draft model."""
+
+    loads_external_weights: bool = False
+    num_draft_modules: int = 0
+    owns_independent_kv_cache: bool = True
+    num_draft_kv_layers: int = 0
+    shares_target_kv_cache: bool = False
+    freezes_draft_attention_state: bool = False
+    requires_external_draft_metadata_view: bool = False
+
+    @classmethod
+    def external_shared_target_kv(cls) -> "DraftModelCapabilities":
+        return cls(
+            loads_external_weights=True,
+            num_draft_modules=1,
+            owns_independent_kv_cache=False,
+            num_draft_kv_layers=0,
+            shares_target_kv_cache=True,
+            freezes_draft_attention_state=True,
+            requires_external_draft_metadata_view=True,
+        )
+
+    @classmethod
+    def from_config(cls, config: Any) -> "DraftModelCapabilities":
+        return cls(
+            loads_external_weights=bool(
+                getattr(config, "loads_external_weights", False)),
+            num_draft_modules=int(getattr(config, "num_draft_modules", 0)),
+            owns_independent_kv_cache=bool(
+                getattr(config, "owns_independent_kv_cache", True)),
+            num_draft_kv_layers=int(getattr(config, "num_draft_kv_layers", 0)),
+            shares_target_kv_cache=bool(
+                getattr(config, "shares_target_kv_cache", False)),
+            freezes_draft_attention_state=bool(
+                getattr(config, "freezes_draft_attention_state", False)),
+            requires_external_draft_metadata_view=bool(
+                getattr(config, "requires_external_draft_metadata_view",
+                        False)),
+        )
+
+
+def get_draft_model_capabilities(
+        spec_config) -> Optional[DraftModelCapabilities]:
+    if spec_config is None:
+        return None
+    return getattr(spec_config, "_draft_model_capabilities", None)
+
+
+def needs_external_draft_weights(spec_config) -> bool:
+    """Whether a one-engine mode loads a separate draft checkpoint."""
+    if spec_config is None:
+        return False
+    capabilities = get_draft_model_capabilities(spec_config)
+    if capabilities is not None and capabilities.loads_external_weights:
+        return True
+    return spec_config.spec_dec_mode.need_load_draft_weights()
+
+
+def should_extend_context(spec_config,
+                          attention_backend: Type[AttentionBackend]) -> bool:
+    """Whether generation verification uses the backend's context kernel."""
+    capabilities = get_draft_model_capabilities(spec_config)
+    if (capabilities is not None and capabilities.shares_target_kv_cache
+            and capabilities.requires_external_draft_metadata_view):
+        from ..attention_backend.flashinfer import FlashInferAttention
+        if issubclass(attention_backend, FlashInferAttention):
+            return True
+    return spec_config.spec_dec_mode.extend_ctx(attention_backend)
+
+
 def should_use_separate_draft_kv_cache(spec_config) -> bool:
     """
     Check if separate draft KV cache should be used for one-engine speculative decoding.
@@ -115,6 +187,9 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
         return False
     if not spec_config.spec_dec_mode.use_one_engine():
         return False
+    capabilities = get_draft_model_capabilities(spec_config)
+    if capabilities is not None and not capabilities.owns_independent_kv_cache:
+        return False
     # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft
     # model does not read the paged draft KV cache managed by attention metadata.
     if spec_config.spec_dec_mode.is_dspark():
diff --git a/tensorrt_llm/_torch/speculative/model_drafter.py b/tensorrt_llm/_torch/speculative/model_drafter.py
index bffdbafd5833..5eae9b7e44cc 100644
--- a/tensorrt_llm/_torch/speculative/model_drafter.py
+++ b/tensorrt_llm/_torch/speculative/model_drafter.py
@@ -1,6 +1,3 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
 from __future__ import annotations
 
 import traceback
@@ -21,7 +18,6 @@
 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:
@@ -97,9 +93,6 @@ def __init__(
         self.guided_decoder = guided_decoder
 
         self.use_static_draft_loop = draft_model_engine.model_is_wrapped
-        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,28 +163,6 @@ def _create_generation_request(self, request: LlmRequest,
         new_request.state = LlmRequestState.GENERATION_IN_PROGRESS
         return new_request
 
-    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."""
-        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(
-                "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)
-            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:
@@ -252,14 +223,6 @@ 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.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,
                                               request,
                                               self.disable_overlap_scheduler)
@@ -267,7 +230,9 @@ 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()
 
-        if is_first_draft:
+        # 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:
             # 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.
@@ -336,10 +301,6 @@ def _prepare_draft_batch(
             for request in scheduled_requests.context_requests:
                 if request.py_disable_speculative_decoding:
                     continue
-                if self.shares_target_kv_cache:
-                    # 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 79aafa0859d4..4a2c8e7e6d51 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -146,8 +146,6 @@ def get_spec_metadata(spec_config,
             eagle3_resource_manager=spec_resource_manager,
             layers_to_capture=None,
             is_mtp_eagle=True,
-            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)
@@ -447,6 +445,10 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
+    """Return the logical number of draft modules executed by the worker."""
+    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
+    if capabilities is not None and capabilities.num_draft_modules:
+        return capabilities.num_draft_modules
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
     if spec_config.spec_dec_mode.is_mtp_vanilla():
@@ -457,6 +459,14 @@ def get_num_spec_layers(spec_config):
     return 0
 
 
+def get_num_draft_kv_layers(spec_config):
+    """Return the number of draft-owned layers requiring KV cache storage."""
+    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
+    if capabilities is not None:
+        return capabilities.num_draft_kv_layers
+    return get_num_spec_layers(spec_config)
+
+
 def update_spec_config_from_draft_model_config(spec_config,
                                                draft_pretrained_config) -> None:
     """Populate Eagle draft-layer fields from the loaded draft model config."""
@@ -523,6 +533,9 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
+    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
+    if capabilities is not None and capabilities.shares_target_kv_cache:
+        return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
     return 0
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index ab5f561070c6..0bea7b5439bd 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -2550,6 +2550,9 @@ class MTPDecodingConfig(DecodingBaseConfig):
 
     # Internal max batch size for dynamic-tree worker buffers.
     _max_batch_size: Optional[int] = PrivateAttr(default=None)
+    # Runtime-only ownership contract populated after loading an external
+    # assistant config. It is intentionally excluded from serialization.
+    _draft_model_capabilities: Any = PrivateAttr(default=None)
 
     sa_config: Optional[SAEnhancerConfig] = Field(
         default=None,
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 610676911cc0..109460a5a57b 100644
--- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
+++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py
@@ -14,13 +14,11 @@
 # 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
@@ -67,36 +65,6 @@ def _make_creator(
 
 
 class TestSplitGpuBudgetForDraft:
-    def test_shared_target_cache_skips_draft_manager(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(return_value="target")
-
-        resources = {}
-        c.build_managers(resources, estimating_kv_cache=False)
-
-        c._split_kv_cache_budget_for_draft.assert_not_called()
-        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 resources[ResourceManagerType.KV_CACHE_MANAGER] == "target"
-        assert resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] is None
-
     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 b49ec0b4ba0d..884fb8e2eebb 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -21,7 +21,7 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.model_engine import (
     PyTorchModelEngine, _build_request_multimodal_input,
-    _filter_cuda_graph_batch_sizes, _make_single_token_context_graph_batch)
+    _make_single_token_context_graph_batch)
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           SeqLenAwareSparseAttentionConfig,
                                           TorchLlmArgs)
@@ -36,7 +36,6 @@
 
 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._torch.speculative.spec_sampler_base import \
     SampleStateTensorsSpec
 from tensorrt_llm.bindings.executor import KvCacheConfig
@@ -1116,56 +1115,6 @@ def test_promoted_context_precedes_speculative_overlap_generation(
             [generation.py_seq_slot], 0)
         kv_cache_manager.shutdown()
 
-    def test_external_draft_len_graph_key_ignores_first_draft_state(
-            self) -> None:
-        runner = object.__new__(CUDAGraphRunner)
-        runner.config = SimpleNamespace(
-            is_draft_model=True,
-            draft_model_external_draft_len=0,
-            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_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/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
index c51b328951ff..089e575148b8 100644
--- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
+++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
@@ -774,6 +774,22 @@ def test_encoder_cache_reuses_image_embedding_across_requests(self):
         torch.testing.assert_close(second, first)
         self.assertEqual(len(model._multimodal_encoder_cache), 1)
 
+    def test_speculative_runtime_contract_is_proxied_to_language_model(self):
+        model = self._make_model()
+
+        self.assertIs(model.model, model.llm.model)
+        self.assertIs(model.lm_head, model.llm.lm_head)
+        self.assertIs(model.epilogue, model.llm.epilogue)
+        self.assertIs(model.spec_worker, model.llm.spec_worker)
+        self.assertIs(model.draft_config, model.llm.draft_config)
+        self.assertIs(model.draft_model, model.llm.draft_model)
+
+        weights = {"draft": torch.ones(1)}
+        mapper = object()
+        with unittest.mock.patch.object(model.llm, "load_draft_weights") as loader:
+            model.load_draft_weights(weights, mapper)
+        loader.assert_called_once_with(weights, mapper)
+
     def test_chunked_prefill_reuses_cached_vision_embeddings(self):
         """Later active chunks slice cached features without rerunning vision."""
         model = self._make_model()
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index c52258c08cae..72fd330ff141 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -588,7 +588,7 @@ def test_assistant_config_auto_config_round_trip(self):
             config.save_pretrained(directory)
             restored = AutoConfig.from_pretrained(directory)
 
-        self.assertIsInstance(restored, Gemma4AssistantConfig)
+        self.assertEqual(restored.model_type, "gemma4_assistant")
         self.assertEqual(restored.backbone_hidden_size, 256)
         self.assertEqual(restored.text_config.num_kv_shared_layers, 4)
 
@@ -2621,6 +2621,56 @@ def _expected_decode_block_table(
             source_offset += page_count
         return expected
 
+    @torch.no_grad()
+    @unittest.mock.patch(
+        "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
+    )
+    def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) -> None:
+        """The assistant gets private decode state over immutable target KV."""
+        page_counts = [3, 2]
+        kv_cache_manager, layers, metadata, queries, _, _ = self._make_trtllm_gen_decode_case(
+            page_counts
+        )
+        target_kv = {
+            layer.layer_idx: kv_cache_manager.get_buffers(layer.layer_idx).clone()
+            for layer in layers
+        }
+
+        accepted_tokens = torch.tensor([1, 3], dtype=torch.int, device="cuda")
+        draft_metadata = metadata.get_shared_kv_draft_metadata()
+        draft_metadata.update_shared_kv_draft_lengths(
+            metadata,
+            accepted_tokens,
+            num_contexts=0,
+        )
+        expected_kv_lens = metadata._cached_token_lens[:2] + accepted_tokens
+
+        for layer, query in zip(layers, queries, strict=True):
+            layer.forward(query, None, None, draft_metadata)
+
+        self.assertIs(draft_metadata.kv_cache_manager, metadata.kv_cache_manager)
+        self.assertIsNot(draft_metadata, metadata)
+        torch.testing.assert_close(
+            draft_metadata._shared_kv_runtime_lens[:2],
+            expected_kv_lens,
+            atol=0,
+            rtol=0,
+        )
+        for wrappers in draft_metadata._plan_params_to_wrappers.values():
+            torch.testing.assert_close(
+                wrappers.decode_wrapper._kv_lens_buffer[:2],
+                expected_kv_lens,
+                atol=0,
+                rtol=0,
+            )
+        for layer in layers:
+            torch.testing.assert_close(
+                kv_cache_manager.get_buffers(layer.layer_idx),
+                target_kv[layer.layer_idx],
+                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 b6eb66fcb6b1..a32d26c60a40 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
@@ -3,127 +3,170 @@
 
 from types import SimpleNamespace
 
+import pytest
 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",
-            shares_target_kv_cache=True,
-            freezes_draft_attention_state=True,
-        )
-        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
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+from tensorrt_llm._torch.speculative.interface import (
+    DraftModelCapabilities,
+    needs_external_draft_weights,
+    should_use_separate_draft_kv_cache,
+)
+from tensorrt_llm._torch.speculative.utils import (
+    get_num_draft_kv_layers,
+    get_num_extra_kv_tokens,
+    get_num_spec_layers,
+)
+from tensorrt_llm.llmapi import MTPDecodingConfig
+
+
+def _shared_kv_capabilities() -> DraftModelCapabilities:
+    return DraftModelCapabilities.external_shared_target_kv()
+
+
+def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
+    spec_config = MTPDecodingConfig(
+        max_draft_len=kwargs.pop("max_draft_len", 3),
+        speculative_model="/tmp/gemma4-assistant",
+        mtp_eagle_one_model=True,
+        **kwargs,
+    )
+    spec_config._draft_model_capabilities = _shared_kv_capabilities()
+    return spec_config
 
 
-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)
+def test_external_shared_kv_capability_separates_module_and_kv_counts():
+    spec_config = _shared_kv_spec_config()
 
-    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 needs_external_draft_weights(spec_config)
+    assert get_num_spec_layers(spec_config) == 1
+    assert get_num_draft_kv_layers(spec_config) == 0
+    assert get_num_extra_kv_tokens(spec_config) == 0
+    assert not should_use_separate_draft_kv_cache(spec_config)
+
+
+def test_embedded_one_model_mtp_does_not_load_external_weights():
+    spec_config = _shared_kv_spec_config()
+    spec_config._draft_model_capabilities = None
 
-    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]
+    assert not needs_external_draft_weights(spec_config)
 
 
-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),
+def test_external_shared_kv_worker_rejects_unverified_modes():
+    spec_config = _shared_kv_spec_config(
+        use_dynamic_tree=True,
+        dynamic_tree_max_topK=2,
     )
-    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,
+    worker = MTPEagleWorker(spec_config)
+
+    with pytest.raises(ValueError, match="linear draft path"):
+        worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
+
+
+@pytest.mark.parametrize(
+    "num_accepted_tokens, expected_hidden_rows",
+    [
+        ([1, 1, 1], [1, 2, 6]),
+        ([1, 2, 3], [1, 3, 8]),
+        ([1, 4, 4], [1, 5, 9]),
+    ],
+)
+def test_external_shared_kv_selects_last_accepted_target_state(
+    num_accepted_tokens,
+    expected_hidden_rows,
+):
+    accepted_tokens = torch.tensor(
+        [
+            [10, 11, 12, 13],
+            [20, 21, 22, 23],
+            [30, 31, 32, 33],
+        ],
+        dtype=torch.int32,
+    )
+    accepted_counts = torch.tensor(num_accepted_tokens, dtype=torch.long)
+    hidden_states = torch.arange(20, dtype=torch.float32).unsqueeze(1)
+    position_ids = torch.arange(10, dtype=torch.int32).unsqueeze(0)
+
+    draft_ids, recurrent_hidden, draft_positions = (
+        MTPEagleWorker._prepare_external_shared_target_kv_draft_inputs(
+            accepted_tokens=accepted_tokens,
+            num_accepted_tokens=accepted_counts,
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            sequence_lengths=torch.tensor([2, 4, 4]),
+            num_contexts=1,
+            batch_indices=torch.arange(3),
+        )
     )
 
-    assert (
-        drafter._create_shared_target_kv_request(request, [1, 2], is_first_draft=True)
-        is draft_request
+    expected_tokens = accepted_tokens[
+        torch.arange(3),
+        accepted_counts - 1,
+    ]
+    assert torch.equal(draft_ids, expected_tokens)
+    assert torch.equal(
+        recurrent_hidden.squeeze(1),
+        torch.tensor(expected_hidden_rows, dtype=torch.float32),
+    )
+    assert torch.equal(
+        draft_positions,
+        torch.tensor(expected_hidden_rows, dtype=torch.int32).unsqueeze(0) + 1,
     )
-    assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 9
 
-    drafter._create_shared_target_kv_request(request, [1, 2], is_first_draft=False)
-    assert drafter.spec_resource_manager.draft_hidden_state_offsets[17] == 3
 
+def test_gemma4_target_forward_dispatches_one_model_worker():
+    hidden_states = torch.tensor(
+        [
+            [1.0, 2.0],
+            [3.0, 4.0],
+            [5.0, 6.0],
+        ]
+    )
+    worker_calls = []
+
+    def spec_worker(**kwargs):
+        worker_calls.append(kwargs)
+        return {"logits": kwargs["logits"], "new_tokens": torch.tensor([[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(freezes_draft_attention_state=True))
-    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
+    model = SimpleNamespace(
+        layer_idx=-1,
+        config=SimpleNamespace(final_logit_softcapping=None),
+        model=lambda **kwargs: hidden_states,
+        logits_processor=SimpleNamespace(forward=lambda selected, *args: selected),
+        lm_head=object(),
+        spec_worker=spec_worker,
+        draft_model=object(),
     )
-    request = SimpleNamespace(py_is_first_draft=True, py_draft_tokens=[1])
-    batch = SimpleNamespace(generation_requests=[request])
+    spec_metadata = SimpleNamespace(
+        gather_ids=torch.tensor([2]),
+        is_layer_capture=lambda layer_idx: False,
+    )
+    attn_metadata = SimpleNamespace(padded_num_tokens=None)
 
-    engine._update_draft_inference_state_for_warmup(
-        batch, is_first_draft=True, resource_manager=resource_manager
+    outputs = Gemma4ForCausalLM.forward(
+        model,
+        attn_metadata=attn_metadata,
+        input_ids=torch.tensor([1, 2, 3]),
+        position_ids=torch.tensor([[0, 1, 2]]),
+        spec_metadata=spec_metadata,
     )
 
-    assert not spec_resource_manager.is_first_draft
-    assert not request.py_is_first_draft
-    assert request.py_draft_tokens == []
+    assert torch.equal(outputs["new_tokens"], torch.tensor([[7]]))
+    assert len(worker_calls) == 1
+    assert torch.equal(worker_calls[0]["hidden_states"], hidden_states)
+    assert torch.equal(worker_calls[0]["logits"], hidden_states[[2]])
+    assert worker_calls[0]["draft_model"] is model.draft_model
 
 
-def test_gemma4_target_forward_captures_speculative_hidden_states():
+def test_gemma4_target_forward_still_captures_hidden_states_without_worker():
     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(),
+        spec_worker=None,
     )
     captured = []
     spec_metadata = SimpleNamespace(

From e567d430d8b16417ed71b8ec2f3782964b00448e Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Sun, 26 Jul 2026 16:33:52 +0000
Subject: [PATCH 11/26] [None][refactor] minimize Gemma4 MTP integration

Localize Gemma4 shared-KV drafting to the MTP Eagle worker.

Remove generic capability and KV-pool plumbing, prune redundant tests, and preserve the FlashInfer CUDA graph refresh validated by A/B coverage.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 examples/llm-api/quickstart_advanced.py       |   3 -
 examples/models/core/gemma/README.md          |   7 +-
 .../_torch/attention_backend/flashinfer.py    |  23 +-
 tensorrt_llm/_torch/configs/gemma4.py         |  14 -
 tensorrt_llm/_torch/models/modeling_gemma4.py |  60 +--
 .../_torch/models/modeling_gemma4mm.py        |   3 +-
 .../_torch/models/modeling_speculative.py     |  35 +-
 tensorrt_llm/_torch/models/modeling_utils.py  |   4 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   8 +-
 .../_torch/pyexecutor/model_loader.py         |  16 +-
 .../_torch/pyexecutor/py_executor_creator.py  |  14 +-
 .../_torch/pyexecutor/resource_manager.py     |   4 +-
 tensorrt_llm/_torch/speculative/__init__.py   |  15 +-
 tensorrt_llm/_torch/speculative/eagle3.py     | 361 +++++++++---------
 tensorrt_llm/_torch/speculative/interface.py  |  63 +--
 tensorrt_llm/_torch/speculative/utils.py      |  31 +-
 tensorrt_llm/llmapi/llm_args.py               |   5 +-
 .../_torch/modeling/test_modeling_gemma4.py   |  68 +---
 .../hw_agnostic/test_gemma4_drafting_loop.py  |  57 +--
 19 files changed, 284 insertions(+), 507 deletions(-)

diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py
index 07d04206bde5..3757d49486d6 100644
--- a/examples/llm-api/quickstart_advanced.py
+++ b/examples/llm-api/quickstart_advanced.py
@@ -308,9 +308,6 @@ 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")
         speculative_model = (args.draft_model_dir if args.draft_model_dir
                              is not None else args.model_dir)
         spec_config = MTPDecodingConfig(
diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md
index da4b166b76cb..b32f75ac076b 100644
--- a/examples/models/core/gemma/README.md
+++ b/examples/models/core/gemma/README.md
@@ -50,10 +50,9 @@ The `/v1/chat/completions` endpoint applies the Gemma 4 chat template automatica
 ### MTP speculative decoding
 
 Gemma 4 supports Multi-Token Prediction (MTP) speculative decoding through the
-one-model PyTorch execution path. The target and its matching assistant
-checkpoint are loaded into one engine, and the Q-only assistant reads the
-target model's KV cache. Create a server configuration for the
-target/assistant pair:
+PyTorch execution path. The target loads its matching assistant checkpoint,
+and the Q-only assistant reads the target model's KV cache. Create a server
+configuration for the target/assistant pair:
 
 ```bash
 cat > gemma4_mtp.yaml <<'EOF'
diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 8c222b98bce7..051c1165213b 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -166,7 +166,6 @@ 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.
@@ -849,6 +848,15 @@ 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():
+                        head_dim = mgr.head_dim_per_layer[
+                            mgr.layer_offsets[layer_idx]]
+                        if head_dim not in self._vswa_head_dim_to_pool:
+                            self._vswa_head_dim_to_pool[head_dim] = 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.
@@ -1138,8 +1146,10 @@ def _build_decode_block_tables(
         num_gens = self.num_generations
         if num_gens == 0:
             return None
+        pool_id = getattr(self, "_vswa_head_dim_to_pool",
+                          {}).get(plan_params.head_dim)
         host_paged_kv_indices = self._host_pool_indices.get(
-            plan_params.kv_pool_id, self._host_paged_kv_indices)
+            pool_id, self._host_paged_kv_indices)
         if host_paged_kv_indices is None:
             return None
         gen_num_blocks = np.asarray(self.num_blocks[self.num_contexts:],
@@ -1471,6 +1481,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
         # table 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
@@ -1479,7 +1490,8 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                         or prefill_wrapper._backend != "trtllm-gen"):
                     continue
                 block_tables = prefill_wrapper._block_tables
-                pool_id = plan_params.kv_pool_id
+                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]
@@ -1515,6 +1527,7 @@ 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
@@ -1522,7 +1535,8 @@ 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 = plan_params.kv_pool_id
+                pool_id = (head_dim_to_pool.get(plan_params.head_dim)
+                           if head_dim_to_pool else None)
                 if pool_id is None:
                     continue
                 batch_size, table_width = block_tables.shape
@@ -1625,7 +1639,6 @@ 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 7609ae3c7fa6..bb455e2d3f83 100644
--- a/tensorrt_llm/_torch/configs/gemma4.py
+++ b/tensorrt_llm/_torch/configs/gemma4.py
@@ -27,15 +27,6 @@ class Gemma4AssistantConfig(PreTrainedConfig):
     model_type = "gemma4_assistant"
     sub_configs = {"text_config": Gemma4TextConfig}
 
-    # Runtime ownership contract for the one-model speculative pipeline.
-    loads_external_weights = True
-    num_draft_modules = 1
-    owns_independent_kv_cache = False
-    num_draft_kv_layers = 0
-    shares_target_kv_cache = True
-    freezes_draft_attention_state = True
-    requires_external_draft_metadata_view = True
-
     def __init__(
         self,
         text_config=None,
@@ -102,8 +93,3 @@ 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 4d46246dbb32..054225a0d28e 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1227,6 +1227,17 @@ def forward(
 # ---------------------------------------------------------------------------
 # Gemma4 For Causal LM
 # ---------------------------------------------------------------------------
+def _configure_gemma4_mtp_assistant(model_config: ModelConfig) -> None:
+    spec_config = model_config.spec_config
+    if (
+        spec_config is not None
+        and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+        and spec_config.speculative_model is not None
+    ):
+        spec_config._is_gemma4_mtp_assistant = True
+        spec_config._allow_separate_draft_kv_cache = False
+
+
 @register_auto_model("Gemma4ForCausalLM")
 class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
     def __init__(
@@ -1248,6 +1259,7 @@ def __init__(
                 "moe_ep_size>1 requires a Gemma4 MoE variant (only 26B-A4B-it today)."
             )
 
+        _configure_gemma4_mtp_assistant(model_config)
         super().__init__(Gemma4TextModel(model_config), model_config)
 
     @classmethod
@@ -1584,10 +1596,6 @@ def __init__(self, model_config: ModelConfig):
         )
         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)
@@ -1627,21 +1635,6 @@ def _constant_position_ids(
             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_draft_step(
         self,
         input_ids: torch.IntTensor,
@@ -1668,35 +1661,6 @@ def forward_draft_step(
             logits = self.lm_head(assistant_hidden_states).float()
         return logits, projected_hidden_states
 
-    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")
-        logits, projected_hidden_states = self.forward_draft_step(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            recurrent_hidden_states=spec_metadata.get_hidden_states(),
-            attn_metadata=attn_metadata,
-            spec_metadata=spec_metadata,
-        )
-        spec_metadata.maybe_capture_hidden_states(
-            self.config.num_hidden_layers - 1,
-            projected_hidden_states,
-        )
-
-        if return_context_logits:
-            return logits
-        last_token_indices = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
-        return logits[last_token_indices]
-
     def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         weights = weight_mapper.preprocess_weights(weights)
         ordering_weight_name = "masked_embedding.token_ordering"
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index a344a800d984..30ca39eadfa9 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -49,7 +49,7 @@
 from ...sampling_params import SamplingParams
 from ..modules.embedding import Embedding
 from ..modules.linear import Linear
-from .modeling_gemma4 import Gemma4ForCausalLM
+from .modeling_gemma4 import Gemma4ForCausalLM, _configure_gemma4_mtp_assistant
 from .modeling_gemma4_audio import Gemma4AudioModel
 from .modeling_gemma4_vision import Gemma4VisionModel
 from .modeling_multimodal_mixin import MultimodalModelMixin, PreparedLlmInputs
@@ -869,6 +869,7 @@ def __init__(self, model_config: ModelConfig[Gemma4Config]):
             )
 
         config = model_config.pretrained_config
+        _configure_gemma4_mtp_assistant(model_config)
         super().__init__(config)
 
         # Pin multimodal tensors to the local rank so each rank of a multi-GPU
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index bca7d4e1cf53..1112cdc19ef9 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -30,8 +30,9 @@
 except ImportError:
     _flashinfer_rope = None
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
-from ..speculative import (DraftModelCapabilities, SpecMetadata,
-                           get_spec_worker, should_use_separate_draft_kv_cache)
+from ..speculative import (SpecMetadata, get_spec_worker,
+                           should_use_separate_draft_kv_cache)
+from ..speculative.interface import is_gemma4_mtp_assistant
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_auto import AutoModelForCausalLM
@@ -1890,9 +1891,7 @@ def get_draft_model(model_config, draft_config, lm_head, model):
             )
 
     elif (spec_dec_mode.is_mtp_eagle_one_model() and draft_config is not None
-          and getattr(model_config.spec_config, "_draft_model_capabilities",
-                      None) is not None and model_config.spec_config.
-          _draft_model_capabilities.loads_external_weights):
+          and is_gemma4_mtp_assistant(model_config.spec_config)):
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
         return MTPForCausalLM(model_config,
@@ -2017,30 +2016,10 @@ def __init__(self,
                         self.draft_config.pretrained_config)
                     draft_architectures = getattr(draft_pretrained_config,
                                                   "architectures", None) or []
-                    if "Gemma4AssistantForCausalLM" in draft_architectures:
-                        # Newer Transformers releases provide a native
-                        # Gemma4AssistantConfig without TRT-LLM runtime
-                        # ownership attributes. Its architecture has the same
-                        # all-Q-only, shared-target-KV contract.
-                        capabilities = (
-                            DraftModelCapabilities.external_shared_target_kv())
-                    else:
-                        capabilities = DraftModelCapabilities.from_config(
-                            draft_pretrained_config)
-                    if not (capabilities.loads_external_weights
-                            and capabilities.shares_target_kv_cache
-                            and not capabilities.owns_independent_kv_cache
-                            and capabilities.num_draft_modules == 1
-                            and capabilities.num_draft_kv_layers == 0
-                            and capabilities.freezes_draft_attention_state and
-                            capabilities.requires_external_draft_metadata_view):
+                    if "Gemma4AssistantForCausalLM" not in draft_architectures:
                         raise ValueError(
-                            "External one-model MTP assistants must load "
-                            "external weights, expose one logical draft module, "
-                            "own zero KV layers, share the target KV cache, "
-                            "freeze draft attention state, and require an "
-                            "external draft metadata view.")
-                    spec_config._draft_model_capabilities = capabilities
+                            "Gemma4 MTP requires a "
+                            "Gemma4AssistantForCausalLM checkpoint.")
 
                 elif spec_config.spec_dec_mode.is_external_drafter():
                     self.draft_config = ModelConfig.from_pretrained(
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 51be81a4cf4c..eb91a6456209 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -312,9 +312,9 @@ def __pp_init__(self):
         total_num_layers = num_hidden_layers
         spec_config = getattr(self.model_config, "spec_config", None)
         if spec_config is not None:
-            from ..speculative.utils import get_num_draft_kv_layers
+            from ..speculative.utils import get_num_spec_layers
 
-            num_spec_layers = get_num_draft_kv_layers(spec_config) or 0
+            num_spec_layers = get_num_spec_layers(spec_config) or 0
             total_num_layers += num_spec_layers
             if num_spec_layers > 0 and mapping.is_last_pp_rank():
                 pp_layer_list.extend(
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index d5bc614a0862..9c0023cac331 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -44,7 +44,7 @@
 from ..hostfunc import set_low_latency_dispatch
 from ..model_config import ModelConfig
 from ..models.modeling_multimodal_mixin import MultimodalModelMixin
-from ..speculative import (get_num_draft_kv_layers, get_num_extra_kv_tokens,
+from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers,
                            get_spec_decoder, should_use_separate_draft_kv_cache)
 from ..utils import is_gdn_replay_enabled
 from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params,
@@ -1207,7 +1207,7 @@ def _get_num_draft_layers(self) -> int:
         """
         if self._speculative_config.spec_dec_mode.is_external_drafter():
             return self._draft_config.pretrained_config.num_hidden_layers
-        return get_num_draft_kv_layers(self._speculative_config)
+        return get_num_spec_layers(self._speculative_config)
 
     def _create_one_model_draft_kv_cache_manager(
         self,
@@ -1797,7 +1797,7 @@ def _build_per_layer_num_kv_heads(
     if spec_config is None or draft_config is None:
         return num_key_value_heads
 
-    from ..speculative.utils import get_num_draft_kv_layers
+    from ..speculative.utils import get_num_spec_layers
     draft_pretrained = draft_config.pretrained_config
     draft_num_kv_heads = getattr(
         draft_pretrained, 'num_key_value_heads',
@@ -1806,7 +1806,7 @@ def _build_per_layer_num_kv_heads(
     if draft_num_kv_heads is None or draft_num_kv_heads == num_key_value_heads:
         return num_key_value_heads
 
-    num_spec_layers = get_num_draft_kv_layers(spec_config)
+    num_spec_layers = get_num_spec_layers(spec_config)
     logger.info(f"Per-layer KV heads for speculative decoding: "
                 f"target={num_key_value_heads} x {num_hidden_layers} layers, "
                 f"draft={draft_num_kv_heads} x {num_spec_layers} layers, "
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 238053adcaf9..8d3367b9cbed 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -417,13 +417,6 @@ def load_config_and_apply_defaults(
 
         config = checkpoint_loader.load_config(checkpoint_dir, **config_kwargs)
 
-        if llm_args.speculative_config is not None:
-            from tensorrt_llm._torch.speculative import \
-                update_spec_config_from_model_config
-
-            update_spec_config_from_model_config(llm_args.speculative_config,
-                                                 config.pretrained_config)
-
         model_cls = AutoModelForCausalLM._resolve_class(config)
         use_kv_cache_manager_v2 = (
             llm_args.kv_cache_config.use_kv_cache_manager_v2)
@@ -440,6 +433,15 @@ def load_config_and_apply_defaults(
                         f"Applied model defaults for {model_cls.__name__}: {applied_defaults}"
                     )
 
+        if llm_args.speculative_config is not None:
+            from tensorrt_llm._torch.speculative import \
+                update_spec_config_from_model_config
+
+            # Model defaults reconstruct nested Pydantic configs. Populate
+            # runtime-only speculative state after that reconstruction.
+            update_spec_config_from_model_config(llm_args.speculative_config,
+                                                 config.pretrained_config)
+
         # The transceiver preference follows the checkpoint's original
         # architecture: _resolve_class may rewrite it to an execution class
         # (e.g. MTPDraftModelForCausalLM), which must not drop the target
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 35733e8bc03b..ce3c6962ceba 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -40,8 +40,7 @@
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
                     validate_feature_combination)
-from .config_utils import (is_hybrid_linear, is_minimax_m3,
-                           load_pretrained_config)
+from .config_utils import is_hybrid_linear, is_minimax_m3
 from .connectors.kv_cache_connector import KvCacheConnectorManager
 from .dwdp import DwdpManager
 from .guided_decoder import CapturableGuidedDecoder, GuidedDecoder
@@ -455,14 +454,9 @@ def create_py_executor(
         # External MTP assistants with shared target KV use a dedicated
         # FlashInfer decode metadata view. Other one-engine modes still rely
         # on the one-query-per-sequence decode contract.
-        supports_shared_kv_flashinfer = False
-        if (spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-                and spec_config.speculative_model is not None
-                and checkpoint_dir is not None):
-            target_config = load_pretrained_config(
-                checkpoint_dir, trust_remote_code=llm_args.trust_remote_code)
-            supports_shared_kv_flashinfer = target_config.model_type in (
-                "gemma4", "gemma4_text")
+        supports_shared_kv_flashinfer = getattr(spec_config,
+                                                "_is_gemma4_mtp_assistant",
+                                                False)
         if (llm_args.attn_backend == "FLASHINFER"
                 and not supports_shared_kv_flashinfer):
             raise ValueError(
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 4a4abe503583..b4faf969d0da 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -168,7 +168,7 @@ def get_pp_layers(
     spec_config: Optional["DecodingBaseConfig"] = None,
     layer_mask: Optional[List[bool]] = None,
 ) -> Tuple[List[int], int]:
-    from ..speculative.utils import get_num_draft_kv_layers
+    from ..speculative.utils import get_num_spec_layers
 
     total_num_layers = num_layers
     if layer_mask is not None:
@@ -193,7 +193,7 @@ def get_pp_layers(
     # When layer_mask is provided, the caller explicitly controls which layers
     # to include, so we should not add extra layers automatically.
     if spec_config is not None and layer_mask is None:
-        num_spec_layers = get_num_draft_kv_layers(spec_config)
+        num_spec_layers = get_num_spec_layers(spec_config)
         total_num_layers += num_spec_layers
         if mapping.is_last_pp_rank():
             pp_layers.extend(
diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py
index 07f2127c4e00..96185991175c 100644
--- a/tensorrt_llm/_torch/speculative/__init__.py
+++ b/tensorrt_llm/_torch/speculative/__init__.py
@@ -3,8 +3,7 @@
 from .draft_target import (DraftTargetOneModelSpecMetadata,
                            DraftTargetOneModelWorker)
 from .eagle3 import Eagle3SpecMetadata, MTPEagleWorker
-from .interface import (DraftModelCapabilities, SpecMetadata, SpecWorkerBase,
-                        get_draft_model_capabilities,
+from .interface import (SpecMetadata, SpecWorkerBase,
                         needs_external_draft_weights,
                         prepare_attn_metadata_for_draft_replay,
                         restore_attn_metadata_after_draft_replay,
@@ -21,18 +20,16 @@
                                 SpecSamplerBase)
 from .spec_tree_manager import SpecTreeManager
 from .suffix_automaton import SuffixAutomatonManager
-from .utils import (get_draft_kv_cache_manager, get_num_draft_kv_layers,
-                    get_num_extra_kv_tokens, get_num_spec_layers,
-                    get_spec_decoder, get_spec_drafter, get_spec_metadata,
-                    get_spec_resource_manager, get_spec_worker,
-                    update_spec_config_from_draft_model_config,
+from .utils import (get_draft_kv_cache_manager, get_num_extra_kv_tokens,
+                    get_num_spec_layers, get_spec_decoder, get_spec_drafter,
+                    get_spec_metadata, get_spec_resource_manager,
+                    get_spec_worker, update_spec_config_from_draft_model_config,
                     update_spec_config_from_loaded_model,
                     update_spec_config_from_model_config)
 
 __all__ = [
     "DFlashSpecMetadata",
     "DFlashWorker",
-    "DraftModelCapabilities",
     "DraftTargetOneModelSpecMetadata",
     "DraftTargetOneModelWorker",
     "Eagle3SpecMetadata",
@@ -57,8 +54,6 @@
     "SpecSamplerBase",
     "SpecWorkerBase",
     "get_draft_kv_cache_manager",
-    "get_draft_model_capabilities",
-    "get_num_draft_kv_layers",
     "get_num_extra_kv_tokens",
     "get_num_spec_layers",
     "get_spec_decoder",
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 9ac7afe65836..1a3b7c346669 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -20,7 +20,7 @@
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata, SpecWorkerBase
+from .interface import SpecMetadata, SpecWorkerBase, is_gemma4_mtp_assistant
 from .mtp import MTPSampler, _select_mtp_position_ids
 from .sa_enhancer import SADraftEnhancer
 from .spec_tree_manager import SpecTreeManager
@@ -657,41 +657,6 @@ def __init__(self,
         self._saved_position_offsets = None
         self._saved_position_offsets_cpp = None
         self._saved_generation_lengths = None
-        self._uses_external_shared_target_kv = False
-
-    def set_draft_model(self, draft_model) -> None:
-        super().set_draft_model(draft_model)
-        capabilities = getattr(self.spec_config, "_draft_model_capabilities",
-                               None)
-        self._uses_external_shared_target_kv = bool(
-            capabilities is not None and capabilities.loads_external_weights
-            and capabilities.shares_target_kv_cache
-            and not capabilities.owns_independent_kv_cache)
-        if not self._uses_external_shared_target_kv:
-            return
-        if self.use_dynamic_tree:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP supports only the "
-                "linear draft path.")
-        if self.spec_config.draft_len_schedule is not None:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP does not support a "
-                "draft length schedule.")
-        if self.sa_enhancer is not None:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP does not support the "
-                "suffix automaton enhancer.")
-        if self.spec_config.use_rejection_sampling:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP does not support "
-                "rejection sampling.")
-
-    def set_guided_decoder(self, guided_decoder) -> bool:
-        if self._uses_external_shared_target_kv:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP does not support "
-                "guided decoding.")
-        return super().set_guided_decoder(guided_decoder)
 
     @property
     def max_draft_len(self) -> int:
@@ -759,11 +724,6 @@ def _forward_impl(self,
                       resource_manager=None):
 
         runtime_draft_len = spec_metadata.runtime_draft_len
-        if self._uses_external_shared_target_kv:
-            return self._forward_external_shared_target_kv(
-                input_ids, position_ids, hidden_states, logits, attn_metadata,
-                spec_metadata, draft_model)
-
         # skip the draft forward if the runtime draft length is 0
         if runtime_draft_len == 0:
             return self.skip_drafting(input_ids, position_ids, hidden_states,
@@ -858,134 +818,6 @@ def _forward_impl(self,
             'next_new_tokens': next_new_tokens,
         }
 
-    def _forward_external_shared_target_kv(
-        self,
-        input_ids,
-        position_ids,
-        hidden_states,
-        logits,
-        attn_metadata,
-        spec_metadata,
-        draft_model,
-    ):
-        """Draft with an external Q-only assistant over accepted target KV."""
-        if not isinstance(attn_metadata, FlashInferAttentionMetadata):
-            raise TypeError(
-                "Gemma4 shared-target-KV one-model MTP currently requires "
-                "FlashInfer attention metadata.")
-        if self.guided_decoder is not None:
-            raise ValueError(
-                "Gemma4 shared-target-KV one-model MTP does not support "
-                "guided decoding.")
-
-        batch_size = attn_metadata.num_seqs
-        num_contexts = batch_size - spec_metadata.num_generations
-        runtime_draft_len = spec_metadata.runtime_draft_len
-        if runtime_draft_len == 0:
-            target_tokens = self._sample_tokens_for_batch(
-                logits, spec_metadata, num_contexts, batch_size)
-            accepted_tokens = target_tokens.unsqueeze(1)
-            num_accepted_tokens = torch.ones(batch_size,
-                                             dtype=torch.int,
-                                             device=logits.device)
-            next_draft_tokens = torch.empty((batch_size, 0),
-                                            dtype=torch.int32,
-                                            device=logits.device)
-            return {
-                "logits": logits,
-                "new_tokens": accepted_tokens,
-                "new_tokens_lens": num_accepted_tokens,
-                "next_draft_tokens": next_draft_tokens,
-                "next_new_tokens": accepted_tokens,
-            }
-
-        raw_logits = logits
-        accepted_tokens, num_accepted_tokens = (
-            self.sample_and_accept_draft_tokens(input_ids, logits,
-                                                attn_metadata, spec_metadata))
-
-        (
-            draft_input_ids,
-            recurrent_hidden_states,
-            draft_position_ids,
-        ) = self._prepare_external_shared_target_kv_draft_inputs(
-            accepted_tokens=accepted_tokens,
-            num_accepted_tokens=num_accepted_tokens,
-            hidden_states=hidden_states,
-            position_ids=position_ids,
-            sequence_lengths=attn_metadata.seq_lens_cuda[:batch_size],
-            num_contexts=num_contexts,
-            batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
-        )
-
-        draft_metadata = attn_metadata.get_shared_kv_draft_metadata()
-        draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
-                                                      num_accepted_tokens,
-                                                      num_contexts)
-        draft_metadata.use_spec_decoding = False
-        draft_metadata.padded_num_tokens = None
-        draft_metadata.all_rank_num_tokens = (
-            spec_metadata.subseq_all_rank_num_tokens)
-
-        next_draft_tokens = []
-        for draft_step in range(runtime_draft_len):
-            draft_logits, recurrent_hidden_states = (
-                draft_model.forward_draft_step(
-                    input_ids=draft_input_ids,
-                    position_ids=draft_position_ids,
-                    recurrent_hidden_states=recurrent_hidden_states,
-                    attn_metadata=draft_metadata,
-                    spec_metadata=spec_metadata,
-                ))
-            draft_input_ids = self.sample_draft_tokens(
-                draft_logits,
-                spec_metadata,
-                batch_size,
-                draft_step=draft_step,
-            )
-            next_draft_tokens.append(draft_input_ids)
-
-        next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
-        batch_indices = spec_metadata.batch_indices_cuda[:batch_size]
-        next_new_tokens = self._prepare_next_new_tokens(accepted_tokens,
-                                                        next_draft_tokens,
-                                                        batch_indices,
-                                                        batch_size,
-                                                        num_accepted_tokens)
-        attn_metadata.use_spec_decoding = True
-        return {
-            "logits": raw_logits,
-            "new_tokens": accepted_tokens,
-            "new_tokens_lens": num_accepted_tokens,
-            "next_draft_tokens": next_draft_tokens,
-            "next_new_tokens": next_new_tokens,
-        }
-
-    @staticmethod
-    def _prepare_external_shared_target_kv_draft_inputs(
-        *,
-        accepted_tokens: torch.Tensor,
-        num_accepted_tokens: torch.Tensor,
-        hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
-        sequence_lengths: torch.Tensor,
-        num_contexts: int,
-        batch_indices: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Select the accepted token, hidden row, and fixed draft position."""
-        sequence_starts = torch.cumsum(
-            sequence_lengths, dim=0, dtype=torch.long) - sequence_lengths
-        recurrent_indices = sequence_starts + sequence_lengths - 1
-        recurrent_indices[num_contexts:].copy_(
-            sequence_starts[num_contexts:] +
-            num_accepted_tokens[num_contexts:] - 1)
-        draft_input_ids = accepted_tokens[batch_indices,
-                                          num_accepted_tokens - 1]
-        recurrent_hidden_states = hidden_states[recurrent_indices]
-        draft_position_ids = (
-            _select_mtp_position_ids(position_ids, recurrent_indices) + 1)
-        return draft_input_ids, recurrent_hidden_states, draft_position_ids
-
     def _forward_draft_loop(self, inputs, attn_metadata, spec_metadata,
                             draft_model, draft_kv_cache_manager, num_contexts,
                             num_gens, batch_size, num_accepted_tokens,
@@ -1307,6 +1139,7 @@ def sample_and_accept_draft_tokens(
         logits: torch.Tensor,
         attn_metadata: AttentionMetadata,
         spec_metadata: Eagle3OneModelSpecMetadata,
+        num_contexts: Optional[int] = None,
     ):
         """Sample the golden token and verify previously proposed draft tokens.
 
@@ -1314,9 +1147,8 @@ def sample_and_accept_draft_tokens(
         acceptance is enabled (both Eagle3 and MTP Eagle); ignored otherwise.
         """
         batch_size = attn_metadata.num_seqs
-        num_contexts = (batch_size - spec_metadata.num_generations
-                        if self._uses_external_shared_target_kv else
-                        attn_metadata.num_contexts)
+        if num_contexts is None:
+            num_contexts = attn_metadata.num_contexts
         num_gens = batch_size - num_contexts
 
         runtime_draft_len = spec_metadata.runtime_draft_len
@@ -1444,13 +1276,7 @@ def prepare_1st_drafter_inputs(
 
 
 class MTPEagleWorker(Eagle3OneModelWorker):
-    """Backward-compatible alias for ``Eagle3OneModelWorker`` in MTP Eagle mode.
-
-    The constructor matches the historical positional signature
-    ``(spec_config, model_config, use_separate_draft_kv_cache)`` so callers
-    that import ``MTPEagleWorker`` from ``mtp.py`` or instantiate it directly
-    keep working. All logic is inherited from :class:`Eagle3OneModelWorker`.
-    """
+    """MTP worker built on the shared Eagle3 one-model drafting loop."""
 
     def __init__(self,
                  spec_config,
@@ -1465,3 +1291,180 @@ def __init__(self,
             use_separate_draft_kv_cache=use_separate_draft_kv_cache)
         # Preserved for callers/tests that still expect this attribute.
         self.is_thop = False
+        self._uses_external_shared_target_kv = is_gemma4_mtp_assistant(
+            spec_config)
+
+    def set_draft_model(self, draft_model) -> None:
+        super().set_draft_model(draft_model)
+        if not self._uses_external_shared_target_kv:
+            return
+        if self.use_dynamic_tree:
+            raise ValueError(
+                "Gemma4 shared-target-KV MTP supports only the linear draft "
+                "path.")
+        if self.spec_config.draft_len_schedule is not None:
+            raise ValueError(
+                "Gemma4 shared-target-KV MTP does not support a draft length "
+                "schedule.")
+        if self.sa_enhancer is not None:
+            raise ValueError(
+                "Gemma4 shared-target-KV MTP does not support the suffix "
+                "automaton enhancer.")
+        if self.spec_config.use_rejection_sampling:
+            raise ValueError(
+                "Gemma4 shared-target-KV MTP does not support rejection "
+                "sampling.")
+
+    def set_guided_decoder(self, guided_decoder) -> bool:
+        if self._uses_external_shared_target_kv:
+            raise ValueError(
+                "Gemma4 shared-target-KV MTP does not support guided "
+                "decoding.")
+        return super().set_guided_decoder(guided_decoder)
+
+    def _forward_impl(self,
+                      input_ids,
+                      position_ids,
+                      hidden_states,
+                      logits,
+                      attn_metadata,
+                      spec_metadata,
+                      draft_model,
+                      resource_manager=None):
+        if not self._uses_external_shared_target_kv:
+            return super()._forward_impl(input_ids, position_ids, hidden_states,
+                                         logits, attn_metadata, spec_metadata,
+                                         draft_model, resource_manager)
+        return self._forward_external_shared_target_kv(input_ids, position_ids,
+                                                       hidden_states, logits,
+                                                       attn_metadata,
+                                                       spec_metadata,
+                                                       draft_model)
+
+    def _forward_external_shared_target_kv(
+        self,
+        input_ids,
+        position_ids,
+        hidden_states,
+        logits,
+        attn_metadata,
+        spec_metadata,
+        draft_model,
+    ):
+        """Draft with a Gemma4 Q-only assistant over accepted target KV."""
+        if not isinstance(attn_metadata, FlashInferAttentionMetadata):
+            raise TypeError(
+                "Gemma4 shared-target-KV MTP requires FlashInfer attention "
+                "metadata.")
+
+        batch_size = attn_metadata.num_seqs
+        num_contexts = batch_size - spec_metadata.num_generations
+        runtime_draft_len = spec_metadata.runtime_draft_len
+        if runtime_draft_len == 0:
+            target_tokens = self._sample_tokens_for_batch(
+                logits, spec_metadata, num_contexts, batch_size)
+            accepted_tokens = target_tokens.unsqueeze(1)
+            num_accepted_tokens = torch.ones(batch_size,
+                                             dtype=torch.int,
+                                             device=logits.device)
+            next_draft_tokens = torch.empty((batch_size, 0),
+                                            dtype=torch.int32,
+                                            device=logits.device)
+            return {
+                "logits": logits,
+                "new_tokens": accepted_tokens,
+                "new_tokens_lens": num_accepted_tokens,
+                "next_draft_tokens": next_draft_tokens,
+                "next_new_tokens": accepted_tokens,
+            }
+
+        accepted_tokens, num_accepted_tokens = (
+            self.sample_and_accept_draft_tokens(
+                input_ids,
+                logits,
+                attn_metadata,
+                spec_metadata,
+                num_contexts=num_contexts,
+            ))
+        (
+            draft_input_ids,
+            recurrent_hidden_states,
+            draft_position_ids,
+        ) = self._prepare_external_shared_target_kv_draft_inputs(
+            accepted_tokens=accepted_tokens,
+            num_accepted_tokens=num_accepted_tokens,
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            sequence_lengths=attn_metadata.seq_lens_cuda[:batch_size],
+            num_contexts=num_contexts,
+            batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
+        )
+
+        draft_metadata = attn_metadata.get_shared_kv_draft_metadata()
+        draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
+                                                      num_accepted_tokens,
+                                                      num_contexts)
+        draft_metadata.use_spec_decoding = False
+        draft_metadata.padded_num_tokens = None
+        draft_metadata.all_rank_num_tokens = (
+            spec_metadata.subseq_all_rank_num_tokens)
+
+        next_draft_tokens = []
+        for draft_step in range(runtime_draft_len):
+            draft_logits, recurrent_hidden_states = (
+                draft_model.forward_draft_step(
+                    input_ids=draft_input_ids,
+                    position_ids=draft_position_ids,
+                    recurrent_hidden_states=recurrent_hidden_states,
+                    attn_metadata=draft_metadata,
+                    spec_metadata=spec_metadata,
+                ))
+            draft_input_ids = self.sample_draft_tokens(
+                draft_logits,
+                spec_metadata,
+                batch_size,
+                draft_step=draft_step,
+            )
+            next_draft_tokens.append(draft_input_ids)
+
+        next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
+        next_new_tokens = self._prepare_next_new_tokens(
+            accepted_tokens,
+            next_draft_tokens,
+            spec_metadata.batch_indices_cuda[:batch_size],
+            batch_size,
+            num_accepted_tokens,
+        )
+        attn_metadata.use_spec_decoding = True
+        return {
+            "logits": logits,
+            "new_tokens": accepted_tokens,
+            "new_tokens_lens": num_accepted_tokens,
+            "next_draft_tokens": next_draft_tokens,
+            "next_new_tokens": next_new_tokens,
+        }
+
+    @staticmethod
+    def _prepare_external_shared_target_kv_draft_inputs(
+        *,
+        accepted_tokens: torch.Tensor,
+        num_accepted_tokens: torch.Tensor,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        sequence_lengths: torch.Tensor,
+        num_contexts: int,
+        batch_indices: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Select the accepted token, hidden row, and fixed draft position."""
+        sequence_starts = torch.cumsum(
+            sequence_lengths, dim=0, dtype=torch.long) - sequence_lengths
+        recurrent_indices = sequence_starts + sequence_lengths - 1
+        recurrent_indices[num_contexts:].copy_(
+            sequence_starts[num_contexts:] +
+            num_accepted_tokens[num_contexts:] - 1)
+        draft_input_ids = accepted_tokens[batch_indices,
+                                          num_accepted_tokens - 1]
+        recurrent_hidden_states = hidden_states[recurrent_indices]
+        draft_position_ids = (
+            _select_mtp_position_ids(position_ids, recurrent_indices) + 1)
+        return draft_input_ids, recurrent_hidden_states, draft_position_ids
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 683414763511..b8edf7c40f62 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import IntEnum, auto
-from typing import TYPE_CHECKING, Any, List, Optional, Type
+from typing import TYPE_CHECKING, List, Optional, Type
 
 import torch
 from packaging.version import Version
@@ -107,62 +107,16 @@ def rejection_sampling_one_model(
 _FORCE_ACCEPT_RNG_SLOT_STRIDE = 1009
 
 
-@dataclass(frozen=True)
-class DraftModelCapabilities:
-    """Runtime ownership contract for an externally loaded draft model."""
-
-    loads_external_weights: bool = False
-    num_draft_modules: int = 0
-    owns_independent_kv_cache: bool = True
-    num_draft_kv_layers: int = 0
-    shares_target_kv_cache: bool = False
-    freezes_draft_attention_state: bool = False
-    requires_external_draft_metadata_view: bool = False
-
-    @classmethod
-    def external_shared_target_kv(cls) -> "DraftModelCapabilities":
-        return cls(
-            loads_external_weights=True,
-            num_draft_modules=1,
-            owns_independent_kv_cache=False,
-            num_draft_kv_layers=0,
-            shares_target_kv_cache=True,
-            freezes_draft_attention_state=True,
-            requires_external_draft_metadata_view=True,
-        )
-
-    @classmethod
-    def from_config(cls, config: Any) -> "DraftModelCapabilities":
-        return cls(
-            loads_external_weights=bool(
-                getattr(config, "loads_external_weights", False)),
-            num_draft_modules=int(getattr(config, "num_draft_modules", 0)),
-            owns_independent_kv_cache=bool(
-                getattr(config, "owns_independent_kv_cache", True)),
-            num_draft_kv_layers=int(getattr(config, "num_draft_kv_layers", 0)),
-            shares_target_kv_cache=bool(
-                getattr(config, "shares_target_kv_cache", False)),
-            freezes_draft_attention_state=bool(
-                getattr(config, "freezes_draft_attention_state", False)),
-            requires_external_draft_metadata_view=bool(
-                getattr(config, "requires_external_draft_metadata_view",
-                        False)),
-        )
-
-
-def get_draft_model_capabilities(
-        spec_config) -> Optional[DraftModelCapabilities]:
-    if spec_config is None:
-        return None
-    return getattr(spec_config, "_draft_model_capabilities", None)
+def is_gemma4_mtp_assistant(spec_config) -> bool:
+    return bool(spec_config is not None
+                and getattr(spec_config, "_is_gemma4_mtp_assistant", False))
 
 
 def needs_external_draft_weights(spec_config) -> bool:
     """Whether a one-engine mode loads a separate draft checkpoint."""
     if spec_config is None:
         return False
-    capabilities = get_draft_model_capabilities(spec_config)
-    if capabilities is not None and capabilities.loads_external_weights:
+    if is_gemma4_mtp_assistant(spec_config):
         return True
     return spec_config.spec_dec_mode.need_load_draft_weights()
 
@@ -170,9 +124,7 @@ def needs_external_draft_weights(spec_config) -> bool:
 def should_extend_context(spec_config,
                           attention_backend: Type[AttentionBackend]) -> bool:
     """Whether generation verification uses the backend's context kernel."""
-    capabilities = get_draft_model_capabilities(spec_config)
-    if (capabilities is not None and capabilities.shares_target_kv_cache
-            and capabilities.requires_external_draft_metadata_view):
+    if is_gemma4_mtp_assistant(spec_config):
         from ..attention_backend.flashinfer import FlashInferAttention
         if issubclass(attention_backend, FlashInferAttention):
             return True
@@ -187,9 +139,6 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
         return False
     if not spec_config.spec_dec_mode.use_one_engine():
         return False
-    capabilities = get_draft_model_capabilities(spec_config)
-    if capabilities is not None and not capabilities.owns_independent_kv_cache:
-        return False
     # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft
     # model does not read the paged draft KV cache managed by attention metadata.
     if spec_config.spec_dec_mode.is_dspark():
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 4a2c8e7e6d51..e31512227154 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -15,7 +15,7 @@
 from ..pyexecutor.guided_decoder import GuidedDecoder
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.seq_slot_manager import SeqSlotManager
-from ..speculative.interface import SpecMetadata
+from ..speculative.interface import SpecMetadata, is_gemma4_mtp_assistant
 from .dflash import DFlashSpecMetadata, DFlashWorker
 from .draft_target import (DraftTargetOneModelSampler,
                            DraftTargetOneModelSpecMetadata,
@@ -131,15 +131,13 @@ def get_spec_metadata(spec_config,
             draft_vocab_size=draft_vocab_size,
         )
     if spec_config.spec_dec_mode.is_mtp_eagle():
-        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,
             spec_dec_mode=spec_config.spec_dec_mode,
             max_num_requests=max_num_requests,
             num_layers=model_config.num_hidden_layers,
-            hidden_size=hidden_size,
+            hidden_size=model_config.hidden_size,
             max_num_tokens=max_num_tokens,
             dtype=model_config.torch_dtype,
             is_draft_model=is_draft_model,
@@ -445,10 +443,8 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
-    """Return the logical number of draft modules executed by the worker."""
-    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
-    if capabilities is not None and capabilities.num_draft_modules:
-        return capabilities.num_draft_modules
+    if is_gemma4_mtp_assistant(spec_config):
+        return 0
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
     if spec_config.spec_dec_mode.is_mtp_vanilla():
@@ -459,14 +455,6 @@ def get_num_spec_layers(spec_config):
     return 0
 
 
-def get_num_draft_kv_layers(spec_config):
-    """Return the number of draft-owned layers requiring KV cache storage."""
-    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
-    if capabilities is not None:
-        return capabilities.num_draft_kv_layers
-    return get_num_spec_layers(spec_config)
-
-
 def update_spec_config_from_draft_model_config(spec_config,
                                                draft_pretrained_config) -> None:
     """Populate Eagle draft-layer fields from the loaded draft model config."""
@@ -533,8 +521,7 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
-    capabilities = getattr(spec_config, "_draft_model_capabilities", None)
-    if capabilities is not None and capabilities.shares_target_kv_cache:
+    if is_gemma4_mtp_assistant(spec_config):
         return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
@@ -595,6 +582,14 @@ def update_spec_config_from_model_config(spec_config, model_config):
     if not spec_config.use_dynamic_tree:
         spec_config.max_total_draft_tokens = spec_config.max_draft_len
 
+    model_type = getattr(model_config, "model_type", None)
+    spec_config._is_gemma4_mtp_assistant = bool(
+        model_type in ("gemma4", "gemma4_text")
+        and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+        and spec_config.speculative_model is not None)
+    if spec_config._is_gemma4_mtp_assistant:
+        spec_config._allow_separate_draft_kv_cache = False
+
 
 def update_spec_config_from_loaded_model(spec_config, model) -> None:
     """Populate spec config fields from loaded target and draft model configs."""
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 0bea7b5439bd..406e47038044 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -2550,9 +2550,8 @@ class MTPDecodingConfig(DecodingBaseConfig):
 
     # Internal max batch size for dynamic-tree worker buffers.
     _max_batch_size: Optional[int] = PrivateAttr(default=None)
-    # Runtime-only ownership contract populated after loading an external
-    # assistant config. It is intentionally excluded from serialization.
-    _draft_model_capabilities: Any = PrivateAttr(default=None)
+    # Runtime-only marker populated from the target model config.
+    _is_gemma4_mtp_assistant: bool = PrivateAttr(default=False)
 
     sa_config: Optional[SAEnhancerConfig] = Field(
         default=None,
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 72fd330ff141..504ce47a1530 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -575,12 +575,6 @@ 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))
 
@@ -891,7 +885,6 @@ def _build_gemma4_kv_cache_manager(
     num_blocks=4,
     tokens_per_block=32,
     batch_size=1,
-    force_vswa=False,
 ):
     """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads.
 
@@ -947,7 +940,7 @@ def _build_gemma4_kv_cache_manager(
     # exceeds sliding_window.
     sliding_window = getattr(config, "sliding_window", None)
     max_attn_window = None
-    needs_vswa = force_vswa or (isinstance(head_dim, list) and len(set(head_dim)) > 1)
+    needs_vswa = 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:
@@ -2458,11 +2451,9 @@ 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"],
@@ -2478,13 +2469,12 @@ def _make_trtllm_gen_decode_case(
         if manager_batch_size is None:
             manager_batch_size = batch_size
 
-        config = Gemma4TextConfig(**deepcopy(config_dict or GEMMA4_E2B_REAL_DIMS_CONFIG))
+        config = Gemma4TextConfig(**deepcopy(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")
@@ -2603,13 +2593,14 @@ def _prepare_decode_page_counts(
     def _expected_decode_block_table(
         self,
         metadata: "FlashInferAttentionMetadata",
-        pool_id: int,
+        head_dim: 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
@@ -2694,7 +2685,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.kv_pool_id,
+                    plan_params.head_dim,
                     new_page_counts,
                     rows=len(initial_page_counts),
                     width=max(initial_page_counts),
@@ -2723,12 +2714,11 @@ 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.kv_pool_id] = (
+            initial_state[plan_params.head_dim] = (
                 block_tables.data_ptr(),
                 wrappers.host_decode_block_tables.data_ptr(),
             )
@@ -2740,13 +2730,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.kv_pool_id]
+                old_device_ptr, old_host_ptr = initial_state[plan_params.head_dim]
                 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.kv_pool_id,
+                    plan_params.head_dim,
                     new_page_counts,
                     rows=len(new_page_counts),
                     width=max(new_page_counts),
@@ -2758,48 +2748,6 @@ 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(len(plan_params), 2)
-        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 a32d26c60a40..85181f29ef6b 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
@@ -9,22 +9,13 @@
 from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
 from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
 from tensorrt_llm._torch.speculative.interface import (
-    DraftModelCapabilities,
     needs_external_draft_weights,
     should_use_separate_draft_kv_cache,
 )
-from tensorrt_llm._torch.speculative.utils import (
-    get_num_draft_kv_layers,
-    get_num_extra_kv_tokens,
-    get_num_spec_layers,
-)
+from tensorrt_llm._torch.speculative.utils import get_num_extra_kv_tokens, get_num_spec_layers
 from tensorrt_llm.llmapi import MTPDecodingConfig
 
 
-def _shared_kv_capabilities() -> DraftModelCapabilities:
-    return DraftModelCapabilities.external_shared_target_kv()
-
-
 def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
     spec_config = MTPDecodingConfig(
         max_draft_len=kwargs.pop("max_draft_len", 3),
@@ -32,27 +23,20 @@ def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
         mtp_eagle_one_model=True,
         **kwargs,
     )
-    spec_config._draft_model_capabilities = _shared_kv_capabilities()
+    spec_config._is_gemma4_mtp_assistant = True
+    spec_config._allow_separate_draft_kv_cache = False
     return spec_config
 
 
-def test_external_shared_kv_capability_separates_module_and_kv_counts():
+def test_external_shared_kv_uses_no_draft_kv_cache():
     spec_config = _shared_kv_spec_config()
 
     assert needs_external_draft_weights(spec_config)
-    assert get_num_spec_layers(spec_config) == 1
-    assert get_num_draft_kv_layers(spec_config) == 0
+    assert get_num_spec_layers(spec_config) == 0
     assert get_num_extra_kv_tokens(spec_config) == 0
     assert not should_use_separate_draft_kv_cache(spec_config)
 
 
-def test_embedded_one_model_mtp_does_not_load_external_weights():
-    spec_config = _shared_kv_spec_config()
-    spec_config._draft_model_capabilities = None
-
-    assert not needs_external_draft_weights(spec_config)
-
-
 def test_external_shared_kv_worker_rejects_unverified_modes():
     spec_config = _shared_kv_spec_config(
         use_dynamic_tree=True,
@@ -157,34 +141,3 @@ def spec_worker(**kwargs):
     assert torch.equal(worker_calls[0]["hidden_states"], hidden_states)
     assert torch.equal(worker_calls[0]["logits"], hidden_states[[2]])
     assert worker_calls[0]["draft_model"] is model.draft_model
-
-
-def test_gemma4_target_forward_still_captures_hidden_states_without_worker():
-    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(),
-        spec_worker=None,
-    )
-    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 02f628376bdb81a40b9b975ac716e78517710ff2 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 07:37:50 +0000
Subject: [PATCH 12/26] [None][refactor] minimize Gemma4 MTP VLM proxies

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4mm.py  | 16 ----------------
 .../_torch/modeling/test_gemma4_multimodal.py    |  2 +-
 2 files changed, 1 insertion(+), 17 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 30ca39eadfa9..3989215f0417 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -732,22 +732,6 @@ def post_config(self):
         self.config = self.llm.config
         self.model_config.pretrained_config = self.llm.config
 
-    @property
-    def model(self):
-        return self.llm.model
-
-    @property
-    def lm_head(self):
-        return self.llm.lm_head
-
-    @property
-    def epilogue(self):
-        return self.llm.epilogue
-
-    @property
-    def spec_worker(self):
-        return self.llm.spec_worker
-
     @property
     def draft_config(self):
         return self.llm.draft_config
diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
index 089e575148b8..427ba1bed3d7 100644
--- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
+++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
@@ -774,7 +774,7 @@ def test_encoder_cache_reuses_image_embedding_across_requests(self):
         torch.testing.assert_close(second, first)
         self.assertEqual(len(model._multimodal_encoder_cache), 1)
 
-    def test_speculative_runtime_contract_is_proxied_to_language_model(self):
+    def test_draft_weight_loading_contract_is_proxied_to_language_model(self):
         model = self._make_model()
 
         self.assertIs(model.model, model.llm.model)

From 988d31d48c6292ef4e2f74a8a0d2bb38c7ce6f76 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 08:56:07 +0000
Subject: [PATCH 13/26] [None][refactor] use native FlashInfer decode for
 Gemma4 MTP

Route Gemma4 one-model MTP verification through FlashInfer's native
multi-query decode path. Remove the context-kernel override and paged-
prefill CUDA Graph refresh, and retain coverage for graph request
turnover.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    | 54 ++-----------------
 .../_torch/pyexecutor/model_engine.py         | 34 ++++++------
 .../_torch/pyexecutor/py_executor_creator.py  |  9 ++--
 tensorrt_llm/_torch/speculative/__init__.py   |  2 -
 tensorrt_llm/_torch/speculative/interface.py  | 10 ----
 .../_torch/modeling/test_modeling_gemma4.py   |  4 +-
 6 files changed, 23 insertions(+), 90 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 051c1165213b..247a181d91be 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -226,8 +226,6 @@ 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)
@@ -852,10 +850,10 @@ def _post_init_with_buffers(self, buffers) -> None:
                 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():
-                        head_dim = mgr.head_dim_per_layer[
+                        hd = mgr.head_dim_per_layer[
                             mgr.layer_offsets[layer_idx]]
-                        if head_dim not in self._vswa_head_dim_to_pool:
-                            self._vswa_head_dim_to_pool[head_dim] = pool_id
+                        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
@@ -1475,52 +1473,6 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             self._host_paged_kv_indices = \
                 self._host_pool_indices[primary_pool_id]
 
-        # Decoder graph batches have no scheduled context requests, but linear
-        # speculative verification reclassifies draft-token extensions as
-        # contexts. Refresh each trtllm-gen paged-prefill graph's stable block
-        # table 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
-                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/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 06e55f838931..cfb0360d1c80 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -66,7 +66,6 @@
                            get_num_extra_kv_tokens, get_spec_metadata,
                            prepare_attn_metadata_for_draft_replay,
                            restore_attn_metadata_after_draft_replay,
-                           should_extend_context,
                            update_spec_config_from_loaded_model)
 from ..speculative.drafting_loops import BaseDraftingLoopWrapper
 from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata
@@ -3392,9 +3391,8 @@ def _prepare_incremental_update_metadata(
         attn_metadata.beam_width = 1
         attn_metadata.prompt_lens = prompt_lengths
         attn_metadata.num_contexts = num_extend_ctx_requests if (
-            enable_spec_decode
-            and should_extend_context(spec_config, self.attn_backend)
-            and spec_config.is_linear_tree) else 0
+            enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                self.attn_backend) and spec_config.is_linear_tree) else 0
         attn_metadata.num_chunked_ctx_requests = attn_metadata.num_contexts
 
         # Create KV cache params and prepare metadata
@@ -3673,8 +3671,9 @@ def _apply_incremental_update_target(
         num_extend_dummy_requests = 0
         num_previous_batch = 0
 
-        use_extend_ctx = (self.enable_spec_decode and should_extend_context(
-            spec_config, self.attn_backend) and spec_config.is_linear_tree)
+        use_extend_ctx = (self.enable_spec_decode
+                          and spec_config.spec_dec_mode.extend_ctx(
+                              self.attn_backend) and spec_config.is_linear_tree)
 
         for idx, request in enumerate(extend_requests):
             request_accepted_path[request.py_request_id] = \
@@ -3747,8 +3746,8 @@ def _apply_incremental_update_target(
 
         # Determine if we're using extend_ctx mode for linear tree decoding
         num_extend_ctx_requests = 0
-        if self.enable_spec_decode and should_extend_context(
-                spec_config, self.attn_backend) and spec_config.is_linear_tree:
+        if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                self.attn_backend) and spec_config.is_linear_tree:
             num_extend_ctx_requests = num_extend_requests
 
         virtual_num_tokens = num_generation_tokens
@@ -4331,8 +4330,7 @@ def append_cross_attention_state(request: LlmRequest,
                                        if is_promoted_context else
                                        request.max_beam_num_tokens - 1)
                 draft_lens.append(num_draft_tokens)
-                if self.enable_spec_decode and should_extend_context(
-                        spec_config,
+                if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend) and spec_config.is_linear_tree:
                     # We're treating the prompt lengths as context requests here, so
                     # the the prompt lens should not include the cached tokens.
@@ -4384,8 +4382,7 @@ def append_cross_attention_state(request: LlmRequest,
                     request.py_num_compressed_tokens)
                 request.cached_tokens = (past_seen_token_num +
                                          runtime_tokens_per_gen_step)
-                if self.enable_spec_decode and should_extend_context(
-                        spec_config,
+                if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
                         self.attn_backend) and spec_config.is_linear_tree:
                     prompt_lengths.append(runtime_tokens_per_gen_step)
                 else:
@@ -4983,8 +4980,8 @@ def previous_seq_slots_device():
         # Use num_chunked_ctx_requests to record the number of extend context requests,
         # so that we can update the kv_lens_cuda correctly in _preprocess_inputs.
         attn_metadata.num_chunked_ctx_requests = 0
-        if self.enable_spec_decode and should_extend_context(
-                spec_config, self.attn_backend) and spec_config.is_linear_tree:
+        if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx(
+                self.attn_backend) and spec_config.is_linear_tree:
             # For the tree decoding, we want to use XQA to process the draft tokens for the target model.
             # Therefore, we do not treat them as the chunked context requests.
             attn_metadata.num_contexts += len(extend_requests)
@@ -5629,8 +5626,8 @@ def _get_lora_params_from_requests(
             tokens_per_seq = 1
             if (self.enable_spec_decode and self.runtime_draft_len > 0
                     and self.spec_config.is_linear_tree
-                    and not should_extend_context(self.spec_config,
-                                                  self.attn_backend)):
+                    and not self.spec_config.spec_dec_mode.extend_ctx(
+                        self.attn_backend)):
                 tokens_per_seq = self.runtime_draft_len + 1
             return self.cuda_graph_lora_manager.prepare_cuda_graph_lora_params(
                 scheduled_requests, attn_metadata, peft_cache_manager,
@@ -5749,9 +5746,8 @@ def _get_eager_lora_params_from_requests(
             # count so the kernel correctly expands LoRA weights for all tokens.
             if (self.enable_spec_decode and self.runtime_draft_len > 0
                     and self.spec_config.is_linear_tree
-                    and not should_extend_context(self.spec_config,
-                                                  self.attn_backend)
-                    and num_generations > 0):
+                    and not self.spec_config.spec_dec_mode.extend_ctx(
+                        self.attn_backend) and num_generations > 0):
                 tokens_per_req = self.runtime_draft_len + 1
                 host_request_types = host_request_types.clone()
                 host_request_types[num_contexts:num_seqs].fill_(0)  # kCONTEXT
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index ce3c6962ceba..fc30d8b4eff5 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -451,9 +451,8 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
-        # External MTP assistants with shared target KV use a dedicated
-        # FlashInfer decode metadata view. Other one-engine modes still rely
-        # on the one-query-per-sequence decode contract.
+        # Gemma4 MTP assistants provide the shared-KV metadata required by the
+        # FlashInfer decode path. Other one-engine modes remain unsupported.
         supports_shared_kv_flashinfer = getattr(spec_config,
                                                 "_is_gemma4_mtp_assistant",
                                                 False)
@@ -461,9 +460,7 @@ def create_py_executor(
                 and not supports_shared_kv_flashinfer):
             raise ValueError(
                 f"FLASHINFER attention backend is not supported with one-engine speculative "
-                f"decoding mode '{spec_config.spec_dec_mode.name}'. The FLASHINFER backend's "
-                f"decode path expects exactly 1 token per sequence, but one-engine speculative "
-                f"decoding requires multiple tokens per sequence. Please use 'TRTLLM' attention "
+                f"decoding mode '{spec_config.spec_dec_mode.name}'. Please use 'TRTLLM' attention "
                 f"backend instead by setting attn_backend='TRTLLM'.")
 
     if mm_encoder_only:
diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py
index 96185991175c..dd3088ceee82 100644
--- a/tensorrt_llm/_torch/speculative/__init__.py
+++ b/tensorrt_llm/_torch/speculative/__init__.py
@@ -7,7 +7,6 @@
                         needs_external_draft_weights,
                         prepare_attn_metadata_for_draft_replay,
                         restore_attn_metadata_after_draft_replay,
-                        should_extend_context,
                         should_use_separate_draft_kv_cache)
 from .mtp import MTPSampler, MTPSpecMetadata, MTPWorker
 from .ngram import NGramDrafter, NGramPoolManager
@@ -64,7 +63,6 @@
     "prepare_attn_metadata_for_draft_replay",
     "needs_external_draft_weights",
     "restore_attn_metadata_after_draft_replay",
-    "should_extend_context",
     "should_use_separate_draft_kv_cache",
     "update_spec_config_from_draft_model_config",
     "update_spec_config_from_loaded_model",
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index b8edf7c40f62..ba5f447a0210 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -121,16 +121,6 @@ def needs_external_draft_weights(spec_config) -> bool:
     return spec_config.spec_dec_mode.need_load_draft_weights()
 
 
-def should_extend_context(spec_config,
-                          attention_backend: Type[AttentionBackend]) -> bool:
-    """Whether generation verification uses the backend's context kernel."""
-    if is_gemma4_mtp_assistant(spec_config):
-        from ..attention_backend.flashinfer import FlashInferAttention
-        if issubclass(attention_backend, FlashInferAttention):
-            return True
-    return spec_config.spec_dec_mode.extend_ctx(attention_backend)
-
-
 def should_use_separate_draft_kv_cache(spec_config) -> bool:
     """
     Check if separate draft KV cache should be used for one-engine speculative decoding.
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 504ce47a1530..9846ea5ee6a0 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -3005,7 +3005,7 @@ def test_cuda_graph_decode_hybrid_headdim(self):
         "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."""
+        """Speculative multi-query decode survives graph 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
@@ -3075,7 +3075,7 @@ def make_metadata(
         ):
             return FlashInferAttentionMetadata(
                 seq_lens=torch.tensor([verification_tokens], dtype=torch.int),
-                num_contexts=1,
+                num_contexts=0,
                 is_cuda_graph=is_cuda_graph,
                 kv_cache_params=KVCacheParams(
                     use_cache=True, num_cached_tokens_per_seq=cached_tokens

From 0d2bb3ccbea403d2976ad35b2df400a04b3c9cae Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 09:29:44 +0000
Subject: [PATCH 14/26] [None][refactor] generalize shared-KV MTP integration

Remove model-specific runtime markers and make external shared-KV behavior depend on configuration plus an explicit draft-model capability.

Reuse a local draft-weight decision and retain the FlashInfer guard only for one-engine modes whose draft KV-cache path is unsupported.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py | 17 +------
 .../_torch/models/modeling_gemma4mm.py        |  3 +-
 .../_torch/models/modeling_speculative.py     | 19 ++------
 .../_torch/pyexecutor/model_loader.py         | 17 ++++---
 .../_torch/pyexecutor/py_executor_creator.py  | 14 +++---
 tensorrt_llm/_torch/speculative/__init__.py   |  2 -
 tensorrt_llm/_torch/speculative/eagle3.py     | 31 ++++++++-----
 tensorrt_llm/_torch/speculative/interface.py  | 17 +++----
 tensorrt_llm/_torch/speculative/utils.py      | 14 ++----
 tensorrt_llm/llmapi/llm_args.py               |  2 -
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 44 +++++++++++++++----
 11 files changed, 89 insertions(+), 91 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 054225a0d28e..23ee15950f3b 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1224,20 +1224,6 @@ def forward(
         return hidden_states
 
 
-# ---------------------------------------------------------------------------
-# Gemma4 For Causal LM
-# ---------------------------------------------------------------------------
-def _configure_gemma4_mtp_assistant(model_config: ModelConfig) -> None:
-    spec_config = model_config.spec_config
-    if (
-        spec_config is not None
-        and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-        and spec_config.speculative_model is not None
-    ):
-        spec_config._is_gemma4_mtp_assistant = True
-        spec_config._allow_separate_draft_kv_cache = False
-
-
 @register_auto_model("Gemma4ForCausalLM")
 class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
     def __init__(
@@ -1259,7 +1245,6 @@ def __init__(
                 "moe_ep_size>1 requires a Gemma4 MoE variant (only 26B-A4B-it today)."
             )
 
-        _configure_gemma4_mtp_assistant(model_config)
         super().__init__(Gemma4TextModel(model_config), model_config)
 
     @classmethod
@@ -1564,6 +1549,8 @@ def forward(self, hidden_states: torch.Tensor, lm_head: nn.Module) -> torch.Tens
 class Gemma4AssistantForCausalLM(DecoderModelForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
     """Gemma4 MTP assistant that attends directly to the target model KV cache."""
 
+    shares_target_kv_cache = True
+
     def __init__(self, model_config: ModelConfig):
         assistant_config = model_config.pretrained_config
         text_model_config = dataclasses.replace(
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 3989215f0417..a5d08d4d06b8 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -49,7 +49,7 @@
 from ...sampling_params import SamplingParams
 from ..modules.embedding import Embedding
 from ..modules.linear import Linear
-from .modeling_gemma4 import Gemma4ForCausalLM, _configure_gemma4_mtp_assistant
+from .modeling_gemma4 import Gemma4ForCausalLM
 from .modeling_gemma4_audio import Gemma4AudioModel
 from .modeling_gemma4_vision import Gemma4VisionModel
 from .modeling_multimodal_mixin import MultimodalModelMixin, PreparedLlmInputs
@@ -853,7 +853,6 @@ def __init__(self, model_config: ModelConfig[Gemma4Config]):
             )
 
         config = model_config.pretrained_config
-        _configure_gemma4_mtp_assistant(model_config)
         super().__init__(config)
 
         # Pin multimodal tensors to the local rank so each rank of a multi-GPU
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 1112cdc19ef9..42de6fbdac9e 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -32,7 +32,7 @@
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
 from ..speculative import (SpecMetadata, get_spec_worker,
                            should_use_separate_draft_kv_cache)
-from ..speculative.interface import is_gemma4_mtp_assistant
+from ..speculative.interface import uses_external_shared_kv_mtp
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_auto import AutoModelForCausalLM
@@ -1890,8 +1890,8 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
-    elif (spec_dec_mode.is_mtp_eagle_one_model() and draft_config is not None
-          and is_gemma4_mtp_assistant(model_config.spec_config)):
+    elif (uses_external_shared_kv_mtp(model_config.spec_config)
+          and draft_config is not None):
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
         return MTPForCausalLM(model_config,
@@ -1996,10 +1996,7 @@ def __init__(self,
                     model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
 
-                elif (spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-                      and spec_config.speculative_model is not None
-                      and model_config.pretrained_config.model_type
-                      in ("gemma4", "gemma4_text")):
+                elif uses_external_shared_kv_mtp(spec_config):
                     self.draft_config = ModelConfig.from_pretrained(
                         spec_config.speculative_model,
                         trust_remote_code=True,
@@ -2012,14 +2009,6 @@ def __init__(self,
                     self.draft_config.quant_config.kv_cache_quant_algo = \
                         model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
-                    draft_pretrained_config = (
-                        self.draft_config.pretrained_config)
-                    draft_architectures = getattr(draft_pretrained_config,
-                                                  "architectures", None) or []
-                    if "Gemma4AssistantForCausalLM" not in draft_architectures:
-                        raise ValueError(
-                            "Gemma4 MTP requires a "
-                            "Gemma4AssistantForCausalLM checkpoint.")
 
                 elif spec_config.spec_dec_mode.is_external_drafter():
                     self.draft_config = ModelConfig.from_pretrained(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 8d3367b9cbed..31a7a2e77c75 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -41,7 +41,6 @@
                                      timing)
 from ..modules.fused_moe.moe_load_balancer import (
     MoeLoadBalancer, maybe_create_moe_load_balancer)
-from ..speculative import needs_external_draft_weights
 from ..virtual_memory import RestoreMode
 from ..virtual_memory import scope as virtual_memory_scope
 from .config_utils import (is_hybrid_linear, resolve_hf_torch_dtype,
@@ -437,8 +436,8 @@ def load_config_and_apply_defaults(
             from tensorrt_llm._torch.speculative import \
                 update_spec_config_from_model_config
 
-            # Model defaults reconstruct nested Pydantic configs. Populate
-            # runtime-only speculative state after that reconstruction.
+            # Model defaults reconstruct nested Pydantic configs and drop
+            # init=False runtime fields such as num_nextn_predict_layers.
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
 
@@ -554,7 +553,11 @@ def load(
                 model = AutoModelForCausalLM.from_config(config)
                 is_meta_init = False
 
-            loads_draft_weights = needs_external_draft_weights(self.spec_config)
+            loads_draft_weights = (
+                self.spec_config is not None
+                and (self.spec_config.spec_dec_mode.need_load_draft_weights() or
+                     (self.spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                      and self.spec_config.speculative_model is not None)))
             speculative_mode = self._speculative_mode_name(self.spec_config)
             post_transform_qualification = self._qualify_post_transform_profile(
                 model,
@@ -705,7 +708,7 @@ def init_meta_tensor(t: torch.Tensor):
                     self._call_load_weights(model.load_weights, weights,
                                             self.weight_mapper)
 
-                if needs_external_draft_weights(self.spec_config):
+                if loads_draft_weights:
                     weights = checkpoint_loader.load_weights(
                         self.spec_config.speculative_model,
                         mapping=self.mapping)
@@ -840,7 +843,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                                     "commit an unpopulated model to the GMS "
                                     "pool.")
 
-                            if needs_external_draft_weights(self.spec_config):
+                            if loads_draft_weights:
                                 draft_weights = checkpoint_loader.load_weights(
                                     self.spec_config.speculative_model,
                                     mapping=self.mapping)
@@ -967,7 +970,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                 self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
                     model, config)
                 initialize_dummy_weights(model)
-                if needs_external_draft_weights(self.spec_config):
+                if loads_draft_weights:
                     model.draft_model.load_weights_from_target_model(model)
 
             elif load_format == LoadFormat.VISION_ONLY:
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index fc30d8b4eff5..7fe8b1083ca0 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -36,6 +36,7 @@
 from ..distributed import Distributed
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager)
+from ..speculative.interface import uses_external_shared_kv_mtp
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
@@ -451,16 +452,15 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
-        # Gemma4 MTP assistants provide the shared-KV metadata required by the
-        # FlashInfer decode path. Other one-engine modes remain unsupported.
-        supports_shared_kv_flashinfer = getattr(spec_config,
-                                                "_is_gemma4_mtp_assistant",
-                                                False)
+        # Other one-engine modes may require a separate draft KV cache, which
+        # FlashInfer attention metadata does not yet support.
         if (llm_args.attn_backend == "FLASHINFER"
-                and not supports_shared_kv_flashinfer):
+                and spec_config.spec_dec_mode.use_one_engine()
+                and not uses_external_shared_kv_mtp(spec_config)):
             raise ValueError(
                 f"FLASHINFER attention backend is not supported with one-engine speculative "
-                f"decoding mode '{spec_config.spec_dec_mode.name}'. Please use 'TRTLLM' attention "
+                f"decoding mode '{spec_config.spec_dec_mode.name}' because its draft KV-cache "
+                f"path is not supported. Please use 'TRTLLM' attention "
                 f"backend instead by setting attn_backend='TRTLLM'.")
 
     if mm_encoder_only:
diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py
index dd3088ceee82..347f3aae5dc2 100644
--- a/tensorrt_llm/_torch/speculative/__init__.py
+++ b/tensorrt_llm/_torch/speculative/__init__.py
@@ -4,7 +4,6 @@
                            DraftTargetOneModelWorker)
 from .eagle3 import Eagle3SpecMetadata, MTPEagleWorker
 from .interface import (SpecMetadata, SpecWorkerBase,
-                        needs_external_draft_weights,
                         prepare_attn_metadata_for_draft_replay,
                         restore_attn_metadata_after_draft_replay,
                         should_use_separate_draft_kv_cache)
@@ -61,7 +60,6 @@
     "get_spec_resource_manager",
     "get_spec_worker",
     "prepare_attn_metadata_for_draft_replay",
-    "needs_external_draft_weights",
     "restore_attn_metadata_after_draft_replay",
     "should_use_separate_draft_kv_cache",
     "update_spec_config_from_draft_model_config",
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 1a3b7c346669..06b37eb43a00 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -20,7 +20,7 @@
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata, SpecWorkerBase, is_gemma4_mtp_assistant
+from .interface import SpecMetadata, SpecWorkerBase, uses_external_shared_kv_mtp
 from .mtp import MTPSampler, _select_mtp_position_ids
 from .sa_enhancer import SADraftEnhancer
 from .spec_tree_manager import SpecTreeManager
@@ -1291,34 +1291,43 @@ def __init__(self,
             use_separate_draft_kv_cache=use_separate_draft_kv_cache)
         # Preserved for callers/tests that still expect this attribute.
         self.is_thop = False
-        self._uses_external_shared_target_kv = is_gemma4_mtp_assistant(
-            spec_config)
+        self._uses_external_shared_target_kv = False
 
     def set_draft_model(self, draft_model) -> None:
         super().set_draft_model(draft_model)
+        expects_external_shared_target_kv = uses_external_shared_kv_mtp(
+            self.spec_config)
+        supports_shared_target_kv = bool(
+            getattr(draft_model, "shares_target_kv_cache", False))
+        if expects_external_shared_target_kv and not supports_shared_target_kv:
+            raise ValueError(
+                "External shared-target-KV MTP requires a draft model that "
+                "declares shares_target_kv_cache=True.")
+        self._uses_external_shared_target_kv = (
+            expects_external_shared_target_kv and supports_shared_target_kv)
         if not self._uses_external_shared_target_kv:
             return
         if self.use_dynamic_tree:
             raise ValueError(
-                "Gemma4 shared-target-KV MTP supports only the linear draft "
+                "External shared-target-KV MTP supports only the linear draft "
                 "path.")
         if self.spec_config.draft_len_schedule is not None:
             raise ValueError(
-                "Gemma4 shared-target-KV MTP does not support a draft length "
-                "schedule.")
+                "External shared-target-KV MTP does not support a draft "
+                "length schedule.")
         if self.sa_enhancer is not None:
             raise ValueError(
-                "Gemma4 shared-target-KV MTP does not support the suffix "
+                "External shared-target-KV MTP does not support the suffix "
                 "automaton enhancer.")
         if self.spec_config.use_rejection_sampling:
             raise ValueError(
-                "Gemma4 shared-target-KV MTP does not support rejection "
+                "External shared-target-KV MTP does not support rejection "
                 "sampling.")
 
     def set_guided_decoder(self, guided_decoder) -> bool:
         if self._uses_external_shared_target_kv:
             raise ValueError(
-                "Gemma4 shared-target-KV MTP does not support guided "
+                "External shared-target-KV MTP does not support guided "
                 "decoding.")
         return super().set_guided_decoder(guided_decoder)
 
@@ -1351,10 +1360,10 @@ def _forward_external_shared_target_kv(
         spec_metadata,
         draft_model,
     ):
-        """Draft with a Gemma4 Q-only assistant over accepted target KV."""
+        """Draft with an external Q-only assistant over accepted target KV."""
         if not isinstance(attn_metadata, FlashInferAttentionMetadata):
             raise TypeError(
-                "Gemma4 shared-target-KV MTP requires FlashInfer attention "
+                "External shared-target-KV MTP requires FlashInfer attention "
                 "metadata.")
 
         batch_size = attn_metadata.num_seqs
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index ba5f447a0210..ad9e3ed0aa8e 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -107,18 +107,11 @@ def rejection_sampling_one_model(
 _FORCE_ACCEPT_RNG_SLOT_STRIDE = 1009
 
 
-def is_gemma4_mtp_assistant(spec_config) -> bool:
+def uses_external_shared_kv_mtp(spec_config) -> bool:
+    """Whether one-engine MTP uses an external assistant over target KV."""
     return bool(spec_config is not None
-                and getattr(spec_config, "_is_gemma4_mtp_assistant", False))
-
-
-def needs_external_draft_weights(spec_config) -> bool:
-    """Whether a one-engine mode loads a separate draft checkpoint."""
-    if spec_config is None:
-        return False
-    if is_gemma4_mtp_assistant(spec_config):
-        return True
-    return spec_config.spec_dec_mode.need_load_draft_weights()
+                and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                and spec_config.speculative_model is not None)
 
 
 def should_use_separate_draft_kv_cache(spec_config) -> bool:
@@ -129,6 +122,8 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
         return False
     if not spec_config.spec_dec_mode.use_one_engine():
         return False
+    if uses_external_shared_kv_mtp(spec_config):
+        return False
     # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft
     # model does not read the paged draft KV cache managed by attention metadata.
     if spec_config.spec_dec_mode.is_dspark():
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index e31512227154..313cd31c87af 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -15,7 +15,7 @@
 from ..pyexecutor.guided_decoder import GuidedDecoder
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.seq_slot_manager import SeqSlotManager
-from ..speculative.interface import SpecMetadata, is_gemma4_mtp_assistant
+from ..speculative.interface import SpecMetadata, uses_external_shared_kv_mtp
 from .dflash import DFlashSpecMetadata, DFlashWorker
 from .draft_target import (DraftTargetOneModelSampler,
                            DraftTargetOneModelSpecMetadata,
@@ -443,7 +443,7 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
-    if is_gemma4_mtp_assistant(spec_config):
+    if uses_external_shared_kv_mtp(spec_config):
         return 0
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
@@ -521,7 +521,7 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
-    if is_gemma4_mtp_assistant(spec_config):
+    if uses_external_shared_kv_mtp(spec_config):
         return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
@@ -582,14 +582,6 @@ def update_spec_config_from_model_config(spec_config, model_config):
     if not spec_config.use_dynamic_tree:
         spec_config.max_total_draft_tokens = spec_config.max_draft_len
 
-    model_type = getattr(model_config, "model_type", None)
-    spec_config._is_gemma4_mtp_assistant = bool(
-        model_type in ("gemma4", "gemma4_text")
-        and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-        and spec_config.speculative_model is not None)
-    if spec_config._is_gemma4_mtp_assistant:
-        spec_config._allow_separate_draft_kv_cache = False
-
 
 def update_spec_config_from_loaded_model(spec_config, model) -> None:
     """Populate spec config fields from loaded target and draft model configs."""
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index 406e47038044..ab5f561070c6 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -2550,8 +2550,6 @@ class MTPDecodingConfig(DecodingBaseConfig):
 
     # Internal max batch size for dynamic-tree worker buffers.
     _max_batch_size: Optional[int] = PrivateAttr(default=None)
-    # Runtime-only marker populated from the target model config.
-    _is_gemma4_mtp_assistant: bool = PrivateAttr(default=False)
 
     sa_config: Optional[SAEnhancerConfig] = Field(
         default=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 85181f29ef6b..baa199961b44 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
@@ -6,12 +6,10 @@
 import pytest
 import torch
 
+from tensorrt_llm._torch.models import modeling_speculative
 from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
 from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm._torch.speculative.interface import (
-    needs_external_draft_weights,
-    should_use_separate_draft_kv_cache,
-)
+from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
 from tensorrt_llm._torch.speculative.utils import get_num_extra_kv_tokens, get_num_spec_layers
 from tensorrt_llm.llmapi import MTPDecodingConfig
 
@@ -23,20 +21,45 @@ def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
         mtp_eagle_one_model=True,
         **kwargs,
     )
-    spec_config._is_gemma4_mtp_assistant = True
-    spec_config._allow_separate_draft_kv_cache = False
     return spec_config
 
 
 def test_external_shared_kv_uses_no_draft_kv_cache():
     spec_config = _shared_kv_spec_config()
 
-    assert needs_external_draft_weights(spec_config)
     assert get_num_spec_layers(spec_config) == 0
     assert get_num_extra_kv_tokens(spec_config) == 0
     assert not should_use_separate_draft_kv_cache(spec_config)
 
 
+def test_external_shared_kv_builds_draft_from_external_config(monkeypatch):
+    draft_config = object()
+    expected_model = object()
+    monkeypatch.setattr(
+        modeling_speculative.AutoModelForCausalLM,
+        "from_config",
+        lambda config: expected_model,
+    )
+
+    model_config = SimpleNamespace(spec_config=_shared_kv_spec_config())
+    assert (
+        modeling_speculative.get_draft_model(
+            model_config,
+            draft_config,
+            lm_head=None,
+            model=None,
+        )
+        is expected_model
+    )
+
+
+def test_external_shared_kv_worker_requires_draft_model_capability():
+    worker = MTPEagleWorker(_shared_kv_spec_config())
+
+    with pytest.raises(ValueError, match="shares_target_kv_cache=True"):
+        worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
+
+
 def test_external_shared_kv_worker_rejects_unverified_modes():
     spec_config = _shared_kv_spec_config(
         use_dynamic_tree=True,
@@ -45,7 +68,12 @@ def test_external_shared_kv_worker_rejects_unverified_modes():
     worker = MTPEagleWorker(spec_config)
 
     with pytest.raises(ValueError, match="linear draft path"):
-        worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
+        worker.set_draft_model(
+            SimpleNamespace(
+                model=SimpleNamespace(),
+                shares_target_kv_cache=True,
+            )
+        )
 
 
 @pytest.mark.parametrize(

From da665e4f7e826d0dbe5a2ada2e165841e91142f8 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 10:23:08 +0000
Subject: [PATCH 15/26] [None][refactor] generalize shared-KV draft capability

Represent shared target KV as a runtime speculative-decoding capability instead of inferring it from an external MTP checkpoint. Let Gemma4 declare that capability, keep draft-weight loading and KV resource decisions capability-based, and remove the obsolete FlashInfer one-engine restriction.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py |  2 +
 .../_torch/models/modeling_gemma4mm.py        |  2 +
 .../_torch/models/modeling_speculative.py     |  6 +--
 .../_torch/pyexecutor/model_loader.py         | 11 +++--
 .../_torch/pyexecutor/py_executor_creator.py  | 12 ------
 tensorrt_llm/_torch/speculative/eagle3.py     |  4 +-
 tensorrt_llm/_torch/speculative/interface.py  |  9 ++--
 tensorrt_llm/_torch/speculative/utils.py      | 10 +++--
 tensorrt_llm/llmapi/llm_args.py               |  2 +
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 43 ++++++++++++++++++-
 10 files changed, 71 insertions(+), 30 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 23ee15950f3b..7dd9b2c64a46 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1226,6 +1226,8 @@ def forward(
 
 @register_auto_model("Gemma4ForCausalLM")
 class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
+    external_draft_shares_target_kv_cache = True
+
     def __init__(
         self,
         model_config: ModelConfig[Gemma4TextConfig],
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index a5d08d4d06b8..1bf5035f1f9d 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -843,6 +843,8 @@ class Gemma4ForConditionalGeneration(Gemma4MultimodalModelBase):
     - mm_token_type_ids-based bidirectional masking
     """
 
+    external_draft_shares_target_kv_cache = True
+
     def __init__(self, model_config: ModelConfig[Gemma4Config]):
         if _is_mm_disagg():
             raise NotImplementedError(
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 42de6fbdac9e..cc6560965fa9 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -32,7 +32,7 @@
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
 from ..speculative import (SpecMetadata, get_spec_worker,
                            should_use_separate_draft_kv_cache)
-from ..speculative.interface import uses_external_shared_kv_mtp
+from ..speculative.interface import uses_shared_kv_cache
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_auto import AutoModelForCausalLM
@@ -1890,7 +1890,7 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
-    elif (uses_external_shared_kv_mtp(model_config.spec_config)
+    elif (uses_shared_kv_cache(model_config.spec_config)
           and draft_config is not None):
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
@@ -1996,7 +1996,7 @@ def __init__(self,
                     model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
 
-                elif uses_external_shared_kv_mtp(spec_config):
+                elif uses_shared_kv_cache(spec_config):
                     self.draft_config = ModelConfig.from_pretrained(
                         spec_config.speculative_model,
                         trust_remote_code=True,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 31a7a2e77c75..33ca49d2be19 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -440,6 +440,12 @@ def load_config_and_apply_defaults(
             # init=False runtime fields such as num_nextn_predict_layers.
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
+            spec_config = llm_args.speculative_config
+            spec_config._use_shared_kv_cache = bool(
+                getattr(model_cls, 'external_draft_shares_target_kv_cache',
+                        False)
+                and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
+                and spec_config.speculative_model is not None)
 
         # The transceiver preference follows the checkpoint's original
         # architecture: _resolve_class may rewrite it to an execution class
@@ -555,9 +561,8 @@ def load(
 
             loads_draft_weights = (
                 self.spec_config is not None
-                and (self.spec_config.spec_dec_mode.need_load_draft_weights() or
-                     (self.spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-                      and self.spec_config.speculative_model is not None)))
+                and (self.spec_config.spec_dec_mode.need_load_draft_weights()
+                     or self.spec_config._use_shared_kv_cache))
             speculative_mode = self._speculative_mode_name(self.spec_config)
             post_transform_qualification = self._qualify_post_transform_profile(
                 model,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 7fe8b1083ca0..8cc7c8c0730d 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -36,7 +36,6 @@
 from ..distributed import Distributed
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager)
-from ..speculative.interface import uses_external_shared_kv_mtp
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
@@ -452,17 +451,6 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
-        # Other one-engine modes may require a separate draft KV cache, which
-        # FlashInfer attention metadata does not yet support.
-        if (llm_args.attn_backend == "FLASHINFER"
-                and spec_config.spec_dec_mode.use_one_engine()
-                and not uses_external_shared_kv_mtp(spec_config)):
-            raise ValueError(
-                f"FLASHINFER attention backend is not supported with one-engine speculative "
-                f"decoding mode '{spec_config.spec_dec_mode.name}' because its draft KV-cache "
-                f"path is not supported. Please use 'TRTLLM' attention "
-                f"backend instead by setting attn_backend='TRTLLM'.")
-
     if mm_encoder_only:
         llm_args.mm_encoder_only = True
         llm_args.disable_overlap_scheduler = True
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 06b37eb43a00..b655e0ef63c5 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -20,7 +20,7 @@
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata, SpecWorkerBase, uses_external_shared_kv_mtp
+from .interface import SpecMetadata, SpecWorkerBase, uses_shared_kv_cache
 from .mtp import MTPSampler, _select_mtp_position_ids
 from .sa_enhancer import SADraftEnhancer
 from .spec_tree_manager import SpecTreeManager
@@ -1295,7 +1295,7 @@ def __init__(self,
 
     def set_draft_model(self, draft_model) -> None:
         super().set_draft_model(draft_model)
-        expects_external_shared_target_kv = uses_external_shared_kv_mtp(
+        expects_external_shared_target_kv = uses_shared_kv_cache(
             self.spec_config)
         supports_shared_target_kv = bool(
             getattr(draft_model, "shares_target_kv_cache", False))
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index ad9e3ed0aa8e..b0f3065c5948 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -107,11 +107,10 @@ def rejection_sampling_one_model(
 _FORCE_ACCEPT_RNG_SLOT_STRIDE = 1009
 
 
-def uses_external_shared_kv_mtp(spec_config) -> bool:
-    """Whether one-engine MTP uses an external assistant over target KV."""
+def uses_shared_kv_cache(spec_config) -> bool:
+    """Whether the draft model attends directly over the target KV cache."""
     return bool(spec_config is not None
-                and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-                and spec_config.speculative_model is not None)
+                and getattr(spec_config, "_use_shared_kv_cache", False))
 
 
 def should_use_separate_draft_kv_cache(spec_config) -> bool:
@@ -122,7 +121,7 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
         return False
     if not spec_config.spec_dec_mode.use_one_engine():
         return False
-    if uses_external_shared_kv_mtp(spec_config):
+    if uses_shared_kv_cache(spec_config):
         return False
     # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft
     # model does not read the paged draft KV cache managed by attention metadata.
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 313cd31c87af..f63f5c89621c 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -15,7 +15,7 @@
 from ..pyexecutor.guided_decoder import GuidedDecoder
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.seq_slot_manager import SeqSlotManager
-from ..speculative.interface import SpecMetadata, uses_external_shared_kv_mtp
+from ..speculative.interface import SpecMetadata, uses_shared_kv_cache
 from .dflash import DFlashSpecMetadata, DFlashWorker
 from .draft_target import (DraftTargetOneModelSampler,
                            DraftTargetOneModelSpecMetadata,
@@ -443,7 +443,7 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
-    if uses_external_shared_kv_mtp(spec_config):
+    if uses_shared_kv_cache(spec_config):
         return 0
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
@@ -521,7 +521,7 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
-    if uses_external_shared_kv_mtp(spec_config):
+    if uses_shared_kv_cache(spec_config):
         return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
@@ -586,6 +586,10 @@ def update_spec_config_from_model_config(spec_config, model_config):
 def update_spec_config_from_loaded_model(spec_config, model) -> None:
     """Populate spec config fields from loaded target and draft model configs."""
     update_spec_config_from_model_config(spec_config, model.config)
+    draft_model = getattr(model, 'draft_model', None)
+    spec_config._use_shared_kv_cache = bool(
+        draft_model is not None
+        and getattr(draft_model, 'shares_target_kv_cache', False))
     draft_config = getattr(model, 'draft_config', None)
     if draft_config is not None:
         update_spec_config_from_draft_model_config(
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index ab5f561070c6..c4c8682e3819 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -1808,6 +1808,8 @@ class DecodingBaseConfig(StrictBaseModel):
     _decoding_type_alias: Optional[str] = PrivateAttr(default=None)
     # If set, drafting will use separate KV cache in one-model speculative decoding.
     _allow_separate_draft_kv_cache: bool = PrivateAttr(True)
+    # If set, the draft model attends directly over the target model KV cache.
+    _use_shared_kv_cache: bool = PrivateAttr(False)
     # Internal: true when draft_len_schedule was auto-translated from max_concurrency.
     _translated_from_max_concurrency: bool = PrivateAttr(False)
 
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 baa199961b44..cfe7f856cf0e 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
@@ -9,8 +9,15 @@
 from tensorrt_llm._torch.models import modeling_speculative
 from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
 from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
-from tensorrt_llm._torch.speculative.utils import get_num_extra_kv_tokens, get_num_spec_layers
+from tensorrt_llm._torch.speculative.interface import (
+    should_use_separate_draft_kv_cache,
+    uses_shared_kv_cache,
+)
+from tensorrt_llm._torch.speculative.utils import (
+    get_num_extra_kv_tokens,
+    get_num_spec_layers,
+    update_spec_config_from_loaded_model,
+)
 from tensorrt_llm.llmapi import MTPDecodingConfig
 
 
@@ -21,12 +28,44 @@ def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
         mtp_eagle_one_model=True,
         **kwargs,
     )
+    spec_config._use_shared_kv_cache = True
     return spec_config
 
 
+def test_external_checkpoint_does_not_imply_shared_kv_cache():
+    spec_config = MTPDecodingConfig(
+        max_draft_len=3,
+        speculative_model="/tmp/assistant",
+        mtp_eagle_one_model=True,
+    )
+
+    assert not uses_shared_kv_cache(spec_config)
+    assert get_num_spec_layers(spec_config) == 1
+    assert get_num_extra_kv_tokens(spec_config) == 2
+    assert should_use_separate_draft_kv_cache(spec_config)
+
+
+def test_loaded_draft_capability_updates_runtime_spec_config():
+    spec_config = MTPDecodingConfig(
+        max_draft_len=3,
+        speculative_model="/tmp/gemma4-assistant",
+        mtp_eagle_one_model=True,
+    )
+    model = SimpleNamespace(
+        config=SimpleNamespace(num_nextn_predict_layers=1),
+        draft_config=None,
+        draft_model=SimpleNamespace(shares_target_kv_cache=True),
+    )
+
+    update_spec_config_from_loaded_model(spec_config, model)
+
+    assert uses_shared_kv_cache(spec_config)
+
+
 def test_external_shared_kv_uses_no_draft_kv_cache():
     spec_config = _shared_kv_spec_config()
 
+    assert uses_shared_kv_cache(spec_config)
     assert get_num_spec_layers(spec_config) == 0
     assert get_num_extra_kv_tokens(spec_config) == 0
     assert not should_use_separate_draft_kv_cache(spec_config)

From f58eac0c282dd2ef2128144ec4add324b5a7034a Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 11:07:36 +0000
Subject: [PATCH 16/26] [None][refactor] refine shared-KV FlashInfer
 integration

Use explicit KV pool identities for FlashInfer VSWA plans, restore shared-KV aliases after one-engine model loading, and reject unsupported non-shared one-engine configurations. Add focused regression coverage and update the Gemma4 example.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 examples/models/core/gemma/README.md          |  1 +
 .../_torch/attention_backend/flashinfer.py    | 19 ++----
 .../_torch/models/modeling_speculative.py     |  9 +++
 .../_torch/pyexecutor/py_executor_creator.py  |  8 +++
 .../_torch/modeling/test_modeling_gemma4.py   | 61 ++++++++++++++++---
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 16 +++++
 .../_torch/speculative/hw_agnostic/test_sa.py | 10 ++-
 7 files changed, 95 insertions(+), 29 deletions(-)

diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md
index b32f75ac076b..721c160aadeb 100644
--- a/examples/models/core/gemma/README.md
+++ b/examples/models/core/gemma/README.md
@@ -85,6 +85,7 @@ python3 examples/llm-api/quickstart_advanced.py \
     --draft_model_dir google/gemma-4-E4B-it-assistant \
     --spec_decode_algo MTP \
     --spec_decode_max_draft_len 3 \
+    --use_one_model \
     --disable_kv_cache_reuse \
     --apply_chat_template
 ```
diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 247a181d91be..508ce7ee3cf9 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -166,6 +166,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.
@@ -846,15 +847,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.
@@ -1144,10 +1136,8 @@ def _build_decode_block_tables(
         num_gens = self.num_generations
         if num_gens == 0:
             return None
-        pool_id = getattr(self, "_vswa_head_dim_to_pool",
-                          {}).get(plan_params.head_dim)
         host_paged_kv_indices = self._host_pool_indices.get(
-            pool_id, self._host_paged_kv_indices)
+            plan_params.kv_pool_id, self._host_paged_kv_indices)
         if host_paged_kv_indices is None:
             return None
         gen_num_blocks = np.asarray(self.num_blocks[self.num_contexts:],
@@ -1479,7 +1469,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
@@ -1487,8 +1476,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
@@ -1591,6 +1579,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/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index cc6560965fa9..ad8c5fe142c7 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -2045,6 +2045,15 @@ def __init__(self,
                 self.epilogue.append(self.spec_worker)
         self.layer_idx = -1
 
+    def setup_aliases(self) -> None:
+        if (self.draft_model is not None
+                and getattr(self.draft_model, "shares_target_kv_cache", False)):
+            setup_target_aliases = getattr(self.draft_model,
+                                           "load_weights_from_target_model",
+                                           None)
+            if callable(setup_target_aliases):
+                setup_target_aliases(self)
+
     def forward(
         self,
         attn_metadata: AttentionMetadata,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 8cc7c8c0730d..93da715b417c 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -36,6 +36,7 @@
 from ..distributed import Distributed
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager)
+from ..speculative.interface import uses_shared_kv_cache
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
@@ -451,6 +452,13 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
+    if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
+            and spec_config.spec_dec_mode.use_one_engine()
+            and not uses_shared_kv_cache(spec_config)):
+        raise ValueError(
+            "FLASHINFER attention backend supports one-engine speculative "
+            "decoding only when the draft model shares the target KV cache.")
+
     if mm_encoder_only:
         llm_args.mm_encoder_only = True
         llm_args.disable_overlap_scheduler = True
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 9846ea5ee6a0..54557f116e22 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -885,6 +885,7 @@ def _build_gemma4_kv_cache_manager(
     num_blocks=4,
     tokens_per_block=32,
     batch_size=1,
+    force_vswa=False,
 ):
     """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads.
 
@@ -940,7 +941,7 @@ def _build_gemma4_kv_cache_manager(
     # 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:
@@ -2451,9 +2452,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"],
@@ -2469,12 +2472,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")
@@ -2593,14 +2597,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
@@ -2685,7 +2688,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),
@@ -2714,11 +2717,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(),
             )
@@ -2730,13 +2734,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),
@@ -2748,6 +2752,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"]
+        page_counts = [5, 3]
+        _, _, metadata, _, _, _ = self._make_trtllm_gen_decode_case(
+            page_counts,
+            config_dict=config_dict,
+            force_vswa=True,
+        )
+
+        plan_params = list(metadata._plan_params_to_wrappers)
+        self.assertEqual(len(plan_params), 2)
+        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():
+            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 cfe7f856cf0e..6ce77fa42051 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
@@ -92,6 +92,22 @@ def test_external_shared_kv_builds_draft_from_external_config(monkeypatch):
     )
 
 
+def test_shared_kv_alias_setup_rebinds_target_model():
+    calls = []
+    draft_model = SimpleNamespace(
+        shares_target_kv_cache=True,
+        load_weights_from_target_model=lambda target: calls.append(target),
+    )
+    model = SimpleNamespace(draft_model=draft_model)
+
+    modeling_speculative.SpecDecOneEngineForCausalLM.setup_aliases(model)
+
+    assert calls == [model]
+
+    model.draft_model = SimpleNamespace(shares_target_kv_cache=True)
+    modeling_speculative.SpecDecOneEngineForCausalLM.setup_aliases(model)
+
+
 def test_external_shared_kv_worker_requires_draft_model_capability():
     worker = MTPEagleWorker(_shared_kv_spec_config())
 
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py
index bf2a6dff6376..1486f49d0a72 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py
@@ -34,12 +34,10 @@ def get_perf_metrics(result):
 # - attn_backend: Attention implementation (TRTLLM only - FLASHINFER not supported)
 # - max_matching_ngram_size: SA matching mode (2=fixed size, -1=longest match)
 #
-# NOTE: FLASHINFER backend is NOT supported for one-engine speculative decoding modes
-# (SA, MTP, Eagle3-one-model). The FLASHINFER backend's decode path expects exactly
-# 1 token per sequence, but one-engine speculative decoding requires processing multiple
-# tokens per generation sequence (last_accepted + draft_tokens). This is a fundamental
-# architectural limitation of the FLASHINFER integration that would require significant
-# changes to support multi-token generation sequences.
+# NOTE: FLASHINFER target decode supports multiple queries per request, but
+# non-shared one-engine modes still require a separate draft KV cache. The
+# draft KV metadata/manager swap is currently implemented only for TRTLLM
+# attention. Shared-target-KV modes use a separate FlashInfer metadata view.
 @pytest.mark.parametrize(
     "disable_overlap_scheduler,use_cuda_graph,attn_backend,max_matching_ngram_size",
     [

From 192d76b04f462e0a7e808a101b36332c5057cb2c Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 27 Jul 2026 15:39:39 +0000
Subject: [PATCH 17/26] [None][refactor] scope shared-KV config to Gemma4

Derive the one-model shared-KV setting from Gemma4 target architectures during speculative config loading. Remove the generic target-model capability flag and retain focused coverage for Gemma4 and non-Gemma4 configurations.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py |  2 --
 .../_torch/models/modeling_gemma4mm.py        |  2 --
 .../_torch/pyexecutor/model_loader.py         |  6 -----
 tensorrt_llm/_torch/speculative/utils.py      | 14 +++++++---
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 27 ++++++++++++-------
 5 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 7dd9b2c64a46..23ee15950f3b 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1226,8 +1226,6 @@ def forward(
 
 @register_auto_model("Gemma4ForCausalLM")
 class Gemma4ForCausalLM(SpecDecOneEngineForCausalLM[Gemma4TextModel, Gemma4TextConfig]):
-    external_draft_shares_target_kv_cache = True
-
     def __init__(
         self,
         model_config: ModelConfig[Gemma4TextConfig],
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
index 1bf5035f1f9d..a5d08d4d06b8 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py
@@ -843,8 +843,6 @@ class Gemma4ForConditionalGeneration(Gemma4MultimodalModelBase):
     - mm_token_type_ids-based bidirectional masking
     """
 
-    external_draft_shares_target_kv_cache = True
-
     def __init__(self, model_config: ModelConfig[Gemma4Config]):
         if _is_mm_disagg():
             raise NotImplementedError(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 33ca49d2be19..2a6ea23076aa 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -440,12 +440,6 @@ def load_config_and_apply_defaults(
             # init=False runtime fields such as num_nextn_predict_layers.
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
-            spec_config = llm_args.speculative_config
-            spec_config._use_shared_kv_cache = bool(
-                getattr(model_cls, 'external_draft_shares_target_kv_cache',
-                        False)
-                and spec_config.spec_dec_mode.is_mtp_eagle_one_model()
-                and spec_config.speculative_model is not None)
 
         # The transceiver preference follows the checkpoint's original
         # architecture: _resolve_class may rewrite it to an execution class
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index f63f5c89621c..821c0c7f08e4 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -37,6 +37,11 @@
                                 SaveHiddenStatesSpecMetadata)
 from .suffix_automaton import SuffixAutomatonManager
 
+_GEMMA4_SHARED_KV_TARGET_ARCHITECTURES = (
+    "Gemma4ForCausalLM",
+    "Gemma4ForConditionalGeneration",
+)
+
 
 def _is_effective_dynamic_tree(spec_config) -> bool:
     # At dynamic_tree_max_topK == 1 the tree collapses to a linear chain; route
@@ -547,6 +552,11 @@ def update_spec_config_from_model_config(spec_config, model_config):
     from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
     if not isinstance(spec_config, MTPDecodingConfig):
         return
+    architectures = getattr(model_config, "architectures", None) or ()
+    if (architectures
+            and architectures[0] in _GEMMA4_SHARED_KV_TARGET_ARCHITECTURES):
+        spec_config._use_shared_kv_cache = (
+            spec_config.spec_dec_mode.is_mtp_eagle_one_model())
     # Read the MTP layer count from the model's pretrained config. This
     # determines the actual MTP layer count in the checkpoint and drives the
     # spec_dec_mode decision (EAGLE vs vanilla MTP). Different checkpoints expose
@@ -586,10 +596,6 @@ def update_spec_config_from_model_config(spec_config, model_config):
 def update_spec_config_from_loaded_model(spec_config, model) -> None:
     """Populate spec config fields from loaded target and draft model configs."""
     update_spec_config_from_model_config(spec_config, model.config)
-    draft_model = getattr(model, 'draft_model', None)
-    spec_config._use_shared_kv_cache = bool(
-        draft_model is not None
-        and getattr(draft_model, 'shares_target_kv_cache', False))
     draft_config = getattr(model, 'draft_config', None)
     if draft_config is not None:
         update_spec_config_from_draft_model_config(
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 6ce77fa42051..9cb7379a35d6 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
@@ -16,7 +16,7 @@
 from tensorrt_llm._torch.speculative.utils import (
     get_num_extra_kv_tokens,
     get_num_spec_layers,
-    update_spec_config_from_loaded_model,
+    update_spec_config_from_model_config,
 )
 from tensorrt_llm.llmapi import MTPDecodingConfig
 
@@ -38,6 +38,12 @@ def test_external_checkpoint_does_not_imply_shared_kv_cache():
         speculative_model="/tmp/assistant",
         mtp_eagle_one_model=True,
     )
+    model_config = SimpleNamespace(
+        architectures=["LlamaForCausalLM"],
+        num_nextn_predict_layers=1,
+    )
+
+    update_spec_config_from_model_config(spec_config, model_config)
 
     assert not uses_shared_kv_cache(spec_config)
     assert get_num_spec_layers(spec_config) == 1
@@ -45,21 +51,24 @@ def test_external_checkpoint_does_not_imply_shared_kv_cache():
     assert should_use_separate_draft_kv_cache(spec_config)
 
 
-def test_loaded_draft_capability_updates_runtime_spec_config():
+@pytest.mark.parametrize("one_model,expected", [(True, True), (False, False)])
+def test_gemma4_config_sets_shared_kv_cache_for_one_model_only(
+    one_model,
+    expected,
+):
     spec_config = MTPDecodingConfig(
         max_draft_len=3,
         speculative_model="/tmp/gemma4-assistant",
-        mtp_eagle_one_model=True,
+        mtp_eagle_one_model=one_model,
     )
-    model = SimpleNamespace(
-        config=SimpleNamespace(num_nextn_predict_layers=1),
-        draft_config=None,
-        draft_model=SimpleNamespace(shares_target_kv_cache=True),
+    model_config = SimpleNamespace(
+        architectures=["Gemma4ForConditionalGeneration"],
+        num_nextn_predict_layers=1,
     )
 
-    update_spec_config_from_loaded_model(spec_config, model)
+    update_spec_config_from_model_config(spec_config, model_config)
 
-    assert uses_shared_kv_cache(spec_config)
+    assert uses_shared_kv_cache(spec_config) is expected
 
 
 def test_external_shared_kv_uses_no_draft_kv_cache():

From 9529bbb4c46d174e02158601fa21529e40b49edd Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 28 Jul 2026 04:31:31 +0000
Subject: [PATCH 18/26] [None][fix] stabilize Gemma4 shared-KV MTP

Keep the shared-KV decision scoped to Gemma4 model loading, remove the redundant helper, and preserve FlashInfer CUDA graph plans while refreshing block tables.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    | 12 +++++---
 .../_torch/models/modeling_speculative.py     |  5 ++--
 .../_torch/pyexecutor/model_loader.py         | 14 ++++++++-
 .../_torch/pyexecutor/py_executor_creator.py  |  3 +-
 tensorrt_llm/_torch/speculative/eagle3.py     |  6 ++--
 tensorrt_llm/_torch/speculative/interface.py  |  8 +----
 tensorrt_llm/_torch/speculative/utils.py      | 16 ++--------
 .../_torch/modeling/test_modeling_gemma4.py   |  5 ++++
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 29 ++-----------------
 9 files changed, 39 insertions(+), 59 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 508ce7ee3cf9..4c5c5a9f9574 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -687,10 +687,14 @@ def _sync_shared_kv_draft_view(
                         target.seq_lens_kv_cuda[:num_seqs])
         self._shared_kv_runtime_lens[:num_seqs].copy_(full_kv_lens)
 
-        # Re-plan every known assistant wrapper outside graph capture. The
-        # assistant forces trtllm-gen, whose decode plan does not retain
-        # workspace state shared with another wrapper.
-        self._clean_cached_plans(defer_plan=False)
+        if self.is_cuda_graph:
+            # Graph replay keeps the captured trtllm-gen plan. Refresh only
+            # its stable block-table buffers; re-planning the wrapper would
+            # mutate state owned by the captured kernel.
+            for plan_params, wrappers in self._plan_params_to_wrappers.items():
+                self._build_decode_block_tables(plan_params, wrappers)
+        else:
+            self._clean_cached_plans(defer_plan=False)
 
     def update_shared_kv_draft_lengths(
         self,
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index ad8c5fe142c7..35b477960bf0 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -32,7 +32,6 @@
 from ..pyexecutor.guided_decoder import CapturableGuidedDecoder
 from ..speculative import (SpecMetadata, get_spec_worker,
                            should_use_separate_draft_kv_cache)
-from ..speculative.interface import uses_shared_kv_cache
 from ..utils import AuxStreamType
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .modeling_auto import AutoModelForCausalLM
@@ -1890,7 +1889,7 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
-    elif (uses_shared_kv_cache(model_config.spec_config)
+    elif (model_config.spec_config._use_shared_kv_cache
           and draft_config is not None):
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
@@ -1996,7 +1995,7 @@ def __init__(self,
                     model_config.quant_config.kv_cache_quant_algo
                     self.draft_config.extra_attrs = model_config.extra_attrs
 
-                elif uses_shared_kv_cache(spec_config):
+                elif spec_config._use_shared_kv_cache:
                     self.draft_config = ModelConfig.from_pretrained(
                         spec_config.speculative_model,
                         trust_remote_code=True,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 2a6ea23076aa..15d175ff8d53 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -21,7 +21,7 @@
 from tensorrt_llm._utils import str_dtype_to_torch
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           ExecutorMemoryType,
-                                          ModelExpressConfig,
+                                          ModelExpressConfig, MTPDecodingConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
 from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
@@ -52,6 +52,10 @@
     "auto": "auto"
 }
 _VALID_KV_CACHE_DTYPES = ("fp8", "nvfp4", "auto")
+_GEMMA4_TARGET_ARCHITECTURES = (
+    "Gemma4ForCausalLM",
+    "Gemma4ForConditionalGeneration",
+)
 
 
 def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig,
@@ -440,6 +444,14 @@ def load_config_and_apply_defaults(
             # init=False runtime fields such as num_nextn_predict_layers.
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
+            architectures = getattr(config.pretrained_config, "architectures",
+                                    None) or ()
+            if (isinstance(llm_args.speculative_config, MTPDecodingConfig)
+                    and architectures
+                    and architectures[0] in _GEMMA4_TARGET_ARCHITECTURES):
+                llm_args.speculative_config._use_shared_kv_cache = (
+                    llm_args.speculative_config.spec_dec_mode.
+                    is_mtp_eagle_one_model())
 
         # The transceiver preference follows the checkpoint's original
         # architecture: _resolve_class may rewrite it to an execution class
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 93da715b417c..03ca7ac061c1 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -36,7 +36,6 @@
 from ..distributed import Distributed
 from ..speculative import (get_num_extra_kv_tokens, get_spec_drafter,
                            get_spec_resource_manager)
-from ..speculative.interface import uses_shared_kv_cache
 from ..virtual_memory import scope as virtual_memory_scope
 from ._util import (KvCacheCreator, _adjust_torch_mem_fraction,
                     create_py_executor_instance, instantiate_sampler, is_mla,
@@ -454,7 +453,7 @@ def create_py_executor(
 
     if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
             and spec_config.spec_dec_mode.use_one_engine()
-            and not uses_shared_kv_cache(spec_config)):
+            and not spec_config._use_shared_kv_cache):
         raise ValueError(
             "FLASHINFER attention backend supports one-engine speculative "
             "decoding only when the draft model shares the target KV cache.")
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index b655e0ef63c5..58fcee144e06 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -20,7 +20,7 @@
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.scheduler import ScheduledRequests
-from .interface import SpecMetadata, SpecWorkerBase, uses_shared_kv_cache
+from .interface import SpecMetadata, SpecWorkerBase
 from .mtp import MTPSampler, _select_mtp_position_ids
 from .sa_enhancer import SADraftEnhancer
 from .spec_tree_manager import SpecTreeManager
@@ -1295,8 +1295,8 @@ def __init__(self,
 
     def set_draft_model(self, draft_model) -> None:
         super().set_draft_model(draft_model)
-        expects_external_shared_target_kv = uses_shared_kv_cache(
-            self.spec_config)
+        expects_external_shared_target_kv = (
+            self.spec_config._use_shared_kv_cache)
         supports_shared_target_kv = bool(
             getattr(draft_model, "shares_target_kv_cache", False))
         if expects_external_shared_target_kv and not supports_shared_target_kv:
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index b0f3065c5948..9c9d01271956 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -107,12 +107,6 @@ def rejection_sampling_one_model(
 _FORCE_ACCEPT_RNG_SLOT_STRIDE = 1009
 
 
-def uses_shared_kv_cache(spec_config) -> bool:
-    """Whether the draft model attends directly over the target KV cache."""
-    return bool(spec_config is not None
-                and getattr(spec_config, "_use_shared_kv_cache", False))
-
-
 def should_use_separate_draft_kv_cache(spec_config) -> bool:
     """
     Check if separate draft KV cache should be used for one-engine speculative decoding.
@@ -121,7 +115,7 @@ def should_use_separate_draft_kv_cache(spec_config) -> bool:
         return False
     if not spec_config.spec_dec_mode.use_one_engine():
         return False
-    if uses_shared_kv_cache(spec_config):
+    if spec_config._use_shared_kv_cache:
         return False
     # DSpark owns a dedicated rolling-window cache in DSparkWorker. Its draft
     # model does not read the paged draft KV cache managed by attention metadata.
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 821c0c7f08e4..78aca6b8c498 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -15,7 +15,7 @@
 from ..pyexecutor.guided_decoder import GuidedDecoder
 from ..pyexecutor.sampler import TorchSampler
 from ..pyexecutor.seq_slot_manager import SeqSlotManager
-from ..speculative.interface import SpecMetadata, uses_shared_kv_cache
+from ..speculative.interface import SpecMetadata
 from .dflash import DFlashSpecMetadata, DFlashWorker
 from .draft_target import (DraftTargetOneModelSampler,
                            DraftTargetOneModelSpecMetadata,
@@ -37,11 +37,6 @@
                                 SaveHiddenStatesSpecMetadata)
 from .suffix_automaton import SuffixAutomatonManager
 
-_GEMMA4_SHARED_KV_TARGET_ARCHITECTURES = (
-    "Gemma4ForCausalLM",
-    "Gemma4ForConditionalGeneration",
-)
-
 
 def _is_effective_dynamic_tree(spec_config) -> bool:
     # At dynamic_tree_max_topK == 1 the tree collapses to a linear chain; route
@@ -448,7 +443,7 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
-    if uses_shared_kv_cache(spec_config):
+    if spec_config._use_shared_kv_cache:
         return 0
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
@@ -526,7 +521,7 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
-    if uses_shared_kv_cache(spec_config):
+    if spec_config._use_shared_kv_cache:
         return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
@@ -552,11 +547,6 @@ def update_spec_config_from_model_config(spec_config, model_config):
     from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
     if not isinstance(spec_config, MTPDecodingConfig):
         return
-    architectures = getattr(model_config, "architectures", None) or ()
-    if (architectures
-            and architectures[0] in _GEMMA4_SHARED_KV_TARGET_ARCHITECTURES):
-        spec_config._use_shared_kv_cache = (
-            spec_config.spec_dec_mode.is_mtp_eagle_one_model())
     # Read the MTP layer count from the model's pretrained config. This
     # determines the actual MTP layer count in the checkpoint and drives the
     # spec_dec_mode decision (EAGLE vs vanilla MTP). Different checkpoints expose
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 54557f116e22..479c29dc4ee5 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -2664,6 +2664,11 @@ def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) ->
                 atol=0,
                 rtol=0,
             )
+        with unittest.mock.patch.object(
+            draft_metadata, "_plan_with_params", wraps=draft_metadata._plan_with_params
+        ) as replan:
+            metadata.prepare()
+        replan.assert_not_called()
 
     @torch.no_grad()
     @unittest.mock.patch(
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 9cb7379a35d6..c23f8d227cb3 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
@@ -9,10 +9,7 @@
 from tensorrt_llm._torch.models import modeling_speculative
 from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
 from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm._torch.speculative.interface import (
-    should_use_separate_draft_kv_cache,
-    uses_shared_kv_cache,
-)
+from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
 from tensorrt_llm._torch.speculative.utils import (
     get_num_extra_kv_tokens,
     get_num_spec_layers,
@@ -45,36 +42,16 @@ def test_external_checkpoint_does_not_imply_shared_kv_cache():
 
     update_spec_config_from_model_config(spec_config, model_config)
 
-    assert not uses_shared_kv_cache(spec_config)
+    assert not spec_config._use_shared_kv_cache
     assert get_num_spec_layers(spec_config) == 1
     assert get_num_extra_kv_tokens(spec_config) == 2
     assert should_use_separate_draft_kv_cache(spec_config)
 
 
-@pytest.mark.parametrize("one_model,expected", [(True, True), (False, False)])
-def test_gemma4_config_sets_shared_kv_cache_for_one_model_only(
-    one_model,
-    expected,
-):
-    spec_config = MTPDecodingConfig(
-        max_draft_len=3,
-        speculative_model="/tmp/gemma4-assistant",
-        mtp_eagle_one_model=one_model,
-    )
-    model_config = SimpleNamespace(
-        architectures=["Gemma4ForConditionalGeneration"],
-        num_nextn_predict_layers=1,
-    )
-
-    update_spec_config_from_model_config(spec_config, model_config)
-
-    assert uses_shared_kv_cache(spec_config) is expected
-
-
 def test_external_shared_kv_uses_no_draft_kv_cache():
     spec_config = _shared_kv_spec_config()
 
-    assert uses_shared_kv_cache(spec_config)
+    assert spec_config._use_shared_kv_cache
     assert get_num_spec_layers(spec_config) == 0
     assert get_num_extra_kv_tokens(spec_config) == 0
     assert not should_use_separate_draft_kv_cache(spec_config)

From b89c10b7f4e37acb0df4ca2642e736a50bb790d3 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 28 Jul 2026 06:04:22 +0000
Subject: [PATCH 19/26] [None][feat] support separate-KV FlashInfer drafting

Add a manager-specific FlashInfer draft metadata view for one-engine speculative decoding while preserving shared-KV behavior and CUDA graph buffer stability.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    | 181 ++++++++++++++----
 .../_torch/pyexecutor/py_executor_creator.py  |   7 -
 tensorrt_llm/_torch/speculative/eagle3.py     |  13 +-
 tensorrt_llm/_torch/speculative/interface.py  |  21 +-
 .../attention/test_flashinfer_attention.py    |  85 ++++++++
 .../_torch/modeling/test_modeling_gemma4.py   |   4 +-
 6 files changed, 256 insertions(+), 55 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 4c5c5a9f9574..94c9fee825cb 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -269,12 +269,18 @@ class FlashInferAttentionMetadata(AttentionMetadata):
 
     _multi_item_params: Optional[FlashInferMultiItemParams] = field(
         init=False, default=None)
-    _shared_kv_draft_metadata: Optional["FlashInferAttentionMetadata"] = field(
+    _draft_metadata: Optional["FlashInferAttentionMetadata"] = field(
         init=False, default=None, repr=False)
-    _shared_kv_runtime_lens: torch.Tensor = field(init=False, repr=False)
+    _draft_kv_runtime_lens: torch.Tensor = field(init=False, repr=False)
     _is_shared_kv_draft_view: bool = field(init=False,
                                            default=False,
                                            repr=False)
+    _is_separate_kv_draft_view: bool = field(init=False,
+                                             default=False,
+                                             repr=False)
+    _uses_full_draft_page_table: bool = field(init=False,
+                                              default=False,
+                                              repr=False)
 
     def needs_plan(self, plan_params: PlanParams) -> bool:
         if plan_params not in self._plan_params_to_wrappers:
@@ -296,14 +302,14 @@ def get_decode_wrapper(
     ) -> flashinfer.BatchDecodeWithPagedKVCacheWrapper:
         assert plan_params in self._plan_params_to_wrappers, "Plan params not found, make sure to call plan()"
         result = self._plan_params_to_wrappers[plan_params].decode_wrapper
-        if self._is_shared_kv_draft_view:
+        if self._is_shared_kv_draft_view or self._is_separate_kv_draft_view:
             if result._backend != "trtllm-gen":
                 raise ValueError(
-                    "The shared-target-KV draft metadata view requires the "
-                    "FlashInfer trtllm-gen decode backend.")
+                    "FlashInfer draft metadata views require the trtllm-gen "
+                    "decode backend.")
             num_seqs = self.num_seqs
             result._kv_lens_buffer[:num_seqs].copy_(
-                self._shared_kv_runtime_lens[:num_seqs])
+                self._draft_kv_runtime_lens[:num_seqs])
         return result
 
     def get_ragged_prefill_wrapper(
@@ -594,42 +600,82 @@ def batch_indices(self) -> torch.Tensor:
     def positions(self) -> torch.Tensor:
         return self._positions[:self.num_tokens]
 
-    def get_shared_kv_draft_metadata(self) -> "FlashInferAttentionMetadata":
-        """Return a one-query decode view over this metadata's target KV."""
-        if self._shared_kv_draft_metadata is None:
+    def get_draft_metadata(
+        self,
+        draft_kv_cache_manager: Optional[Any] = None,
+    ) -> "FlashInferAttentionMetadata":
+        """Return a planned draft view over shared or separate KV cache."""
+        uses_shared_kv_cache = draft_kv_cache_manager is None
+        if self._draft_metadata is None:
             draft_metadata = copy.copy(self)
-            draft_metadata._is_shared_kv_draft_view = True
-            draft_metadata._shared_kv_draft_metadata = None
+            draft_metadata._is_shared_kv_draft_view = uses_shared_kv_cache
+            draft_metadata._is_separate_kv_draft_view = (
+                not uses_shared_kv_cache)
+            draft_metadata._draft_metadata = None
             draft_metadata.workspace_buffer = self.workspace_buffer
             draft_metadata.cuda_graph_buffers = None
             draft_metadata.cross = None
-            draft_metadata.max_num_tokens = self.max_num_requests
+            if uses_shared_kv_cache:
+                draft_metadata.max_num_tokens = self.max_num_requests
+            else:
+                draft_metadata.kv_cache_manager = draft_kv_cache_manager
             draft_metadata._seq_lens = None
             draft_metadata._seq_lens_cuda = None
             draft_metadata._seq_lens_kv = None
             draft_metadata._seq_lens_kv_cuda = None
             draft_metadata._saved_tensors = {}
-            draft_metadata.seq_lens = torch.ones((self.max_num_requests, ),
-                                                 dtype=torch.int)
+            if uses_shared_kv_cache:
+                draft_metadata.seq_lens = torch.ones((self.max_num_requests, ),
+                                                     dtype=torch.int)
+                draft_metadata.num_contexts = 0
+            else:
+                draft_metadata.seq_lens = self.seq_lens
+                draft_metadata.num_contexts = self.num_contexts
             draft_metadata.seq_lens_kv = None
-            draft_metadata.num_contexts = 0
             draft_metadata.__post_init__()
-            self._shared_kv_draft_metadata = draft_metadata
-            draft_metadata._sync_shared_kv_draft_view(self)
-        return self._shared_kv_draft_metadata
+            if not uses_shared_kv_cache:
+                # Keep the existing speculative worker's backend-neutral
+                # kv_lens_cuda protocol scoped to the separate-KV draft view.
+                draft_metadata.kv_lens_cuda = (
+                    draft_metadata._draft_kv_runtime_lens)
+            self._draft_metadata = draft_metadata
+            draft_metadata._sync_draft_view(self)
+        elif (self._draft_metadata._is_shared_kv_draft_view
+              != uses_shared_kv_cache):
+            raise RuntimeError(
+                "FlashInfer metadata cannot mix shared and separate draft KV "
+                "views.")
+        return self._draft_metadata
 
-    def _sync_shared_kv_draft_view(
-            self, target: "FlashInferAttentionMetadata") -> None:
-        """Refresh host-planned page tables before target graph replay."""
-        if not self._is_shared_kv_draft_view:
-            raise RuntimeError("Only a shared-KV draft metadata view can sync")
+    def _sync_draft_view(self, target: "FlashInferAttentionMetadata") -> None:
+        """Refresh a shared- or separate-KV draft metadata view."""
+        if not (self._is_shared_kv_draft_view
+                or self._is_separate_kv_draft_view):
+            raise RuntimeError("Only a draft metadata view can be synchronized")
 
         num_seqs = target.num_seqs
         self.request_ids = target.request_ids
         self.prompt_lens = target.prompt_lens
         self.kv_cache_params = target.kv_cache_params
-        self.kv_cache_manager = target.kv_cache_manager
         self.all_rank_num_tokens = target.all_rank_num_tokens
+
+        if self._is_separate_kv_draft_view:
+            self.padded_num_tokens = target.padded_num_tokens
+            self.seq_lens = target.seq_lens
+            self.seq_lens_kv = None
+            self.num_contexts = target.num_contexts
+            self._uses_full_draft_page_table = False
+            self.prepare()
+            torch.add(
+                target._cached_token_lens[:num_seqs],
+                target.seq_lens_kv_cuda[:num_seqs],
+                out=self._draft_kv_runtime_lens[:num_seqs],
+            )
+            if self.num_contexts == 0:
+                self._prepare_full_draft_page_table()
+            return
+
+        self.kv_cache_manager = target.kv_cache_manager
         self.seq_lens = torch.ones((num_seqs, ), dtype=torch.int)
         self.seq_lens_kv = None
         self.num_contexts = 0
@@ -685,7 +731,7 @@ def _sync_shared_kv_draft_view(
         self._qo_indptr[:num_seqs + 1].copy_(host_qo_indptr, non_blocking=True)
         full_kv_lens = (target._cached_token_lens[:num_seqs] +
                         target.seq_lens_kv_cuda[:num_seqs])
-        self._shared_kv_runtime_lens[:num_seqs].copy_(full_kv_lens)
+        self._draft_kv_runtime_lens[:num_seqs].copy_(full_kv_lens)
 
         if self.is_cuda_graph:
             # Graph replay keeps the captured trtllm-gen plan. Refresh only
@@ -708,7 +754,7 @@ def update_shared_kv_draft_lengths(
                 "Draft KV lengths can only be set on the shared-KV view")
         num_seqs = target.num_seqs
         cached_lens = target._cached_token_lens[:num_seqs]
-        runtime_lens = self._shared_kv_runtime_lens[:num_seqs]
+        runtime_lens = self._draft_kv_runtime_lens[:num_seqs]
         if num_contexts > 0:
             runtime_lens[:num_contexts].copy_(
                 cached_lens[:num_contexts] +
@@ -716,11 +762,72 @@ def update_shared_kv_draft_lengths(
         runtime_lens[num_contexts:num_seqs].copy_(
             cached_lens[num_contexts:num_seqs] +
             num_accepted_tokens[num_contexts:num_seqs])
+        self._update_draft_kv_lengths()
+
+    def _prepare_full_draft_page_table(self) -> None:
+        """Expose every allocated draft page and use device KV lengths."""
+        if self._uses_full_draft_page_table:
+            return
+        assert self.request_ids is not None
+        block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices(
+            self.request_ids)
+        self.num_blocks = [len(block_ids) for block_ids in block_ids_per_seq]
+        self.num_context_blocks = 0
+        self.num_generation_blocks = sum(self.num_blocks)
+        paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat(
+            self.request_ids, self.num_blocks)
+        self._paged_kv_indices[:paged_kv_indices.numel()].copy_(
+            paged_kv_indices, non_blocking=True)
+        self._host_paged_kv_indices = paged_kv_indices
+
+        host_indptr = maybe_pin_memory(
+            torch.from_numpy(
+                np.concatenate([[0], np.cumsum(self.num_blocks)
+                                ]).astype(np.int32, copy=False)))
+        num_seqs = self.num_seqs
+        self.paged_kv_indptr_decode[:num_seqs + 1].copy_(host_indptr,
+                                                         non_blocking=True)
+        self._host_paged_kv_indptr_decode = host_indptr
+        self.paged_kv_indptr = self.paged_kv_indptr_decode[:num_seqs + 1]
+        self._uses_full_draft_page_table = True
+        self._update_draft_kv_lengths()
+
+        if self.is_cuda_graph:
+            for plan_params, wrappers in self._plan_params_to_wrappers.items():
+                self._build_decode_block_tables(plan_params, wrappers)
+        elif not torch.cuda.is_current_stream_capturing():
+            self._clean_cached_plans(defer_plan=False)
+
+    def _update_draft_kv_lengths(self) -> None:
+        """Publish runtime KV lengths to a shared or separate draft view."""
+        num_seqs = self.num_seqs
+        runtime_lens = self._draft_kv_runtime_lens[:num_seqs]
         self._cached_token_lens[:num_seqs].copy_(runtime_lens)
-        torch.remainder(runtime_lens - 1,
+        self._paged_kv_last_page_len[:num_seqs].copy_(runtime_lens)
+        self._paged_kv_last_page_len[:num_seqs].sub_(1)
+        torch.remainder(self._paged_kv_last_page_len[:num_seqs],
                         self.page_size,
                         out=self._paged_kv_last_page_len[:num_seqs])
         self._paged_kv_last_page_len[:num_seqs].add_(1)
+        if self._is_shared_kv_draft_view:
+            return
+
+        self._cached_token_lens[:num_seqs].sub_(1)
+        self._positions[:num_seqs].copy_(self._cached_token_lens[:num_seqs])
+        torch.arange(num_seqs + 1,
+                     dtype=torch.int32,
+                     device=self._qo_indptr.device,
+                     out=self._qo_indptr[:num_seqs + 1])
+        torch.arange(num_seqs,
+                     dtype=torch.int32,
+                     device=self._batch_indices.device,
+                     out=self._batch_indices[:num_seqs])
+
+    def update_for_spec_dec(self) -> None:
+        if not self._is_separate_kv_draft_view:
+            return
+        self._prepare_full_draft_page_table()
+        self._update_draft_kv_lengths()
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -776,9 +883,13 @@ def _post_init_with_buffers(self, buffers) -> None:
         self._cached_token_lens = torch.empty((self.max_num_requests, ),
                                               dtype=torch.int,
                                               device='cuda')
-        self._shared_kv_runtime_lens = torch.empty((self.max_num_requests, ),
-                                                   dtype=torch.int,
-                                                   device='cuda')
+        self._draft_kv_runtime_lens = self.get_empty(
+            buffers,
+            (self.max_num_requests, ),
+            dtype=torch.int,
+            cache_name="_draft_kv_runtime_lens",
+            capture_graph=capture_graph,
+        )
         self._batch_indices = torch.empty((self.max_num_tokens, ),
                                           dtype=torch.int,
                                           device='cuda')
@@ -905,7 +1016,7 @@ def create_cuda_graph_metadata(self,
         metadata.max_num_tokens = max_batch_size * (1 + max_draft_tokens)
         # The graph owns distinct assistant wrappers and stable metadata
         # buffers; never inherit the eager view through the shallow copy.
-        metadata._shared_kv_draft_metadata = None
+        metadata._draft_metadata = None
         # Post init again to make sure all tensors are allocated
         metadata.__post_init__()
         return metadata
@@ -1417,7 +1528,8 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
             if pp.attention_mask_data is None
         ]
         defer_plan = len(active_wrappers) > 1
-        self._clean_cached_plans(defer_plan=defer_plan)
+        if not (self._is_separate_kv_draft_view and self.is_cuda_graph):
+            self._clean_cached_plans(defer_plan=defer_plan)
 
         # Re-plan MLA wrappers outside of forward/capture using the params
         # cached by prior warmup forwards. Forward still handles first-use or
@@ -1549,8 +1661,9 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor:
                     if self.num_generations < batch_size:
                         kv_lens_buf[self.num_generations:batch_size].zero_()
         if (not self._is_shared_kv_draft_view
-                and self._shared_kv_draft_metadata is not None):
-            self._shared_kv_draft_metadata._sync_shared_kv_draft_view(self)
+                and not self._is_separate_kv_draft_view
+                and self._draft_metadata is not None):
+            self._draft_metadata._sync_draft_view(self)
 
         if self.cross is not None and self.cross is not self:
             self.cross.prepare()
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 03ca7ac061c1..8cc7c8c0730d 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -451,13 +451,6 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
-    if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
-            and spec_config.spec_dec_mode.use_one_engine()
-            and not spec_config._use_shared_kv_cache):
-        raise ValueError(
-            "FLASHINFER attention backend supports one-engine speculative "
-            "decoding only when the draft model shares the target KV cache.")
-
     if mm_encoder_only:
         llm_args.mm_encoder_only = True
         llm_args.disable_overlap_scheduler = True
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 58fcee144e06..c55dad080e79 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -850,7 +850,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
             # Accepted counts let the indexer stash each gen's last-accepted row.
             attn_metadata.set_mtp_num_accepted(num_accepted_tokens)
 
-        with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager):
+        with self.draft_kv_cache_context(
+                attn_metadata, draft_kv_cache_manager) as draft_attn_metadata:
+            attn_metadata = draft_attn_metadata
+            inputs["attn_metadata"] = draft_attn_metadata
             for i in range(runtime_draft_len):
                 if reuse_mtp_topk:
                     attn_metadata.set_skip_topk(i > 0)
@@ -1013,8 +1016,10 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                     has_kv_cache = inputs[
                         "attn_metadata"].kv_cache_manager is not None
                     if has_kv_cache:
-                        attn_metadata.host_request_types[:attn_metadata.
-                                                         num_contexts].fill_(1)
+                        if hasattr(attn_metadata, "host_request_types"):
+                            attn_metadata.host_request_types[:attn_metadata.
+                                                             num_contexts].fill_(
+                                                                 1)
                         attn_metadata.num_contexts = 0
                     if hasattr(attn_metadata, 'kv_lens_cuda'):
                         attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
@@ -1409,7 +1414,7 @@ def _forward_external_shared_target_kv(
             batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
         )
 
-        draft_metadata = attn_metadata.get_shared_kv_draft_metadata()
+        draft_metadata = attn_metadata.get_draft_metadata()
         draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
                                                       num_accepted_tokens,
                                                       num_contexts)
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 9c9d01271956..bdb122ac4457 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -2137,20 +2137,25 @@ def get_draft_kv_cache_manager(self, resource_manager):
     @contextmanager
     def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
         """
-        Context manager to temporarily switch to draft KV cache manager in one-engine speculative decoding.
+        Select draft attention metadata for one-engine speculative decoding.
 
-        This swaps both the kv_cache_manager reference AND the block offset tensors,
-        since the target and draft KV caches have different block layouts.
+        TRTLLM metadata temporarily swaps its manager and block offsets.
+        FlashInfer uses an independently planned metadata view because its page
+        tables and kernel wrappers are manager-specific.
         """
 
         # draft_kv_cache_manager is None if using two-engine speculative decoding or not enabling separate draft KV cache.
         if draft_kv_cache_manager is None:
-            yield
+            yield attn_metadata
+            return
+
+        from ..attention_backend.flashinfer import FlashInferAttentionMetadata
+        if isinstance(attn_metadata, FlashInferAttentionMetadata):
+            yield attn_metadata.get_draft_metadata(draft_kv_cache_manager)
             return
 
-        # Only TrtllmAttentionMetadata supports separate draft KV cache layouts
         if not isinstance(attn_metadata, TrtllmAttentionMetadata):
-            yield
+            yield attn_metadata
             return
 
         # Check if draft KV cache block offsets are allocated
@@ -2158,7 +2163,7 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
                                       'draft_kv_cache_block_offsets', None)
         if draft_block_offsets is None:
             # Draft KV cache block offsets not allocated, skip switching
-            yield
+            yield attn_metadata
             return
 
         # Save main KV cache manager and block offsets
@@ -2174,7 +2179,7 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
             attn_metadata.prepare_flash_mla()
 
         try:
-            yield
+            yield attn_metadata
         finally:
             # Restore main KV cache manager and block offsets
             attn_metadata.kv_cache_manager = target_kv_cache_manager
diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py
index e512ca0196f7..9ce210ef28c9 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_attention.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py
@@ -66,6 +66,91 @@ class CUDAGraphTestScenario:
 
 class TestFlashInferAttention(unittest.TestCase):
 
+    def test_separate_kv_draft_metadata_uses_draft_manager(self):
+        if not torch.cuda.is_available():
+            self.skipTest("CUDA is required for FlashInfer metadata")
+
+        def create_manager():
+            return KVCacheManager(
+                KvCacheConfig(max_tokens=256),
+                tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+                num_layers=1,
+                num_kv_heads=1,
+                head_dim=128,
+                tokens_per_block=32,
+                max_seq_len=64,
+                max_batch_size=2,
+                mapping=Mapping(world_size=1, tp_size=1, rank=0),
+                dtype=tensorrt_llm.bindings.DataType.BF16,
+            )
+
+        target_manager = create_manager()
+        draft_manager = create_manager()
+        try:
+            target_manager.add_dummy_requests([0, 1], [31, 45], is_gen=True)
+            draft_manager.add_dummy_requests([0, 1], [31, 45],
+                                             is_gen=True,
+                                             max_num_draft_tokens=3)
+            metadata = FlashInferAttentionMetadata(
+                seq_lens=torch.ones(2, dtype=torch.int32),
+                num_contexts=0,
+                kv_cache_params=KVCacheParams(
+                    use_cache=True,
+                    num_cached_tokens_per_seq=[30, 44],
+                ),
+                max_num_requests=2,
+                max_num_tokens=8,
+                kv_cache_manager=target_manager,
+                request_ids=[0, 1],
+                is_cuda_graph=True,
+            )
+            metadata.prepare()
+            draft_metadata = metadata.get_draft_metadata(draft_manager)
+
+            self.assertIs(draft_metadata.kv_cache_manager, draft_manager)
+            self.assertFalse(hasattr(metadata, "kv_lens_cuda"))
+            torch.testing.assert_close(
+                draft_metadata.kv_lens_cuda[:2],
+                torch.tensor([31, 45], dtype=torch.int32, device="cuda"),
+            )
+            draft_blocks = draft_manager.get_batch_cache_indices([0, 1])
+            self.assertEqual(draft_metadata.num_blocks,
+                             list(map(len, draft_blocks)))
+
+            layer = FlashInferAttention(
+                layer_idx=0,
+                num_heads=1,
+                num_kv_heads=1,
+                head_dim=128,
+                flashinfer_backend="trtllm-gen",
+            )
+            q = torch.randn(2, 128, dtype=torch.bfloat16, device="cuda")
+            k = torch.randn_like(q)
+            v = torch.randn_like(q)
+            self.assertEqual(
+                layer.forward(q, k, v, draft_metadata).shape, q.shape)
+            for wrappers in draft_metadata._plan_params_to_wrappers.values():
+                torch.testing.assert_close(
+                    wrappers.decode_wrapper._kv_lens_buffer[:2],
+                    draft_metadata.kv_lens_cuda[:2],
+                )
+
+            with mock.patch.object(
+                    draft_metadata,
+                    "_plan_with_params",
+                    wraps=draft_metadata._plan_with_params,
+            ) as replan, mock.patch.object(
+                    draft_metadata,
+                    "_build_decode_block_tables",
+            ) as refresh_block_tables:
+                metadata.prepare()
+            replan.assert_not_called()
+            self.assertEqual(refresh_block_tables.call_count,
+                             len(draft_metadata._plan_params_to_wrappers))
+        finally:
+            target_manager.shutdown()
+            draft_manager.shutdown()
+
     def test_ragged_no_kv_cuda_graph_uses_stable_indptr_aliases(self):
         if not torch.cuda.is_available():
             self.skipTest("CUDA is required for FlashInfer metadata")
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 479c29dc4ee5..63a1c36d10d4 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -2631,7 +2631,7 @@ def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) ->
         }
 
         accepted_tokens = torch.tensor([1, 3], dtype=torch.int, device="cuda")
-        draft_metadata = metadata.get_shared_kv_draft_metadata()
+        draft_metadata = metadata.get_draft_metadata()
         draft_metadata.update_shared_kv_draft_lengths(
             metadata,
             accepted_tokens,
@@ -2645,7 +2645,7 @@ def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) ->
         self.assertIs(draft_metadata.kv_cache_manager, metadata.kv_cache_manager)
         self.assertIsNot(draft_metadata, metadata)
         torch.testing.assert_close(
-            draft_metadata._shared_kv_runtime_lens[:2],
+            draft_metadata._draft_kv_runtime_lens[:2],
             expected_kv_lens,
             atol=0,
             rtol=0,

From 2e8c4debe3579fa0871de5bd1fc9b69b5c445819 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 28 Jul 2026 07:18:25 +0000
Subject: [PATCH 20/26] [None][refactor] simplify Gemma4 shared-KV MTP
 integration

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_loader.py         |  14 +-
 tensorrt_llm/_torch/speculative/eagle3.py     | 368 ++++++++----------
 tensorrt_llm/_torch/speculative/utils.py      |  10 +
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 107 ++++-
 4 files changed, 257 insertions(+), 242 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index 15d175ff8d53..2a6ea23076aa 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -21,7 +21,7 @@
 from tensorrt_llm._utils import str_dtype_to_torch
 from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig,
                                           ExecutorMemoryType,
-                                          ModelExpressConfig, MTPDecodingConfig,
+                                          ModelExpressConfig,
                                           SparseAttentionConfig, TorchLlmArgs)
 from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto,
                                            _resolve_transceiver_runtime_auto,
@@ -52,10 +52,6 @@
     "auto": "auto"
 }
 _VALID_KV_CACHE_DTYPES = ("fp8", "nvfp4", "auto")
-_GEMMA4_TARGET_ARCHITECTURES = (
-    "Gemma4ForCausalLM",
-    "Gemma4ForConditionalGeneration",
-)
 
 
 def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig,
@@ -444,14 +440,6 @@ def load_config_and_apply_defaults(
             # init=False runtime fields such as num_nextn_predict_layers.
             update_spec_config_from_model_config(llm_args.speculative_config,
                                                  config.pretrained_config)
-            architectures = getattr(config.pretrained_config, "architectures",
-                                    None) or ()
-            if (isinstance(llm_args.speculative_config, MTPDecodingConfig)
-                    and architectures
-                    and architectures[0] in _GEMMA4_TARGET_ARCHITECTURES):
-                llm_args.speculative_config._use_shared_kv_cache = (
-                    llm_args.speculative_config.spec_dec_mode.
-                    is_mtp_eagle_one_model())
 
         # The transceiver preference follows the checkpoint's original
         # architecture: _resolve_class may rewrite it to an execution class
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index c55dad080e79..9b9132e5384c 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -645,6 +645,8 @@ def __init__(self,
         if getattr(spec_config, 'sa_config', None) is not None:
             self.sa_enhancer = SADraftEnhancer(spec_config.sa_config.threshold)
         self.use_dynamic_tree = getattr(spec_config, 'use_dynamic_tree', False)
+        self._uses_external_shared_target_kv = (
+            spec_config._use_shared_kv_cache)
         self.spec_tree_manager = None
 
         # MTP Eagle: lazily-resolved flag for Mamba hybrid cache support
@@ -770,38 +772,53 @@ def _forward_impl(self,
                 max_draft_len=runtime_draft_len,
             )
 
-        # Save the old attn_metadata and spec_metadata
-        self._prepare_attn_metadata_for_spec_dec(attn_metadata)
-
-        # Prepare inputs for the 1st draft model forward
-        position_ids = position_ids.squeeze(0)
-        inputs = self.prepare_1st_drafter_inputs(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            hidden_states=hidden_states,
-            accepted_tokens=accepted_tokens,
-            attn_metadata=attn_metadata,
-            spec_metadata=spec_metadata,
-            draft_model=draft_model)
-
-        # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here
-        # so the post-loop restore (below) can put attn_metadata back into a
-        # state the target model expects.
-        original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
-
-        # Get the draft KV cache manager if using separate layouts
-        draft_kv_cache_manager = self.get_draft_kv_cache_manager(
-            resource_manager)
-
-        next_draft_tokens = self._forward_draft_loop(
-            inputs, attn_metadata, spec_metadata, draft_model,
-            draft_kv_cache_manager, num_contexts, num_gens, batch_size,
-            num_accepted_tokens, original_all_rank_num_tokens, resource_manager)
-        # restore attn_metadata to support cuda graph
-        self._restore_attn_metadata_from_spec_dec(attn_metadata)
-        # restore all_rank_num_tokens for attention DP
-        if original_all_rank_num_tokens is not None:
-            attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
+        if self._uses_external_shared_target_kv:
+            next_draft_tokens = (
+                self._forward_external_shared_target_kv_draft_loop(
+                    position_ids=position_ids,
+                    hidden_states=hidden_states,
+                    attn_metadata=attn_metadata,
+                    spec_metadata=spec_metadata,
+                    draft_model=draft_model,
+                    accepted_tokens=accepted_tokens,
+                    num_accepted_tokens=num_accepted_tokens,
+                    num_contexts=num_contexts,
+                    batch_size=batch_size,
+                ))
+        else:
+            # Save the old attn_metadata and spec_metadata
+            self._prepare_attn_metadata_for_spec_dec(attn_metadata)
+
+            # Prepare inputs for the 1st draft model forward
+            position_ids = position_ids.squeeze(0)
+            inputs = self.prepare_1st_drafter_inputs(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                accepted_tokens=accepted_tokens,
+                attn_metadata=attn_metadata,
+                spec_metadata=spec_metadata,
+                draft_model=draft_model)
+
+            # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here
+            # so the post-loop restore (below) can put attn_metadata back into a
+            # state the target model expects.
+            original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
+
+            # Get the draft KV cache manager if using separate layouts
+            draft_kv_cache_manager = self.get_draft_kv_cache_manager(
+                resource_manager)
+
+            next_draft_tokens = self._forward_draft_loop(
+                inputs, attn_metadata, spec_metadata, draft_model,
+                draft_kv_cache_manager, num_contexts, num_gens, batch_size,
+                num_accepted_tokens, original_all_rank_num_tokens,
+                resource_manager)
+            # restore attn_metadata to support cuda graph
+            self._restore_attn_metadata_from_spec_dec(attn_metadata)
+            # restore all_rank_num_tokens for attention DP
+            if original_all_rank_num_tokens is not None:
+                attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens
 
         # prepare next new tokens to support overlap scheduler
         next_new_tokens = self._prepare_next_new_tokens(
@@ -818,6 +835,86 @@ def _forward_impl(self,
             'next_new_tokens': next_new_tokens,
         }
 
+    def _forward_external_shared_target_kv_draft_loop(
+        self,
+        *,
+        position_ids,
+        hidden_states,
+        attn_metadata,
+        spec_metadata,
+        draft_model,
+        accepted_tokens,
+        num_accepted_tokens,
+        num_contexts,
+        batch_size,
+    ):
+        """Draft with an external Q-only assistant over accepted target KV."""
+        if not isinstance(attn_metadata, FlashInferAttentionMetadata):
+            raise TypeError(
+                "External shared-target-KV MTP requires FlashInfer attention "
+                "metadata.")
+
+        (
+            draft_input_ids,
+            recurrent_hidden_states,
+            draft_position_ids,
+        ) = self._prepare_external_shared_target_kv_draft_inputs(
+            accepted_tokens=accepted_tokens,
+            num_accepted_tokens=num_accepted_tokens,
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            sequence_lengths=attn_metadata.seq_lens_cuda[:batch_size],
+            num_contexts=num_contexts,
+            batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
+        )
+
+        draft_metadata = attn_metadata.get_draft_metadata()
+        draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
+                                                      num_accepted_tokens,
+                                                      num_contexts)
+        draft_metadata.use_spec_decoding = False
+        draft_metadata.padded_num_tokens = None
+        draft_metadata.all_rank_num_tokens = (
+            spec_metadata.subseq_all_rank_num_tokens)
+
+        next_draft_tokens = []
+        for draft_step in range(spec_metadata.runtime_draft_len):
+            if self.guided_decoder is not None:
+                self.guided_decoder.add_draft_batch(
+                    draft_input_ids,
+                    num_accepted_tokens,
+                    draft_step=draft_step,
+                )
+
+            draft_logits, recurrent_hidden_states = (
+                draft_model.forward_draft_step(
+                    input_ids=draft_input_ids,
+                    position_ids=draft_position_ids,
+                    recurrent_hidden_states=recurrent_hidden_states,
+                    attn_metadata=draft_metadata,
+                    spec_metadata=spec_metadata,
+                ))
+            if self.guided_decoder is not None:
+                self.guided_decoder.execute_draft_batch(
+                    draft_logits,
+                    draft_step=draft_step,
+                )
+            draft_input_ids = self.sample_draft_tokens(
+                draft_logits,
+                spec_metadata,
+                batch_size,
+                draft_step=draft_step,
+            )
+            next_draft_tokens.append(draft_input_ids)
+
+        next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
+        if self.sa_enhancer is not None:
+            gen_draft_tokens = next_draft_tokens[num_contexts:]
+            gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens(
+                gen_draft_tokens)
+            next_draft_tokens[num_contexts:] = gen_draft_tokens
+        return next_draft_tokens
+
     def _forward_draft_loop(self, inputs, attn_metadata, spec_metadata,
                             draft_model, draft_kv_cache_manager, num_contexts,
                             num_gens, batch_size, num_accepted_tokens,
@@ -1144,7 +1241,6 @@ def sample_and_accept_draft_tokens(
         logits: torch.Tensor,
         attn_metadata: AttentionMetadata,
         spec_metadata: Eagle3OneModelSpecMetadata,
-        num_contexts: Optional[int] = None,
     ):
         """Sample the golden token and verify previously proposed draft tokens.
 
@@ -1152,8 +1248,7 @@ def sample_and_accept_draft_tokens(
         acceptance is enabled (both Eagle3 and MTP Eagle); ignored otherwise.
         """
         batch_size = attn_metadata.num_seqs
-        if num_contexts is None:
-            num_contexts = attn_metadata.num_contexts
+        num_contexts = attn_metadata.num_contexts
         num_gens = batch_size - num_contexts
 
         runtime_draft_len = spec_metadata.runtime_draft_len
@@ -1279,185 +1374,6 @@ def prepare_1st_drafter_inputs(
             "spec_metadata": spec_metadata,
         }
 
-
-class MTPEagleWorker(Eagle3OneModelWorker):
-    """MTP worker built on the shared Eagle3 one-model drafting loop."""
-
-    def __init__(self,
-                 spec_config,
-                 model_config: Optional[ModelConfig] = None,
-                 use_separate_draft_kv_cache: bool = False,
-                 *,
-                 mapping: Optional[Mapping] = None):
-        super().__init__(
-            spec_config,
-            mapping=mapping,
-            model_config=model_config,
-            use_separate_draft_kv_cache=use_separate_draft_kv_cache)
-        # Preserved for callers/tests that still expect this attribute.
-        self.is_thop = False
-        self._uses_external_shared_target_kv = False
-
-    def set_draft_model(self, draft_model) -> None:
-        super().set_draft_model(draft_model)
-        expects_external_shared_target_kv = (
-            self.spec_config._use_shared_kv_cache)
-        supports_shared_target_kv = bool(
-            getattr(draft_model, "shares_target_kv_cache", False))
-        if expects_external_shared_target_kv and not supports_shared_target_kv:
-            raise ValueError(
-                "External shared-target-KV MTP requires a draft model that "
-                "declares shares_target_kv_cache=True.")
-        self._uses_external_shared_target_kv = (
-            expects_external_shared_target_kv and supports_shared_target_kv)
-        if not self._uses_external_shared_target_kv:
-            return
-        if self.use_dynamic_tree:
-            raise ValueError(
-                "External shared-target-KV MTP supports only the linear draft "
-                "path.")
-        if self.spec_config.draft_len_schedule is not None:
-            raise ValueError(
-                "External shared-target-KV MTP does not support a draft "
-                "length schedule.")
-        if self.sa_enhancer is not None:
-            raise ValueError(
-                "External shared-target-KV MTP does not support the suffix "
-                "automaton enhancer.")
-        if self.spec_config.use_rejection_sampling:
-            raise ValueError(
-                "External shared-target-KV MTP does not support rejection "
-                "sampling.")
-
-    def set_guided_decoder(self, guided_decoder) -> bool:
-        if self._uses_external_shared_target_kv:
-            raise ValueError(
-                "External shared-target-KV MTP does not support guided "
-                "decoding.")
-        return super().set_guided_decoder(guided_decoder)
-
-    def _forward_impl(self,
-                      input_ids,
-                      position_ids,
-                      hidden_states,
-                      logits,
-                      attn_metadata,
-                      spec_metadata,
-                      draft_model,
-                      resource_manager=None):
-        if not self._uses_external_shared_target_kv:
-            return super()._forward_impl(input_ids, position_ids, hidden_states,
-                                         logits, attn_metadata, spec_metadata,
-                                         draft_model, resource_manager)
-        return self._forward_external_shared_target_kv(input_ids, position_ids,
-                                                       hidden_states, logits,
-                                                       attn_metadata,
-                                                       spec_metadata,
-                                                       draft_model)
-
-    def _forward_external_shared_target_kv(
-        self,
-        input_ids,
-        position_ids,
-        hidden_states,
-        logits,
-        attn_metadata,
-        spec_metadata,
-        draft_model,
-    ):
-        """Draft with an external Q-only assistant over accepted target KV."""
-        if not isinstance(attn_metadata, FlashInferAttentionMetadata):
-            raise TypeError(
-                "External shared-target-KV MTP requires FlashInfer attention "
-                "metadata.")
-
-        batch_size = attn_metadata.num_seqs
-        num_contexts = batch_size - spec_metadata.num_generations
-        runtime_draft_len = spec_metadata.runtime_draft_len
-        if runtime_draft_len == 0:
-            target_tokens = self._sample_tokens_for_batch(
-                logits, spec_metadata, num_contexts, batch_size)
-            accepted_tokens = target_tokens.unsqueeze(1)
-            num_accepted_tokens = torch.ones(batch_size,
-                                             dtype=torch.int,
-                                             device=logits.device)
-            next_draft_tokens = torch.empty((batch_size, 0),
-                                            dtype=torch.int32,
-                                            device=logits.device)
-            return {
-                "logits": logits,
-                "new_tokens": accepted_tokens,
-                "new_tokens_lens": num_accepted_tokens,
-                "next_draft_tokens": next_draft_tokens,
-                "next_new_tokens": accepted_tokens,
-            }
-
-        accepted_tokens, num_accepted_tokens = (
-            self.sample_and_accept_draft_tokens(
-                input_ids,
-                logits,
-                attn_metadata,
-                spec_metadata,
-                num_contexts=num_contexts,
-            ))
-        (
-            draft_input_ids,
-            recurrent_hidden_states,
-            draft_position_ids,
-        ) = self._prepare_external_shared_target_kv_draft_inputs(
-            accepted_tokens=accepted_tokens,
-            num_accepted_tokens=num_accepted_tokens,
-            hidden_states=hidden_states,
-            position_ids=position_ids,
-            sequence_lengths=attn_metadata.seq_lens_cuda[:batch_size],
-            num_contexts=num_contexts,
-            batch_indices=spec_metadata.batch_indices_cuda[:batch_size],
-        )
-
-        draft_metadata = attn_metadata.get_draft_metadata()
-        draft_metadata.update_shared_kv_draft_lengths(attn_metadata,
-                                                      num_accepted_tokens,
-                                                      num_contexts)
-        draft_metadata.use_spec_decoding = False
-        draft_metadata.padded_num_tokens = None
-        draft_metadata.all_rank_num_tokens = (
-            spec_metadata.subseq_all_rank_num_tokens)
-
-        next_draft_tokens = []
-        for draft_step in range(runtime_draft_len):
-            draft_logits, recurrent_hidden_states = (
-                draft_model.forward_draft_step(
-                    input_ids=draft_input_ids,
-                    position_ids=draft_position_ids,
-                    recurrent_hidden_states=recurrent_hidden_states,
-                    attn_metadata=draft_metadata,
-                    spec_metadata=spec_metadata,
-                ))
-            draft_input_ids = self.sample_draft_tokens(
-                draft_logits,
-                spec_metadata,
-                batch_size,
-                draft_step=draft_step,
-            )
-            next_draft_tokens.append(draft_input_ids)
-
-        next_draft_tokens = torch.stack(next_draft_tokens, dim=1)
-        next_new_tokens = self._prepare_next_new_tokens(
-            accepted_tokens,
-            next_draft_tokens,
-            spec_metadata.batch_indices_cuda[:batch_size],
-            batch_size,
-            num_accepted_tokens,
-        )
-        attn_metadata.use_spec_decoding = True
-        return {
-            "logits": logits,
-            "new_tokens": accepted_tokens,
-            "new_tokens_lens": num_accepted_tokens,
-            "next_draft_tokens": next_draft_tokens,
-            "next_new_tokens": next_new_tokens,
-        }
-
     @staticmethod
     def _prepare_external_shared_target_kv_draft_inputs(
         *,
@@ -1482,3 +1398,27 @@ def _prepare_external_shared_target_kv_draft_inputs(
         draft_position_ids = (
             _select_mtp_position_ids(position_ids, recurrent_indices) + 1)
         return draft_input_ids, recurrent_hidden_states, draft_position_ids
+
+
+class MTPEagleWorker(Eagle3OneModelWorker):
+    """Backward-compatible alias for ``Eagle3OneModelWorker`` in MTP Eagle mode.
+
+    The constructor matches the historical positional signature
+    ``(spec_config, model_config, use_separate_draft_kv_cache)`` so callers
+    that import ``MTPEagleWorker`` from ``mtp.py`` or instantiate it directly
+    keep working. All logic is inherited from :class:`Eagle3OneModelWorker`.
+    """
+
+    def __init__(self,
+                 spec_config,
+                 model_config: Optional[ModelConfig] = None,
+                 use_separate_draft_kv_cache: bool = False,
+                 *,
+                 mapping: Optional[Mapping] = None):
+        super().__init__(
+            spec_config,
+            mapping=mapping,
+            model_config=model_config,
+            use_separate_draft_kv_cache=use_separate_draft_kv_cache)
+        # Preserved for callers/tests that still expect this attribute.
+        self.is_thop = False
diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 78aca6b8c498..887d2bd7cda5 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -37,6 +37,11 @@
                                 SaveHiddenStatesSpecMetadata)
 from .suffix_automaton import SuffixAutomatonManager
 
+_GEMMA4_SHARED_KV_TARGET_ARCHITECTURES = (
+    "Gemma4ForCausalLM",
+    "Gemma4ForConditionalGeneration",
+)
+
 
 def _is_effective_dynamic_tree(spec_config) -> bool:
     # At dynamic_tree_max_topK == 1 the tree collapses to a linear chain; route
@@ -547,6 +552,11 @@ def update_spec_config_from_model_config(spec_config, model_config):
     from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig
     if not isinstance(spec_config, MTPDecodingConfig):
         return
+    architectures = getattr(model_config, "architectures", None) or ()
+    if (architectures
+            and architectures[0] in _GEMMA4_SHARED_KV_TARGET_ARCHITECTURES):
+        spec_config._use_shared_kv_cache = (
+            spec_config.spec_dec_mode.is_mtp_eagle_one_model())
     # Read the MTP layer count from the model's pretrained config. This
     # determines the actual MTP layer count in the checkpoint and drives the
     # spec_dec_mode decision (EAGLE vs vanilla MTP). Different checkpoints expose
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 c23f8d227cb3..120a8df3618b 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
@@ -2,12 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
 
 from tensorrt_llm._torch.models import modeling_speculative
 from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
+from tensorrt_llm._torch.speculative import eagle3
 from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
 from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
 from tensorrt_llm._torch.speculative.utils import (
@@ -48,6 +50,26 @@ def test_external_checkpoint_does_not_imply_shared_kv_cache():
     assert should_use_separate_draft_kv_cache(spec_config)
 
 
+@pytest.mark.parametrize("one_model,expected", [(True, True), (False, False)])
+def test_gemma4_config_sets_shared_kv_cache_for_one_model_only(
+    one_model,
+    expected,
+):
+    spec_config = MTPDecodingConfig(
+        max_draft_len=3,
+        speculative_model="/tmp/gemma4-assistant",
+        mtp_eagle_one_model=one_model,
+    )
+    model_config = SimpleNamespace(
+        architectures=["Gemma4ForConditionalGeneration"],
+        num_nextn_predict_layers=1,
+    )
+
+    update_spec_config_from_model_config(spec_config, model_config)
+
+    assert spec_config._use_shared_kv_cache is expected
+
+
 def test_external_shared_kv_uses_no_draft_kv_cache():
     spec_config = _shared_kv_spec_config()
 
@@ -94,27 +116,82 @@ def test_shared_kv_alias_setup_rebinds_target_model():
     modeling_speculative.SpecDecOneEngineForCausalLM.setup_aliases(model)
 
 
-def test_external_shared_kv_worker_requires_draft_model_capability():
+def test_external_shared_kv_worker_uses_config_and_supports_guided_decoding():
     worker = MTPEagleWorker(_shared_kv_spec_config())
+    guided_decoder = object()
 
-    with pytest.raises(ValueError, match="shares_target_kv_cache=True"):
-        worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
+    assert worker._uses_external_shared_target_kv
+    assert worker.set_guided_decoder(guided_decoder)
+    assert worker.guided_decoder is guided_decoder
+    worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
 
 
-def test_external_shared_kv_worker_rejects_unverified_modes():
-    spec_config = _shared_kv_spec_config(
-        use_dynamic_tree=True,
-        dynamic_tree_max_topK=2,
+def test_external_shared_kv_draft_loop_applies_guided_decoding(monkeypatch):
+    draft_metadata = SimpleNamespace(
+        update_shared_kv_draft_lengths=Mock(),
     )
-    worker = MTPEagleWorker(spec_config)
 
-    with pytest.raises(ValueError, match="linear draft path"):
-        worker.set_draft_model(
-            SimpleNamespace(
-                model=SimpleNamespace(),
-                shares_target_kv_cache=True,
-            )
-        )
+    class FakeFlashInferAttentionMetadata:
+        def __init__(self):
+            self.seq_lens_cuda = torch.tensor([2, 2], dtype=torch.int32)
+
+        def get_draft_metadata(self):
+            return draft_metadata
+
+    monkeypatch.setattr(eagle3, "FlashInferAttentionMetadata", FakeFlashInferAttentionMetadata)
+    worker = MTPEagleWorker(_shared_kv_spec_config(max_draft_len=2))
+    guided_decoder = SimpleNamespace(
+        add_draft_batch=Mock(),
+        execute_draft_batch=Mock(),
+    )
+    worker.set_guided_decoder(guided_decoder)
+
+    sampled_tokens = [
+        torch.tensor([41, 42], dtype=torch.int32),
+        torch.tensor([51, 52], dtype=torch.int32),
+    ]
+    monkeypatch.setattr(
+        worker,
+        "sample_draft_tokens",
+        lambda *args, **kwargs: sampled_tokens.pop(0),
+    )
+    draft_model = SimpleNamespace(
+        forward_draft_step=lambda **kwargs: (torch.zeros(2, 4), kwargs["recurrent_hidden_states"])
+    )
+    attn_metadata = FakeFlashInferAttentionMetadata()
+    spec_metadata = SimpleNamespace(
+        batch_indices_cuda=torch.arange(2),
+        runtime_draft_len=2,
+        subseq_all_rank_num_tokens=None,
+    )
+    accepted_tokens = torch.tensor(
+        [[10, 11, 12], [20, 21, 22]],
+        dtype=torch.int32,
+    )
+    num_accepted_tokens = torch.ones(2, dtype=torch.long)
+
+    next_draft_tokens = worker._forward_external_shared_target_kv_draft_loop(
+        position_ids=torch.arange(4, dtype=torch.int32),
+        hidden_states=torch.arange(8, dtype=torch.float32).unsqueeze(1),
+        attn_metadata=attn_metadata,
+        spec_metadata=spec_metadata,
+        draft_model=draft_model,
+        accepted_tokens=accepted_tokens,
+        num_accepted_tokens=num_accepted_tokens,
+        num_contexts=1,
+        batch_size=2,
+    )
+
+    assert torch.equal(
+        next_draft_tokens,
+        torch.tensor([[41, 51], [42, 52]], dtype=torch.int32),
+    )
+    assert [
+        call.kwargs["draft_step"] for call in guided_decoder.add_draft_batch.call_args_list
+    ] == [0, 1]
+    assert [
+        call.kwargs["draft_step"] for call in guided_decoder.execute_draft_batch.call_args_list
+    ] == [0, 1]
 
 
 @pytest.mark.parametrize(

From 5ade312d44b9ec8a02a2ddcdc6cb58293637062a Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 28 Jul 2026 07:57:23 +0000
Subject: [PATCH 21/26] [None][refactor] simplify Gemma4 MTP helpers and tests

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/speculative/eagle3.py     |  14 +-
 .../_torch/modeling/test_gemma4_multimodal.py |  16 -
 .../_torch/modeling/test_modeling_gemma4.py   | 231 +-------------
 .../hw_agnostic/test_gemma4_drafting_loop.py  | 289 ------------------
 .../speculative/hw_agnostic/test_mtp.py       |  71 +++++
 5 files changed, 89 insertions(+), 532 deletions(-)
 delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py

diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 9b9132e5384c..d2c6e037a85c 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -858,7 +858,7 @@ def _forward_external_shared_target_kv_draft_loop(
             draft_input_ids,
             recurrent_hidden_states,
             draft_position_ids,
-        ) = self._prepare_external_shared_target_kv_draft_inputs(
+        ) = self._prepare_shared_kv_draft_inputs(
             accepted_tokens=accepted_tokens,
             num_accepted_tokens=num_accepted_tokens,
             hidden_states=hidden_states,
@@ -1112,11 +1112,11 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                     attn_metadata.on_update()
                     has_kv_cache = inputs[
                         "attn_metadata"].kv_cache_manager is not None
+                    if has_kv_cache and hasattr(attn_metadata,
+                                                "host_request_types"):
+                        attn_metadata.host_request_types[:attn_metadata.
+                                                         num_contexts].fill_(1)
                     if has_kv_cache:
-                        if hasattr(attn_metadata, "host_request_types"):
-                            attn_metadata.host_request_types[:attn_metadata.
-                                                             num_contexts].fill_(
-                                                                 1)
                         attn_metadata.num_contexts = 0
                     if hasattr(attn_metadata, 'kv_lens_cuda'):
                         attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= (
@@ -1374,8 +1374,8 @@ def prepare_1st_drafter_inputs(
             "spec_metadata": spec_metadata,
         }
 
-    @staticmethod
-    def _prepare_external_shared_target_kv_draft_inputs(
+    def _prepare_shared_kv_draft_inputs(
+        self,
         *,
         accepted_tokens: torch.Tensor,
         num_accepted_tokens: torch.Tensor,
diff --git a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
index 427ba1bed3d7..c51b328951ff 100644
--- a/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
+++ b/tests/unittest/_torch/modeling/test_gemma4_multimodal.py
@@ -774,22 +774,6 @@ def test_encoder_cache_reuses_image_embedding_across_requests(self):
         torch.testing.assert_close(second, first)
         self.assertEqual(len(model._multimodal_encoder_cache), 1)
 
-    def test_draft_weight_loading_contract_is_proxied_to_language_model(self):
-        model = self._make_model()
-
-        self.assertIs(model.model, model.llm.model)
-        self.assertIs(model.lm_head, model.llm.lm_head)
-        self.assertIs(model.epilogue, model.llm.epilogue)
-        self.assertIs(model.spec_worker, model.llm.spec_worker)
-        self.assertIs(model.draft_config, model.llm.draft_config)
-        self.assertIs(model.draft_model, model.llm.draft_model)
-
-        weights = {"draft": torch.ones(1)}
-        mapper = object()
-        with unittest.mock.patch.object(model.llm, "load_draft_weights") as loader:
-            model.load_draft_weights(weights, mapper)
-        loader.assert_called_once_with(weights, mapper)
-
     def test_chunked_prefill_reuses_cached_vision_embeddings(self):
         """Later active chunks slice cached features without rerunning vision."""
         model = self._make_model()
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 63a1c36d10d4..5471e0a6ae9f 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -566,17 +566,17 @@ def test_reject_unrecognized_expert_weights(self):
 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))
+    def test_assistant_config(self):
+        config_dict = deepcopy(GEMMA4_ASSISTANT_CONFIG)
+        config_dict["text_config"].pop("num_kv_shared_layers")
+        config = Gemma4AssistantConfig(**config_dict)
 
         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_config_auto_config_round_trip(self):
-        config = Gemma4AssistantConfig(**deepcopy(GEMMA4_ASSISTANT_CONFIG))
+        self.assertEqual(config.text_config.num_kv_shared_layers, 4)
 
         with tempfile.TemporaryDirectory() as directory:
             config.save_pretrained(directory)
@@ -586,18 +586,7 @@ def test_assistant_config_auto_config_round_trip(self):
         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):
+    def test_assistant_rejects_partial_kv_sharing(self):
         config_dict = deepcopy(GEMMA4_ASSISTANT_CONFIG)
         config_dict["text_config"]["num_kv_shared_layers"] = 2
 
@@ -885,7 +874,6 @@ def _build_gemma4_kv_cache_manager(
     num_blocks=4,
     tokens_per_block=32,
     batch_size=1,
-    force_vswa=False,
 ):
     """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads.
 
@@ -941,7 +929,7 @@ def _build_gemma4_kv_cache_manager(
     # exceeds sliding_window.
     sliding_window = getattr(config, "sliding_window", None)
     max_attn_window = None
-    needs_vswa = force_vswa or (isinstance(head_dim, list) and len(set(head_dim)) > 1)
+    needs_vswa = 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:
@@ -2452,11 +2440,9 @@ 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"],
@@ -2472,13 +2458,12 @@ def _make_trtllm_gen_decode_case(
         if manager_batch_size is None:
             manager_batch_size = batch_size
 
-        config = Gemma4TextConfig(**deepcopy(config_dict or GEMMA4_E2B_REAL_DIMS_CONFIG))
+        config = Gemma4TextConfig(**deepcopy(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")
@@ -2619,11 +2604,10 @@ def _expected_decode_block_table(
     @unittest.mock.patch(
         "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None
     )
-    def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) -> None:
-        """The assistant gets private decode state over immutable target KV."""
-        page_counts = [3, 2]
+    def test_shared_kv_draft_view(self) -> None:
+        """The draft view advances lengths without modifying target KV."""
         kv_cache_manager, layers, metadata, queries, _, _ = self._make_trtllm_gen_decode_case(
-            page_counts
+            [3, 2]
         )
         target_kv = {
             layer.layer_idx: kv_cache_manager.get_buffers(layer.layer_idx).clone()
@@ -2643,20 +2627,12 @@ def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) ->
             layer.forward(query, None, None, draft_metadata)
 
         self.assertIs(draft_metadata.kv_cache_manager, metadata.kv_cache_manager)
-        self.assertIsNot(draft_metadata, metadata)
         torch.testing.assert_close(
             draft_metadata._draft_kv_runtime_lens[:2],
             expected_kv_lens,
             atol=0,
             rtol=0,
         )
-        for wrappers in draft_metadata._plan_params_to_wrappers.values():
-            torch.testing.assert_close(
-                wrappers.decode_wrapper._kv_lens_buffer[:2],
-                expected_kv_lens,
-                atol=0,
-                rtol=0,
-            )
         for layer in layers:
             torch.testing.assert_close(
                 kv_cache_manager.get_buffers(layer.layer_idx),
@@ -2664,11 +2640,6 @@ def test_shared_kv_draft_view_uses_accepted_prefix_without_appending_kv(self) ->
                 atol=0,
                 rtol=0,
             )
-        with unittest.mock.patch.object(
-            draft_metadata, "_plan_with_params", wraps=draft_metadata._plan_with_params
-        ) as replan:
-            metadata.prepare()
-        replan.assert_not_called()
 
     @torch.no_grad()
     @unittest.mock.patch(
@@ -2757,47 +2728,6 @@ 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"]
-        page_counts = [5, 3]
-        _, _, metadata, _, _, _ = self._make_trtllm_gen_decode_case(
-            page_counts,
-            config_dict=config_dict,
-            force_vswa=True,
-        )
-
-        plan_params = list(metadata._plan_params_to_wrappers)
-        self.assertEqual(len(plan_params), 2)
-        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():
-            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
@@ -3050,145 +2980,6 @@ 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 multi-query decode survives graph 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=0,
-                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
deleted file mode 100644
index 120a8df3618b..000000000000
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py
+++ /dev/null
@@ -1,289 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-import pytest
-import torch
-
-from tensorrt_llm._torch.models import modeling_speculative
-from tensorrt_llm._torch.models.modeling_gemma4 import Gemma4ForCausalLM
-from tensorrt_llm._torch.speculative import eagle3
-from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
-from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
-from tensorrt_llm._torch.speculative.utils import (
-    get_num_extra_kv_tokens,
-    get_num_spec_layers,
-    update_spec_config_from_model_config,
-)
-from tensorrt_llm.llmapi import MTPDecodingConfig
-
-
-def _shared_kv_spec_config(**kwargs) -> MTPDecodingConfig:
-    spec_config = MTPDecodingConfig(
-        max_draft_len=kwargs.pop("max_draft_len", 3),
-        speculative_model="/tmp/gemma4-assistant",
-        mtp_eagle_one_model=True,
-        **kwargs,
-    )
-    spec_config._use_shared_kv_cache = True
-    return spec_config
-
-
-def test_external_checkpoint_does_not_imply_shared_kv_cache():
-    spec_config = MTPDecodingConfig(
-        max_draft_len=3,
-        speculative_model="/tmp/assistant",
-        mtp_eagle_one_model=True,
-    )
-    model_config = SimpleNamespace(
-        architectures=["LlamaForCausalLM"],
-        num_nextn_predict_layers=1,
-    )
-
-    update_spec_config_from_model_config(spec_config, model_config)
-
-    assert not spec_config._use_shared_kv_cache
-    assert get_num_spec_layers(spec_config) == 1
-    assert get_num_extra_kv_tokens(spec_config) == 2
-    assert should_use_separate_draft_kv_cache(spec_config)
-
-
-@pytest.mark.parametrize("one_model,expected", [(True, True), (False, False)])
-def test_gemma4_config_sets_shared_kv_cache_for_one_model_only(
-    one_model,
-    expected,
-):
-    spec_config = MTPDecodingConfig(
-        max_draft_len=3,
-        speculative_model="/tmp/gemma4-assistant",
-        mtp_eagle_one_model=one_model,
-    )
-    model_config = SimpleNamespace(
-        architectures=["Gemma4ForConditionalGeneration"],
-        num_nextn_predict_layers=1,
-    )
-
-    update_spec_config_from_model_config(spec_config, model_config)
-
-    assert spec_config._use_shared_kv_cache is expected
-
-
-def test_external_shared_kv_uses_no_draft_kv_cache():
-    spec_config = _shared_kv_spec_config()
-
-    assert spec_config._use_shared_kv_cache
-    assert get_num_spec_layers(spec_config) == 0
-    assert get_num_extra_kv_tokens(spec_config) == 0
-    assert not should_use_separate_draft_kv_cache(spec_config)
-
-
-def test_external_shared_kv_builds_draft_from_external_config(monkeypatch):
-    draft_config = object()
-    expected_model = object()
-    monkeypatch.setattr(
-        modeling_speculative.AutoModelForCausalLM,
-        "from_config",
-        lambda config: expected_model,
-    )
-
-    model_config = SimpleNamespace(spec_config=_shared_kv_spec_config())
-    assert (
-        modeling_speculative.get_draft_model(
-            model_config,
-            draft_config,
-            lm_head=None,
-            model=None,
-        )
-        is expected_model
-    )
-
-
-def test_shared_kv_alias_setup_rebinds_target_model():
-    calls = []
-    draft_model = SimpleNamespace(
-        shares_target_kv_cache=True,
-        load_weights_from_target_model=lambda target: calls.append(target),
-    )
-    model = SimpleNamespace(draft_model=draft_model)
-
-    modeling_speculative.SpecDecOneEngineForCausalLM.setup_aliases(model)
-
-    assert calls == [model]
-
-    model.draft_model = SimpleNamespace(shares_target_kv_cache=True)
-    modeling_speculative.SpecDecOneEngineForCausalLM.setup_aliases(model)
-
-
-def test_external_shared_kv_worker_uses_config_and_supports_guided_decoding():
-    worker = MTPEagleWorker(_shared_kv_spec_config())
-    guided_decoder = object()
-
-    assert worker._uses_external_shared_target_kv
-    assert worker.set_guided_decoder(guided_decoder)
-    assert worker.guided_decoder is guided_decoder
-    worker.set_draft_model(SimpleNamespace(model=SimpleNamespace()))
-
-
-def test_external_shared_kv_draft_loop_applies_guided_decoding(monkeypatch):
-    draft_metadata = SimpleNamespace(
-        update_shared_kv_draft_lengths=Mock(),
-    )
-
-    class FakeFlashInferAttentionMetadata:
-        def __init__(self):
-            self.seq_lens_cuda = torch.tensor([2, 2], dtype=torch.int32)
-
-        def get_draft_metadata(self):
-            return draft_metadata
-
-    monkeypatch.setattr(eagle3, "FlashInferAttentionMetadata", FakeFlashInferAttentionMetadata)
-    worker = MTPEagleWorker(_shared_kv_spec_config(max_draft_len=2))
-    guided_decoder = SimpleNamespace(
-        add_draft_batch=Mock(),
-        execute_draft_batch=Mock(),
-    )
-    worker.set_guided_decoder(guided_decoder)
-
-    sampled_tokens = [
-        torch.tensor([41, 42], dtype=torch.int32),
-        torch.tensor([51, 52], dtype=torch.int32),
-    ]
-    monkeypatch.setattr(
-        worker,
-        "sample_draft_tokens",
-        lambda *args, **kwargs: sampled_tokens.pop(0),
-    )
-    draft_model = SimpleNamespace(
-        forward_draft_step=lambda **kwargs: (torch.zeros(2, 4), kwargs["recurrent_hidden_states"])
-    )
-    attn_metadata = FakeFlashInferAttentionMetadata()
-    spec_metadata = SimpleNamespace(
-        batch_indices_cuda=torch.arange(2),
-        runtime_draft_len=2,
-        subseq_all_rank_num_tokens=None,
-    )
-    accepted_tokens = torch.tensor(
-        [[10, 11, 12], [20, 21, 22]],
-        dtype=torch.int32,
-    )
-    num_accepted_tokens = torch.ones(2, dtype=torch.long)
-
-    next_draft_tokens = worker._forward_external_shared_target_kv_draft_loop(
-        position_ids=torch.arange(4, dtype=torch.int32),
-        hidden_states=torch.arange(8, dtype=torch.float32).unsqueeze(1),
-        attn_metadata=attn_metadata,
-        spec_metadata=spec_metadata,
-        draft_model=draft_model,
-        accepted_tokens=accepted_tokens,
-        num_accepted_tokens=num_accepted_tokens,
-        num_contexts=1,
-        batch_size=2,
-    )
-
-    assert torch.equal(
-        next_draft_tokens,
-        torch.tensor([[41, 51], [42, 52]], dtype=torch.int32),
-    )
-    assert [
-        call.kwargs["draft_step"] for call in guided_decoder.add_draft_batch.call_args_list
-    ] == [0, 1]
-    assert [
-        call.kwargs["draft_step"] for call in guided_decoder.execute_draft_batch.call_args_list
-    ] == [0, 1]
-
-
-@pytest.mark.parametrize(
-    "num_accepted_tokens, expected_hidden_rows",
-    [
-        ([1, 1, 1], [1, 2, 6]),
-        ([1, 2, 3], [1, 3, 8]),
-        ([1, 4, 4], [1, 5, 9]),
-    ],
-)
-def test_external_shared_kv_selects_last_accepted_target_state(
-    num_accepted_tokens,
-    expected_hidden_rows,
-):
-    accepted_tokens = torch.tensor(
-        [
-            [10, 11, 12, 13],
-            [20, 21, 22, 23],
-            [30, 31, 32, 33],
-        ],
-        dtype=torch.int32,
-    )
-    accepted_counts = torch.tensor(num_accepted_tokens, dtype=torch.long)
-    hidden_states = torch.arange(20, dtype=torch.float32).unsqueeze(1)
-    position_ids = torch.arange(10, dtype=torch.int32).unsqueeze(0)
-
-    draft_ids, recurrent_hidden, draft_positions = (
-        MTPEagleWorker._prepare_external_shared_target_kv_draft_inputs(
-            accepted_tokens=accepted_tokens,
-            num_accepted_tokens=accepted_counts,
-            hidden_states=hidden_states,
-            position_ids=position_ids,
-            sequence_lengths=torch.tensor([2, 4, 4]),
-            num_contexts=1,
-            batch_indices=torch.arange(3),
-        )
-    )
-
-    expected_tokens = accepted_tokens[
-        torch.arange(3),
-        accepted_counts - 1,
-    ]
-    assert torch.equal(draft_ids, expected_tokens)
-    assert torch.equal(
-        recurrent_hidden.squeeze(1),
-        torch.tensor(expected_hidden_rows, dtype=torch.float32),
-    )
-    assert torch.equal(
-        draft_positions,
-        torch.tensor(expected_hidden_rows, dtype=torch.int32).unsqueeze(0) + 1,
-    )
-
-
-def test_gemma4_target_forward_dispatches_one_model_worker():
-    hidden_states = torch.tensor(
-        [
-            [1.0, 2.0],
-            [3.0, 4.0],
-            [5.0, 6.0],
-        ]
-    )
-    worker_calls = []
-
-    def spec_worker(**kwargs):
-        worker_calls.append(kwargs)
-        return {"logits": kwargs["logits"], "new_tokens": torch.tensor([[7]])}
-
-    model = SimpleNamespace(
-        layer_idx=-1,
-        config=SimpleNamespace(final_logit_softcapping=None),
-        model=lambda **kwargs: hidden_states,
-        logits_processor=SimpleNamespace(forward=lambda selected, *args: selected),
-        lm_head=object(),
-        spec_worker=spec_worker,
-        draft_model=object(),
-    )
-    spec_metadata = SimpleNamespace(
-        gather_ids=torch.tensor([2]),
-        is_layer_capture=lambda layer_idx: False,
-    )
-    attn_metadata = SimpleNamespace(padded_num_tokens=None)
-
-    outputs = Gemma4ForCausalLM.forward(
-        model,
-        attn_metadata=attn_metadata,
-        input_ids=torch.tensor([1, 2, 3]),
-        position_ids=torch.tensor([[0, 1, 2]]),
-        spec_metadata=spec_metadata,
-    )
-
-    assert torch.equal(outputs["new_tokens"], torch.tensor([[7]]))
-    assert len(worker_calls) == 1
-    assert torch.equal(worker_calls[0]["hidden_states"], hidden_states)
-    assert torch.equal(worker_calls[0]["logits"], hidden_states[[2]])
-    assert worker_calls[0]["draft_model"] is model.draft_model
diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
index a381dbd94907..b5f9a9c1db68 100644
--- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
+++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp.py
@@ -1,6 +1,7 @@
 import os
 import sys
 import unittest
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -10,7 +11,14 @@
 from tensorrt_llm import LLM, SamplingParams
 from tensorrt_llm._torch.attention_backend import TrtllmAttentionMetadata
 from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.speculative.eagle3 import MTPEagleWorker
+from tensorrt_llm._torch.speculative.interface import should_use_separate_draft_kv_cache
 from tensorrt_llm._torch.speculative.mtp import MTPHiddenStatesManager, MTPSpecMetadata, MTPWorker
+from tensorrt_llm._torch.speculative.utils import (
+    get_num_extra_kv_tokens,
+    get_num_spec_layers,
+    update_spec_config_from_model_config,
+)
 from tensorrt_llm.llmapi import KvCacheConfig, MTPDecodingConfig
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
@@ -1733,6 +1741,69 @@ def test_prepare_drafter_inputs(
         torch.testing.assert_close(draft_inputs["hidden_states"], ref_previous_hidden_states)
 
 
+@pytest.mark.parametrize(
+    ("architecture", "one_model", "expected"),
+    [
+        ("Gemma4ForCausalLM", True, True),
+        ("Gemma4ForConditionalGeneration", False, False),
+        ("LlamaForCausalLM", True, False),
+    ],
+)
+def test_mtp_shared_kv_config(architecture, one_model, expected):
+    spec_config = MTPDecodingConfig(
+        max_draft_len=3,
+        speculative_model="/tmp/assistant",
+        mtp_eagle_one_model=one_model,
+    )
+    model_config = SimpleNamespace(
+        architectures=[architecture],
+        num_nextn_predict_layers=1,
+    )
+
+    update_spec_config_from_model_config(spec_config, model_config)
+
+    assert spec_config._use_shared_kv_cache is expected
+    if expected:
+        assert get_num_spec_layers(spec_config) == 0
+        assert get_num_extra_kv_tokens(spec_config) == 0
+        assert not should_use_separate_draft_kv_cache(spec_config)
+
+
+def test_mtp_shared_kv_draft_inputs():
+    spec_config = MTPDecodingConfig(
+        max_draft_len=3,
+        speculative_model="/tmp/assistant",
+        mtp_eagle_one_model=True,
+    )
+    spec_config._use_shared_kv_cache = True
+    worker = MTPEagleWorker(spec_config)
+    accepted_tokens = torch.tensor(
+        [
+            [10, 11, 12],
+            [20, 21, 22],
+            [30, 31, 32],
+        ],
+        dtype=torch.int32,
+    )
+
+    draft_ids, recurrent_hidden, draft_positions = worker._prepare_shared_kv_draft_inputs(
+        accepted_tokens=accepted_tokens,
+        num_accepted_tokens=torch.tensor([1, 2, 3]),
+        hidden_states=torch.arange(20, dtype=torch.float32).unsqueeze(1),
+        position_ids=torch.arange(10, dtype=torch.int32).unsqueeze(0),
+        sequence_lengths=torch.tensor([2, 4, 4]),
+        num_contexts=1,
+        batch_indices=torch.arange(3),
+    )
+
+    torch.testing.assert_close(draft_ids, torch.tensor([10, 21, 32], dtype=torch.int32))
+    torch.testing.assert_close(recurrent_hidden.squeeze(1), torch.tensor([1.0, 3.0, 8.0]))
+    torch.testing.assert_close(
+        draft_positions,
+        torch.tensor([[2, 4, 9]], dtype=torch.int32),
+    )
+
+
 @pytest.mark.high_cuda_memory
 def test_vanilla_mtp_rejection():
     """Vanilla MTP with rejection sampling on: the per-step rejection path

From a9881dcf04f63fb05d26eecbc99dfe521edbee69 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Tue, 28 Jul 2026 16:06:04 +0000
Subject: [PATCH 22/26] [None][refactor] simplify Gemma4 config and embedding
 sharing

Consolidate the Gemma4 compatibility configs in one module and make the target input embedding a direct, explicitly named alias while preserving the assistant embedding tied to its LM head.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/configs/__init__.py       |   4 +-
 tensorrt_llm/_torch/configs/gemma4.py         | 118 ++++++++++++++
 tensorrt_llm/_torch/configs/gemma4_unified.py | 145 ------------------
 tensorrt_llm/_torch/models/modeling_gemma4.py |  19 +--
 .../_torch/models/modeling_gemma4_unified.py  |   2 +-
 .../_torch/modeling/test_modeling_gemma4.py   |   3 +-
 6 files changed, 130 insertions(+), 161 deletions(-)
 delete mode 100644 tensorrt_llm/_torch/configs/gemma4_unified.py

diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py
index 64eedab07343..c5893c21bfae 100644
--- a/tensorrt_llm/_torch/configs/__init__.py
+++ b/tensorrt_llm/_torch/configs/__init__.py
@@ -16,8 +16,8 @@
 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 (
+from tensorrt_llm._torch.configs.gemma4 import (
+    Gemma4AssistantConfig,
     Gemma4UnifiedAudioConfig,
     Gemma4UnifiedConfig,
     Gemma4UnifiedTextConfig,
diff --git a/tensorrt_llm/_torch/configs/gemma4.py b/tensorrt_llm/_torch/configs/gemma4.py
index bb455e2d3f83..b0b2c17e42cb 100644
--- a/tensorrt_llm/_torch/configs/gemma4.py
+++ b/tensorrt_llm/_torch/configs/gemma4.py
@@ -12,6 +12,7 @@
 # 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.
+"""Compatibility configs for Gemma 4 assistant and Unified checkpoints."""
 
 from transformers import Gemma4TextConfig, PreTrainedConfig
 
@@ -93,3 +94,120 @@ def vocab_size(self):
     @property
     def num_hidden_layers(self):
         return self.text_config.num_hidden_layers
+
+
+class Gemma4UnifiedTextConfig(Gemma4TextConfig):
+    """Text sub-config for Gemma 4 12B Unified.
+
+    The 12B text backbone is a standard dense Gemma 4 text model; only the
+    model_type string differs, so this is a pure alias of the native
+    `Gemma4TextConfig`.
+    """
+
+    model_type = "gemma4_unified_text"
+
+
+class Gemma4UnifiedVisionConfig(PreTrainedConfig):
+    """Sub-config for the encoder-free vision projector."""
+
+    model_type = "gemma4_unified_vision"
+
+    def __init__(
+        self,
+        mm_embed_dim: int = 3840,
+        mm_posemb_size: int = 1120,
+        output_proj_dims: int = 3840,
+        patch_size: int = 16,
+        pooling_kernel_size: int = 3,
+        rms_norm_eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.mm_embed_dim = mm_embed_dim
+        self.mm_posemb_size = mm_posemb_size
+        self.output_proj_dims = output_proj_dims
+        self.patch_size = patch_size
+        self.pooling_kernel_size = pooling_kernel_size
+        self.rms_norm_eps = rms_norm_eps
+
+
+class Gemma4UnifiedAudioConfig(PreTrainedConfig):
+    """Sub-config for the encoder-free audio projector.
+
+    `output_proj_dims` and `hidden_size` alias `audio_embed_dim` (the raw audio
+    frame width) when not given, matching the HF implementation; they are plain
+    attributes here so a checkpoint config.json that spells them out loads as-is.
+    """
+
+    model_type = "gemma4_unified_audio"
+
+    def __init__(
+        self,
+        audio_embed_dim: int = 640,
+        rms_norm_eps: float = 1e-6,
+        output_proj_dims: int | None = None,
+        hidden_size: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.audio_embed_dim = audio_embed_dim
+        self.rms_norm_eps = rms_norm_eps
+        self.output_proj_dims = (
+            output_proj_dims if output_proj_dims is not None else audio_embed_dim
+        )
+        self.hidden_size = hidden_size if hidden_size is not None else audio_embed_dim
+
+
+class Gemma4UnifiedConfig(PreTrainedConfig):
+    """Top-level config for Gemma 4 12B Unified (encoder-free multimodal).
+
+    Parses `config.json` fields required by
+    `Gemma4UnifiedForConditionalGeneration` without depending on any
+    natively shipped transformers class. The `text_config`, `vision_config`, and
+    `audio_config` sub-configs are reconstructed from nested dicts using the
+    shim classes above.
+    """
+
+    model_type = "gemma4_unified"
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        audio_config=None,
+        image_token_id: int = 258880,
+        audio_token_id: int = 258881,
+        video_token_id: int = 258884,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.image_token_id = image_token_id
+        self.audio_token_id = audio_token_id
+        self.video_token_id = video_token_id
+
+        # Sub-configs arrive as dicts from AutoConfig.from_pretrained; rebuild
+        # them with the classes above.
+        if text_config is not None:
+            if isinstance(text_config, dict):
+                self.text_config = Gemma4UnifiedTextConfig(**text_config)
+            else:
+                self.text_config = text_config
+        else:
+            self.text_config = None
+
+        if vision_config is not None:
+            if isinstance(vision_config, dict):
+                self.vision_config = Gemma4UnifiedVisionConfig(**vision_config)
+            else:
+                self.vision_config = vision_config
+        else:
+            self.vision_config = None
+
+        if audio_config is not None:
+            if isinstance(audio_config, dict):
+                self.audio_config = Gemma4UnifiedAudioConfig(**audio_config)
+            else:
+                self.audio_config = audio_config
+        else:
+            self.audio_config = None
diff --git a/tensorrt_llm/_torch/configs/gemma4_unified.py b/tensorrt_llm/_torch/configs/gemma4_unified.py
deleted file mode 100644
index 561474ffa891..000000000000
--- a/tensorrt_llm/_torch/configs/gemma4_unified.py
+++ /dev/null
@@ -1,145 +0,0 @@
-# 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.
-"""Config classes for Gemma 4 12B Unified (encoder-free multimodal).
-
-Registered with the transformers CONFIG_MAPPING (see `_torch/configs/__init__.py`)
-so `AutoConfig.from_pretrained` can parse a Gemma 4 12B checkpoint whenever the
-installed transformers does not ship the `gemma4_unified` model_types natively.
-
-All fields are read directly from the checkpoint's config.json, matching the
-attribute names used by `Gemma4UnifiedForConditionalGeneration`. The text
-backbone of the 12B is a standard dense Gemma 4 text model, so its sub-config
-reuses the native `Gemma4TextConfig`.
-"""
-
-from transformers import Gemma4TextConfig
-from transformers.configuration_utils import PretrainedConfig
-
-
-class Gemma4UnifiedTextConfig(Gemma4TextConfig):
-    """Text sub-config for Gemma 4 12B Unified.
-
-    The 12B text backbone is a standard dense Gemma 4 text model; only the
-    model_type string differs, so this is a pure alias of the native
-    `Gemma4TextConfig`.
-    """
-
-    model_type = "gemma4_unified_text"
-
-
-class Gemma4UnifiedVisionConfig(PretrainedConfig):
-    """Sub-config for the encoder-free vision projector."""
-
-    model_type = "gemma4_unified_vision"
-
-    def __init__(
-        self,
-        mm_embed_dim: int = 3840,
-        mm_posemb_size: int = 1120,
-        output_proj_dims: int = 3840,
-        patch_size: int = 16,
-        pooling_kernel_size: int = 3,
-        rms_norm_eps: float = 1e-6,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.mm_embed_dim = mm_embed_dim
-        self.mm_posemb_size = mm_posemb_size
-        self.output_proj_dims = output_proj_dims
-        self.patch_size = patch_size
-        self.pooling_kernel_size = pooling_kernel_size
-        self.rms_norm_eps = rms_norm_eps
-
-
-class Gemma4UnifiedAudioConfig(PretrainedConfig):
-    """Sub-config for the encoder-free audio projector.
-
-    `output_proj_dims` and `hidden_size` alias `audio_embed_dim` (the raw audio
-    frame width) when not given, matching the HF implementation; they are plain
-    attributes here so a checkpoint config.json that spells them out loads as-is.
-    """
-
-    model_type = "gemma4_unified_audio"
-
-    def __init__(
-        self,
-        audio_embed_dim: int = 640,
-        rms_norm_eps: float = 1e-6,
-        output_proj_dims: int | None = None,
-        hidden_size: int | None = None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.audio_embed_dim = audio_embed_dim
-        self.rms_norm_eps = rms_norm_eps
-        self.output_proj_dims = (
-            output_proj_dims if output_proj_dims is not None else audio_embed_dim
-        )
-        self.hidden_size = hidden_size if hidden_size is not None else audio_embed_dim
-
-
-class Gemma4UnifiedConfig(PretrainedConfig):
-    """Top-level config for Gemma 4 12B Unified (encoder-free multimodal).
-
-    Parses `config.json` fields required by
-    `Gemma4UnifiedForConditionalGeneration` without depending on any
-    natively shipped transformers class. The `text_config`, `vision_config`, and
-    `audio_config` sub-configs are reconstructed from nested dicts using the
-    shim classes above.
-    """
-
-    model_type = "gemma4_unified"
-
-    def __init__(
-        self,
-        text_config=None,
-        vision_config=None,
-        audio_config=None,
-        image_token_id: int = 258880,
-        audio_token_id: int = 258881,
-        video_token_id: int = 258884,
-        tie_word_embeddings: bool = True,
-        **kwargs,
-    ):
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-        self.image_token_id = image_token_id
-        self.audio_token_id = audio_token_id
-        self.video_token_id = video_token_id
-
-        # Sub-configs arrive as dicts from AutoConfig.from_pretrained; rebuild
-        # them with the classes above.
-        if text_config is not None:
-            if isinstance(text_config, dict):
-                self.text_config = Gemma4UnifiedTextConfig(**text_config)
-            else:
-                self.text_config = text_config
-        else:
-            self.text_config = None
-
-        if vision_config is not None:
-            if isinstance(vision_config, dict):
-                self.vision_config = Gemma4UnifiedVisionConfig(**vision_config)
-            else:
-                self.vision_config = vision_config
-        else:
-            self.vision_config = None
-
-        if audio_config is not None:
-            if isinstance(audio_config, dict):
-                self.audio_config = Gemma4UnifiedAudioConfig(**audio_config)
-            else:
-                self.audio_config = audio_config
-        else:
-            self.audio_config = None
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 23ee15950f3b..65020e25111f 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -16,7 +16,6 @@
 
 import dataclasses
 import math
-import weakref
 from typing import Dict, Optional, Tuple, Union
 
 import torch
@@ -1581,11 +1580,13 @@ def __init__(self, model_config: ModelConfig):
             if assistant_config.use_ordered_embeddings
             else None
         )
-        self._target_embed_tokens_ref = None
+        # The assistant embedding remains tied to its own LM head. Target input
+        # embeddings have the backbone width and are shared separately.
+        self.target_input_embeddings = None
 
     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)
+        self.target_input_embeddings = target_llm.model.embed_tokens
 
         target_config = target_llm.config
         num_kv_shared = getattr(target_config, "num_kv_shared_layers", 0)
@@ -1600,14 +1601,6 @@ def load_weights_from_target_model(self, target_model: nn.Module) -> None:
             )
             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,
@@ -1631,7 +1624,9 @@ def forward_draft_step(
         spec_metadata=None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Run one Q-only assistant step over a frozen target KV prefix."""
-        target_embeddings = self._get_target_embeddings(input_ids)
+        if self.target_input_embeddings is None:
+            raise RuntimeError("Gemma4 assistant target embeddings have not been initialized")
+        target_embeddings = self.target_input_embeddings(input_ids)
         assistant_inputs = self.pre_projection(
             torch.cat([target_embeddings, recurrent_hidden_states], dim=-1)
         )
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py
index 77a111d05874..0acd09ac63a0 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4_unified.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4_unified.py
@@ -43,7 +43,7 @@
 `_get_audio_features`, and `load_weights` for the encoder-free projections.
 
 TRT-LLM provides its own `gemma4_unified` config classes
-(`_torch/configs/gemma4_unified.py`) and multimodal preprocessing (the vendored
+(`_torch/configs/gemma4.py`) and multimodal preprocessing (the vendored
 section at the end of this file), used whenever the installed transformers does
 not ship them natively — the full model (text + image + audio + video) runs on
 the repo's pinned transformers.
diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
index 5471e0a6ae9f..5913dd3a518d 100644
--- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py
+++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py
@@ -639,7 +639,8 @@ def test_assistant_uses_target_kv_sources(self):
         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)
+        self.assertIs(assistant.target_input_embeddings, target.model.embed_tokens)
+        self.assertIsNot(assistant.model.embed_tokens, target.model.embed_tokens)
 
 
 # ---------------------------------------------------------------------------

From 1937164ede66b9d4616be25a4b17d736af3f8180 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 29 Jul 2026 04:22:33 +0000
Subject: [PATCH 23/26] [None][refactor] address Gemma4 review feedback

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma4.py | 106 ++++++++++++------
 .../_torch/models/modeling_speculative.py     |   6 +-
 2 files changed, 70 insertions(+), 42 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 65020e25111f..23a3d78ba31b 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -16,7 +16,7 @@
 
 import dataclasses
 import math
-from typing import Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -61,6 +61,9 @@
 from .modeling_speculative import SpecDecOneEngineForCausalLM, _slice_spec_position_ids
 from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model
 
+if TYPE_CHECKING:
+    from .modeling_gemma4mm import Gemma4ForConditionalGeneration
+
 _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0"
 if Version(transformers.__version__) < Version(_MIN_TRANSFORMERS_FOR_GEMMA4):
     raise ImportError(
@@ -1367,6 +1370,45 @@ def get_flashinfer_attention_mask(
             token_offset = context_end
         return torch.cat(context_mask_list, dim=0).contiguous()
 
+    def _forward_speculative(
+        self,
+        output: torch.Tensor,
+        input_ids: Optional[torch.IntTensor],
+        orig_input_ids: Optional[torch.IntTensor],
+        position_ids: Optional[torch.IntTensor],
+        attn_metadata: AttentionMetadata,
+        spec_metadata: SpecMetadata,
+        resource_manager,
+    ) -> torch.Tensor:
+        logits = self.logits_processor.forward(
+            output[spec_metadata.gather_ids],
+            self.lm_head,
+            attn_metadata,
+            True,
+        )
+        if self.config.final_logit_softcapping is not None:
+            cap = self.config.final_logit_softcapping
+            logits = torch.tanh(logits / cap) * cap
+
+        spec_input_ids = input_ids if input_ids is not None else orig_input_ids
+        spec_position_ids = position_ids
+        if attn_metadata.padded_num_tokens is not None:
+            if spec_input_ids is not None:
+                spec_input_ids = spec_input_ids[: attn_metadata.num_tokens]
+            if position_ids is not None:
+                spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens)
+
+        return self.spec_worker(
+            input_ids=spec_input_ids,
+            position_ids=spec_position_ids,
+            hidden_states=output,
+            logits=logits,
+            attn_metadata=attn_metadata,
+            spec_metadata=spec_metadata,
+            draft_model=self.draft_model,
+            resource_manager=resource_manager,
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -1377,11 +1419,11 @@ def forward(
         return_context_logits: bool = False,
         mm_token_type_ids: Optional[torch.Tensor] = None,
         spec_metadata: Optional[SpecMetadata] = None,
+        resource_manager=None,
+        orig_input_ids: Optional[torch.IntTensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         local_attention_mask_data = None
-        resource_manager = kwargs.pop("resource_manager", None)
-        orig_input_ids = kwargs.pop("orig_input_ids", None)
         # Only build bidirectional masks when use_bidirectional_attention is
         # set to "vision" (26B, 31B).  E2B/E4B have this as None and should
         # use standard causal attention even for multimodal tokens. Gemma4
@@ -1411,35 +1453,14 @@ def forward(
             output = output[: attn_metadata.num_tokens]
 
         if self.spec_worker is not None:
-            logits = self.logits_processor.forward(
-                output[spec_metadata.gather_ids],
-                self.lm_head,
+            return self._forward_speculative(
+                output,
+                input_ids,
+                orig_input_ids,
+                position_ids,
                 attn_metadata,
-                True,
-            )
-            if self.config.final_logit_softcapping is not None:
-                cap = self.config.final_logit_softcapping
-                logits = torch.tanh(logits / cap) * cap
-
-            spec_input_ids = input_ids if input_ids is not None else orig_input_ids
-            spec_position_ids = position_ids
-            if attn_metadata.padded_num_tokens is not None:
-                if spec_input_ids is not None:
-                    spec_input_ids = spec_input_ids[: attn_metadata.num_tokens]
-                if position_ids is not None:
-                    spec_position_ids = _slice_spec_position_ids(
-                        position_ids, attn_metadata.num_tokens
-                    )
-
-            return self.spec_worker(
-                input_ids=spec_input_ids,
-                position_ids=spec_position_ids,
-                hidden_states=output,
-                logits=logits,
-                attn_metadata=attn_metadata,
-                spec_metadata=spec_metadata,
-                draft_model=self.draft_model,
-                resource_manager=resource_manager,
+                spec_metadata,
+                resource_manager,
             )
 
         logits = self.logits_processor.forward(
@@ -1474,6 +1495,11 @@ def __init__(self, model_config: ModelConfig):
         self.num_centroids = config.num_centroids
         self.centroid_intermediate_top_k = config.centroid_intermediate_top_k
         self.vocab_size = config.vocab_size
+        if self.vocab_size % self.num_centroids != 0:
+            raise ValueError(
+                "Gemma4 assistant vocab_size must be divisible by num_centroids: "
+                f"got {self.vocab_size} and {self.num_centroids}"
+            )
         self.vocab_size_per_centroid = self.vocab_size // self.num_centroids
         self.centroids = Linear(
             self.hidden_size,
@@ -1487,7 +1513,7 @@ def __init__(self, model_config: ModelConfig):
         )
         self.register_buffer(
             "token_ordering",
-            torch.empty(self.vocab_size, dtype=torch.long, device="cuda"),
+            torch.empty(self.vocab_size, dtype=torch.long),
         )
 
     @staticmethod
@@ -1500,7 +1526,7 @@ def _selected_logits_for_vocab_shard(
         """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)
+        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],
@@ -1511,6 +1537,7 @@ def _selected_logits_for_vocab_shard(
         ).squeeze(1)
         return selected_logits.masked_fill(~is_local, 0)
 
+    @torch.inference_mode()
     def forward(self, hidden_states: torch.Tensor, lm_head: nn.Module) -> torch.Tensor:
         centroid_logits = self.centroids(hidden_states)
         _, top_k_indices = torch.topk(
@@ -1584,13 +1611,17 @@ def __init__(self, model_config: ModelConfig):
         # embeddings have the backbone width and are shared separately.
         self.target_input_embeddings = None
 
-    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
+    def load_weights_from_target_model(
+        self,
+        target_model: "Gemma4ForCausalLM | Gemma4ForConditionalGeneration",
+    ) -> None:
+        target_llm = (
+            target_model if isinstance(target_model, Gemma4ForCausalLM) else target_model.llm
+        )
         self.target_input_embeddings = 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
+        num_source_layers = target_config.num_hidden_layers - target_config.num_kv_shared_layers
         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"
@@ -1647,6 +1678,7 @@ def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         weights = weight_mapper.preprocess_weights(weights)
         ordering_weight_name = "masked_embedding.token_ordering"
         if self.masked_embedding is not None and ordering_weight_name in weights:
+            # Copy this checkpoint-backed buffer and consume its exact source key.
             self.masked_embedding.token_ordering.copy_(weights[ordering_weight_name])
             del weights[ordering_weight_name]
         super().load_weights(weights, weight_mapper)
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 35b477960bf0..16f31e5f3f2e 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -2047,11 +2047,7 @@ def __init__(self,
     def setup_aliases(self) -> None:
         if (self.draft_model is not None
                 and getattr(self.draft_model, "shares_target_kv_cache", False)):
-            setup_target_aliases = getattr(self.draft_model,
-                                           "load_weights_from_target_model",
-                                           None)
-            if callable(setup_target_aliases):
-                setup_target_aliases(self)
+            self.draft_model.load_weights_from_target_model(self)
 
     def forward(
         self,

From 660735f02f4ff7402ae76592e9e9f2a18d217d50 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Wed, 29 Jul 2026 13:26:38 +0000
Subject: [PATCH 24/26] [None][fix] fix speculative decoding CI regressions

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/speculative/utils.py                     | 4 ++--
 tests/unittest/_torch/attention/test_flashinfer_attention.py | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py
index 887d2bd7cda5..a8d92fed4f46 100644
--- a/tensorrt_llm/_torch/speculative/utils.py
+++ b/tensorrt_llm/_torch/speculative/utils.py
@@ -448,7 +448,7 @@ def get_spec_drafter(model_engine,
 
 
 def get_num_spec_layers(spec_config):
-    if spec_config._use_shared_kv_cache:
+    if getattr(spec_config, "_use_shared_kv_cache", False):
         return 0
     if spec_config.spec_dec_mode.is_mtp_eagle_one_model():
         return 1
@@ -526,7 +526,7 @@ def get_num_extra_kv_tokens(spec_config):
     """
     if spec_config is None:
         return 0
-    if spec_config._use_shared_kv_cache:
+    if getattr(spec_config, "_use_shared_kv_cache", False):
         return 0
     if spec_config.spec_dec_mode.use_one_engine():
         return spec_config.max_draft_len - 1
diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py
index 9ce210ef28c9..fe00e3d69d07 100644
--- a/tests/unittest/_torch/attention/test_flashinfer_attention.py
+++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py
@@ -69,6 +69,8 @@ class TestFlashInferAttention(unittest.TestCase):
     def test_separate_kv_draft_metadata_uses_draft_manager(self):
         if not torch.cuda.is_available():
             self.skipTest("CUDA is required for FlashInfer metadata")
+        if torch.cuda.get_device_capability() not in ((10, 0), (10, 3)):
+            self.skipTest("FlashInfer trtllm-gen requires SM100 or SM103")
 
         def create_manager():
             return KVCacheManager(

From 5dcaf06aded29036839f35ee369d1260c72cde83 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Fri, 31 Jul 2026 04:05:56 +0000
Subject: [PATCH 25/26] [None][fix] harden Gemma4 speculative draft setup

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/flashinfer.py | 3 +++
 tensorrt_llm/_torch/models/modeling_gemma4.py       | 2 +-
 tensorrt_llm/_torch/models/modeling_speculative.py  | 7 +++++--
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 94c9fee825cb..d49ab084a128 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -632,6 +632,9 @@ def get_draft_metadata(
                 draft_metadata.seq_lens = self.seq_lens
                 draft_metadata.num_contexts = self.num_contexts
             draft_metadata.seq_lens_kv = None
+            if not uses_shared_kv_cache:
+                for pool_id in set((self._vswa_layer_to_pool or {}).values()):
+                    setattr(draft_metadata, f"_vswa_pool_buf_{pool_id}", None)
             draft_metadata.__post_init__()
             if not uses_shared_kv_cache:
                 # Keep the existing speculative worker's backend-neutral
diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py
index 23a3d78ba31b..24aa619b8e34 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma4.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma4.py
@@ -1652,7 +1652,7 @@ def forward_draft_step(
         position_ids: torch.IntTensor,
         recurrent_hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
-        spec_metadata=None,
+        spec_metadata: Optional[SpecMetadata] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Run one Q-only assistant step over a frozen target KV prefix."""
         if self.target_input_embeddings is None:
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index 16f31e5f3f2e..a572ac94f278 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/tensorrt_llm/_torch/models/modeling_speculative.py
@@ -1889,8 +1889,11 @@ def get_draft_model(model_config, draft_config, lm_head, model):
                 f"Unsupported eagle3 model architecture: {spec_dec_mode.eagle3_model_arch}"
             )
 
-    elif (model_config.spec_config._use_shared_kv_cache
-          and draft_config is not None):
+    elif model_config.spec_config._use_shared_kv_cache:
+        if draft_config is None:
+            raise ValueError(
+                "Shared-KV speculative decoding requires an external draft "
+                "model config.")
         return AutoModelForCausalLM.from_config(draft_config)
     elif spec_dec_mode.is_mtp_one_model():
         return MTPForCausalLM(model_config,

From 8df01d7a8540150832b8e1391cdd3d2deadcda11 Mon Sep 17 00:00:00 2001
From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Date: Mon, 3 Aug 2026 04:56:43 +0000
Subject: [PATCH 26/26] [None][fix] scope FlashInfer one-engine guard

Allow only the validated shared-KV one-engine path while preserving the existing guard for other FlashInfer combinations.

Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 8cc7c8c0730d..03ca7ac061c1 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -451,6 +451,13 @@ def create_py_executor(
             )
             llm_args.disable_overlap_scheduler = True
 
+    if (spec_config is not None and llm_args.attn_backend == "FLASHINFER"
+            and spec_config.spec_dec_mode.use_one_engine()
+            and not spec_config._use_shared_kv_cache):
+        raise ValueError(
+            "FLASHINFER attention backend supports one-engine speculative "
+            "decoding only when the draft model shares the target KV cache.")
+
     if mm_encoder_only:
         llm_args.mm_encoder_only = True
         llm_args.disable_overlap_scheduler = True