diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index b03f9961f7cc..dfa283beecde 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` | Gemma 4 MTP assistant | `google/gemma-4-E2B-it-assistant`, `google/gemma-4-E4B-it-assistant`, `google/gemma-4-26B-A4B-it-assistant`, `google/gemma-4-31B-it-assistant` | | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `THUDM/GLM-4-100B-A10B` | | `Glm4MoeLiteForCausalLM` [^5] | GLM-4.7-Flash | `zai-org/GLM-4.7-Flash` | | `GlmMoeDsaForCausalLM` | GLM-5 | `zai-org/GLM-5` | @@ -75,7 +76,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | N/A | Yes | Yes | | `Glm4MoeLiteForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | | `NemotronHForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Untested | Untested | -| `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | +| `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | Yes | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Gemma4UnifiedForConditionalGeneration` | Untested | Untested | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index d2acd490d44a..081ed7138f5a 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -303,6 +303,9 @@ def setup_llm(args, **kwargs): if spec_decode_algo == 'MTP': if not args.use_one_model: print("Running MTP eagle with two model style.") + if args.draft_model_dir is None: + raise ValueError( + "--draft_model_dir is required for two-model MTP") spec_config = MTPDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, use_relaxed_acceptance_for_thinking=args. @@ -313,7 +316,8 @@ def setup_llm(args, **kwargs): use_dynamic_tree=args.use_dynamic_tree, dynamic_tree_max_topK=args.dynamic_tree_max_topK, max_total_draft_tokens=args.max_total_draft_tokens, - speculative_model=args.model_dir) + speculative_model=args.model_dir + if args.use_one_model else args.draft_model_dir) elif spec_decode_algo == "EAGLE3": spec_config = Eagle3DecodingConfig( max_draft_len=args.spec_decode_max_draft_len, diff --git a/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: diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 733d477009cb..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. @@ -225,6 +226,8 @@ class FlashInferWrappers: # and columns before narrowing future updates. decode_block_table_active_rows: int = field(default=0, repr=False) decode_block_table_active_width: int = field(default=0, repr=False) + host_prefill_block_tables: Optional[torch.Tensor] = field(default=None, + repr=False) @dataclass(kw_only=True) @@ -704,15 +707,6 @@ def _post_init_with_buffers(self, buffers) -> None: self._vswa_layer_to_pool[layer_idx] = pool_id if pool_id not in self._vswa_pool_to_rep_layer: self._vswa_pool_to_rep_layer[pool_id] = layer_idx - # Build head_dim → pool_id mapping using V2 per-layer head_dim - self._vswa_head_dim_to_pool: Dict[int, int] = {} - if hasattr(mgr, 'head_dim_per_layer'): - for layer_idx, pool_id in self._vswa_layer_to_pool.items(): - hd = mgr.head_dim_per_layer[ - mgr.layer_offsets[layer_idx]] - if hd not in self._vswa_head_dim_to_pool: - self._vswa_head_dim_to_pool[hd] = pool_id - # Pre-allocate VSWA pool cache buffers. These must be # stable (never reallocated) so that CUDA-graph-recorded # copies reference valid addresses across replays. @@ -1325,13 +1319,54 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: self._host_paged_kv_indices = \ self._host_pool_indices[primary_pool_id] + # trtllm-gen paged-prefill graphs capture one stable block table per KV + # pool. Refresh those tables and KV lengths after request turnover. + if (self.is_cuda_graph and self.num_contexts > 0 + and self._vswa_layer_to_pool is not None): + for plan_params, wrappers in self._plan_params_to_wrappers.items(): + if plan_params.attention_mask_data is not None: + continue + prefill_wrapper = wrappers.prefill_wrapper + if (prefill_wrapper is None + or prefill_wrapper._backend != "trtllm-gen"): + continue + block_tables = prefill_wrapper._block_tables + pool_id = plan_params.kv_pool_id + if block_tables is None or pool_id is None: + continue + host_pool_indices = self._host_pool_indices[pool_id] + host_block_tables = wrappers.host_prefill_block_tables + if (host_block_tables is None + or host_block_tables.shape != block_tables.shape): + host_block_tables = torch.zeros( + block_tables.shape, + dtype=torch.int32, + pin_memory=prefer_pinned(), + ) + wrappers.host_prefill_block_tables = host_block_tables + else: + host_block_tables.zero_() + source_offset = 0 + for row, num_blocks_for_row in enumerate( + num_blocks[:self.num_contexts]): + copy_width = min(int(num_blocks_for_row), + block_tables.size(1)) + host_block_tables[row, :copy_width].copy_( + host_pool_indices[source_offset:source_offset + + copy_width]) + source_offset += int(num_blocks_for_row) + block_tables.copy_(host_block_tables, non_blocking=True) + prefill_wrapper._kv_lens_buffer[:self.num_contexts].copy_( + _to_int32_tensor(kv_lens_host[:self.num_contexts]), + non_blocking=True, + ) + # CUDA graph + trtllm-gen: update _block_tables and _kv_lens_buffer # so the trtllm-gen decode kernel uses current page indices. if (self.is_cuda_graph and self._vswa_layer_to_pool is not None and self._vswa_pool_indices_cache is not None and self.num_generations > 0): decode_blocks = num_blocks[self.num_contexts:] - head_dim_to_pool = getattr(self, '_vswa_head_dim_to_pool', None) for plan_params, wrappers in self._plan_params_to_wrappers.items(): if plan_params.attention_mask_data is not None: continue @@ -1339,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 @@ -1439,6 +1473,7 @@ def plan(self, attention_mask_type=AttentionMaskType(attention_mask_type), attention_mask_data=attention_mask_data, multi_item_params=self._multi_item_params, + kv_pool_id=getattr(self, "_vswa_active_pool_id", None), ) return self._plan_with_params(plan_params, flashinfer_backend) diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index 031c2a534e2a..3bcf6fefdf5b 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -1,6 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from tensorrt_llm._torch.configs.cosmos3 import Cosmos3Config from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config +from tensorrt_llm._torch.configs.gemma4 import Gemma4AssistantConfig from tensorrt_llm._torch.configs.gemma4_unified import ( Gemma4UnifiedAudioConfig, Gemma4UnifiedConfig, @@ -36,6 +52,7 @@ def _register_custom_configs_with_transformers() -> None: "deepseek_v32": DeepseekV3Config, "kimi_k2": DeepseekV3Config, "deepseek_v4": DeepseekV4Config, + "gemma4_assistant": Gemma4AssistantConfig, "laguna": LagunaConfig, "gemma4_unified": Gemma4UnifiedConfig, "gemma4_unified_text": Gemma4UnifiedTextConfig, @@ -59,6 +76,7 @@ def _register_custom_configs_with_transformers() -> None: "Cosmos3Config", "DeepseekV3Config", "DeepseekV4Config", + "Gemma4AssistantConfig", "Gemma4UnifiedAudioConfig", "Gemma4UnifiedConfig", "Gemma4UnifiedTextConfig", diff --git a/tensorrt_llm/_torch/configs/gemma4.py b/tensorrt_llm/_torch/configs/gemma4.py new file mode 100644 index 000000000000..e2942239790e --- /dev/null +++ b/tensorrt_llm/_torch/configs/gemma4.py @@ -0,0 +1,106 @@ +# 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} + + # 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=1536, + use_ordered_embeddings=False, + num_centroids=2048, + centroid_intermediate_top_k=32, + **kwargs, + ): + if text_config is None: + 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 + 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 + + @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/checkpoints/hf/gemma4_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py index e336ef4b894b..3f69b62c9fbb 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py @@ -28,6 +28,7 @@ @register_mapper("HF", "Gemma4ForCausalLM") @register_mapper("HF", "Gemma4ForConditionalGeneration") @register_mapper("HF", "Gemma4UnifiedForConditionalGeneration") +@register_mapper("HF", "Gemma4AssistantForCausalLM") class Gemma4HfWeightMapper(HfWeightMapper): @property def _is_vlm(self) -> bool: diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index bf356d5e3ba1..ea8705030b68 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 @@ -40,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 @@ -54,7 +57,9 @@ 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 _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -656,7 +661,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 +1228,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 +1248,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: @@ -1380,6 +1380,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 @@ -1406,6 +1407,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, @@ -1426,3 +1432,224 @@ 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.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"), + ) + + @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, + 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].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, + ) + 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, 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).float() + return self.lm_head(logits_hidden_states).float() + + 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: + 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_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 0f5a340c0a9a..bf295489f971 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -1098,6 +1098,7 @@ def forward( mm_token_type_ids=mm_token_type_ids, ple_input_ids=ple_input_ids, lora_params=kwargs.get("lora_params", None), + spec_metadata=kwargs.get("spec_metadata", None), ) return logits diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..0078f6c46d49 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1580,11 +1580,16 @@ def build_managers(self, # Split combined KV cache budgets before creating managers. Skip during # estimation — estimation uses max_tokens-based logic and must not # mutate the config. - has_draft = ( - self._draft_model_engine is not None # two-model + draft_shares_target_kv_cache = ( + self._draft_model_engine is not None + and self._draft_model_engine.kv_cache_manager_key + == ResourceManagerType.KV_CACHE_MANAGER) + has_independent_draft_cache = ( + (self._draft_model_engine is not None + and not draft_shares_target_kv_cache) # two-model or self._should_create_separate_draft_kv_cache()) # one-model draft_kv_cache_config = None - if not estimating_kv_cache and has_draft: + if not estimating_kv_cache and has_independent_draft_cache: # Used when each manager sizes pools from max_gpu_total_bytes (V2 # and V1 VSWA). V1 non-VSWA GPU uses shared max_tokens instead. if self._needs_gpu_kv_cache_budget_split(self_kv_cache_config): @@ -1617,8 +1622,10 @@ def build_managers(self, if draft_kv_cache_config is not None else self_kv_cache_config) - # 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/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b8b67ed5ded8..7b4bebdd7f9b 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -111,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: @@ -158,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 @@ -247,11 +250,18 @@ 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.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) 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 50aa317e9e43..e17ea52f2c90 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -66,7 +66,8 @@ prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) -from ..speculative.drafting_loops import BaseDraftingLoopWrapper +from ..speculative.drafting_loops import (BaseDraftingLoopWrapper, + get_draft_model_capability) from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, @@ -281,8 +282,8 @@ def __init__( dist: Optional[Distributed] = None, spec_config: Optional[DecodingBaseConfig] = None, is_draft_model: bool = False, - drafting_loop_wrapper: Optional[Callable[[torch.nn.Module], - torch.nn.Module]] = None, + drafting_loop_wrapper: Optional[Callable[ + [torch.nn.Module], Optional[torch.nn.Module]]] = None, model: Optional[torch.nn.Module] = None, checkpoint_loader: Optional[BaseCheckpointLoader] = None, model_weights_memory_tag: Optional[str] = None, @@ -431,10 +432,19 @@ 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._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'): @@ -654,9 +664,13 @@ def __init__( self._cuda_graph_padding_enabled = cuda_graph_padding_enabled + cuda_graph_max_total_draft_tokens = ( + self._cuda_graph_external_draft_len + if self._cuda_graph_external_draft_len is not None else + self.original_max_total_draft_tokens) self._cuda_graph_batch_sizes = _filter_cuda_graph_batch_sizes( cuda_graph_batch_sizes, self.batch_size, self.max_num_tokens, - self.original_max_total_draft_tokens, + cuda_graph_max_total_draft_tokens, self._cuda_graph_padding_enabled) if cuda_graph_batch_sizes else [] self._max_cuda_graph_batch_size = (self._cuda_graph_batch_sizes[-1] if @@ -746,7 +760,10 @@ def __init__( # We look up this key in resource_manager during forward to find the # kv cache manager. Can be changed to support multiple model engines # with different KV cache managers. - self.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER if is_draft_model else ResourceManagerType.KV_CACHE_MANAGER + self.kv_cache_manager_key = (ResourceManagerType.DRAFT_KV_CACHE_MANAGER + if is_draft_model + and not self._shares_target_kv_cache else + ResourceManagerType.KV_CACHE_MANAGER) self.lora_model_config: Optional[LoraModelConfig] = None self._trtllm_gen_jit_warmup = False @@ -773,6 +790,7 @@ 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) @@ -1469,8 +1487,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._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)) curr_max_num_tokens = kv_cache_manager.get_num_available_tokens( token_num_upper_bound=token_num_upper_bound, max_num_draft_tokens=self.original_max_draft_len) @@ -2405,10 +2427,15 @@ 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 + 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 = True + 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 f82207881c8f..5023afd51103 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -578,22 +578,44 @@ 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 (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 + 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 @@ -633,8 +655,11 @@ 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 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) diff --git a/tensorrt_llm/_torch/speculative/drafting_loops.py b/tensorrt_llm/_torch/speculative/drafting_loops.py index e2be40bfce5e..8cbd6736a589 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. @@ -9,7 +11,7 @@ """ from abc import ABC, abstractmethod -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from typing import Optional, final import torch @@ -57,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: @@ -124,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], @@ -137,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. @@ -149,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) @@ -202,6 +226,31 @@ 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 a1f4f82869bd..5cdf92b76df0 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 + shares_target_kv_cache: bool = False eagle_choices: Optional[List[List[int]]] = None max_total_draft_tokens: int = 0 @@ -274,11 +283,28 @@ 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 - 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..bffdbafd5833 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 @@ -18,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: @@ -93,6 +97,9 @@ 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 @@ -163,6 +170,28 @@ 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: @@ -223,6 +252,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.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) @@ -230,9 +267,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 +336,10 @@ 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 62189035cafb..e9da38e4f38f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -130,19 +130,23 @@ 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=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, 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) 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..610676911cc0 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,36 @@ 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 51c36071274c..9b25578da6d3 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -3,6 +3,7 @@ import unittest from dataclasses import dataclass +from types import SimpleNamespace from unittest.mock import Mock import torch @@ -12,10 +13,12 @@ from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - _restore_spec_decode_capture_state, _save_spec_decode_capture_state) + CUDAGraphRunner, _restore_spec_decode_capture_state, + _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( - PyTorchModelEngine, _build_request_multimodal_input) + PyTorchModelEngine, _build_request_multimodal_input, + _filter_cuda_graph_batch_sizes) from tensorrt_llm.llmapi.llm_args import TorchLlmArgs # isort: off @@ -28,6 +31,7 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.llmapi import CudaGraphConfig, SamplingParams from tensorrt_llm.mapping import CpType, Mapping @@ -188,6 +192,56 @@ def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None: self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2) + def test_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_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index acd7f2d593f5..ca2e6cfb90cb 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -19,19 +19,23 @@ """ import math +import tempfile import unittest import unittest.mock from copy import deepcopy from typing import TYPE_CHECKING import torch -from transformers import Gemma4Config, Gemma4TextConfig +from transformers import AutoConfig, Gemma4Config, Gemma4TextConfig from tensorrt_llm._torch.attention_backend import FlashInferAttention, FlashInferAttentionMetadata +from tensorrt_llm._torch.configs.gemma4 import Gemma4AssistantConfig 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, + Gemma4AssistantMaskedEmbedder, Gemma4Attention, Gemma4DecoderLayer, Gemma4ForCausalLM, @@ -104,6 +108,27 @@ "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, + "vocab_size_per_layer_input": 0, + }, + "backbone_hidden_size": 256, + "use_ordered_embeddings": True, + "num_centroids": 16, + "centroid_intermediate_top_k": 2, + "tie_word_embeddings": True, + "dtype": "bfloat16", +} + def _make_model_config(config_dict): """Build a ModelConfig from a raw config dict.""" @@ -112,6 +137,13 @@ def _make_model_config(config_dict): return ModelConfig(pretrained_config=cfg, mapping=mapping) +def _make_assistant_model_config(config_dict=GEMMA4_ASSISTANT_CONFIG): + """Build a ModelConfig for a standalone Gemma4 assistant.""" + cfg = Gemma4AssistantConfig(**deepcopy(config_dict)) + mapping = Mapping(world_size=1, tp_size=1, rank=0) + return ModelConfig(pretrained_config=cfg, mapping=mapping) + + class TestGemma4Config(unittest.TestCase): """Tests for Gemma4 config classes.""" @@ -445,6 +477,102 @@ 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_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) + 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) # --------------------------------------------------------------------------- @@ -672,7 +800,13 @@ def test_reject_unrecognized_expert_weights(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) @@ -727,7 +861,7 @@ def _build_gemma4_kv_cache_manager(config, num_blocks=4, tokens_per_block=32, ba # exceeds sliding_window. sliding_window = getattr(config, "sliding_window", None) max_attn_window = None - needs_vswa = isinstance(head_dim, list) and len(set(head_dim)) > 1 + needs_vswa = force_vswa or (isinstance(head_dim, list) and len(set(head_dim)) > 1) if not needs_vswa: needs_vswa = isinstance(num_kv_heads, list) and len(set(num_kv_heads)) > 1 if needs_vswa and sliding_window: @@ -2238,9 +2372,11 @@ def _make_trtllm_gen_decode_case( self, initial_page_counts: list[int], *, + config_dict: dict | None = None, reserved_page_counts: list[int] | None = None, max_pages: int = 64, manager_batch_size: int | None = None, + force_vswa: bool = False, ) -> tuple[ "KVCacheManagerV2", list["FlashInferAttention"], @@ -2256,12 +2392,13 @@ def _make_trtllm_gen_decode_case( if manager_batch_size is None: manager_batch_size = batch_size - config = Gemma4TextConfig(**deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG)) + config = Gemma4TextConfig(**deepcopy(config_dict or GEMMA4_E2B_REAL_DIMS_CONFIG)) kv_cache_manager = self._get_kv_cache_manager( config, num_blocks=max_pages, tokens_per_block=_TRTLLM_GEN_TOKENS_PER_BLOCK, batch_size=manager_batch_size, + force_vswa=force_vswa, ) self.addCleanup(kv_cache_manager.shutdown) self.assertTrue(kv_cache_manager.is_vswa, "Expected VSWA manager") @@ -2380,14 +2517,13 @@ def _prepare_decode_page_counts( def _expected_decode_block_table( self, metadata: "FlashInferAttentionMetadata", - head_dim: int, + pool_id: int, page_counts: list[int], *, rows: int, width: int, ) -> torch.Tensor: """Build the expected compact table from one VSWA pool's host indices.""" - pool_id = metadata._vswa_head_dim_to_pool[head_dim] pool_indices = metadata._host_pool_indices[pool_id].numpy() expected = torch.zeros((rows, width), dtype=torch.int32) source_offset = metadata.num_context_blocks @@ -2422,7 +2558,7 @@ def test_cuda_graph_trtllm_gen_block_table_transitions(self) -> None: block_tables = wrappers.decode_wrapper._block_tables expected = self._expected_decode_block_table( metadata, - plan_params.head_dim, + plan_params.kv_pool_id, new_page_counts, rows=len(initial_page_counts), width=max(initial_page_counts), @@ -2451,11 +2587,12 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N initial_state = {} for plan_params, wrappers in metadata._plan_params_to_wrappers.items(): + self.assertIsNotNone(plan_params.kv_pool_id) block_tables = wrappers.decode_wrapper._block_tables self.assertGreaterEqual(block_tables.size(1), 65) self.assertEqual(block_tables.size(1), metadata.kv_cache_manager.max_blocks_per_seq) self.assertEqual(wrappers.host_decode_block_tables.size(1), 64) - initial_state[plan_params.head_dim] = ( + initial_state[plan_params.kv_pool_id] = ( block_tables.data_ptr(), wrappers.host_decode_block_tables.data_ptr(), ) @@ -2467,13 +2604,13 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N for plan_params, wrappers in metadata._plan_params_to_wrappers.items(): with self.subTest(head_dim=plan_params.head_dim): block_tables = wrappers.decode_wrapper._block_tables - old_device_ptr, old_host_ptr = initial_state[plan_params.head_dim] + old_device_ptr, old_host_ptr = initial_state[plan_params.kv_pool_id] self.assertEqual(block_tables.data_ptr(), old_device_ptr) self.assertNotEqual(wrappers.host_decode_block_tables.data_ptr(), old_host_ptr) self.assertGreaterEqual(wrappers.host_decode_block_tables.size(1), 65) expected = self._expected_decode_block_table( metadata, - plan_params.head_dim, + plan_params.kv_pool_id, new_page_counts, rows=len(new_page_counts), width=max(new_page_counts), @@ -2485,6 +2622,47 @@ def test_cuda_graph_trtllm_gen_host_table_growth_keeps_device_pointer(self) -> N rtol=0, ) + @torch.no_grad() + @unittest.mock.patch( + "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None + ) + def test_cuda_graph_trtllm_gen_distinguishes_same_head_dim_pools(self) -> None: + """Plan keys retain the KV pool when sliding and full head dims match.""" + config_dict = deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG) + config_dict["global_head_dim"] = config_dict["head_dim"] + initial_page_counts = [5, 3] + _, _, metadata, _, _, _ = self._make_trtllm_gen_decode_case( + initial_page_counts, + config_dict=config_dict, + force_vswa=True, + ) + + plan_params = list(metadata._plan_params_to_wrappers) + self.assertEqual({params.head_dim for params in plan_params}, {256}) + self.assertEqual(len({params.kv_pool_id for params in plan_params}), 2) + + new_page_counts = [2, 1] + self._prepare_decode_page_counts(metadata, [0, 1], new_page_counts) + torch.cuda.synchronize() + + for params, wrappers in metadata._plan_params_to_wrappers.items(): + with self.subTest(pool_id=params.kv_pool_id): + expected = self._expected_decode_block_table( + metadata, + params.kv_pool_id, + new_page_counts, + rows=len(new_page_counts), + width=max(new_page_counts), + ) + torch.testing.assert_close( + wrappers.decode_wrapper._block_tables[ + : len(new_page_counts), : max(new_page_counts) + ].cpu(), + expected, + atol=0, + rtol=0, + ) + @torch.no_grad() @unittest.mock.patch( "tensorrt_llm.runtime.kv_cache_manager_v2._utils.assert_critical", lambda *a, **kw: None @@ -2737,6 +2915,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..b6eb66fcb6b1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_gemma4_drafting_loop.py @@ -0,0 +1,147 @@ +# 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", + 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 + + +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_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_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_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 + ) + 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]]))