diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 4f4db5525b60..b1ff9698cc8d 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` | @@ -76,7 +77,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | N/A | Yes | Yes | | `Glm4MoeLiteForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | | `NemotronHForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Untested | Untested | -| `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | +| `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | Yes | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Gemma4UnifiedForConditionalGeneration` | Untested | Untested | Untested | No | Yes | No | No | No | No | Yes | Untested | No | Yes | Untested | Untested | | `Step3p7ForConditionalGeneration`| Yes | Yes | Yes | Untested | Untested | Yes | No | No | No | Yes | Untested | Untested | Yes | Untested | Untested | | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes | Yes | Yes | Untested | Untested | No | No | No | No | Yes | Untested | No | N/A | Untested | Untested | diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index bb06009bdc33..3757d49486d6 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -308,6 +308,8 @@ def setup_llm(args, **kwargs): if spec_decode_algo == 'MTP': if not args.use_one_model: print("Running MTP eagle with two model style.") + speculative_model = (args.draft_model_dir if args.draft_model_dir + is not None else args.model_dir) spec_config = MTPDecodingConfig( max_draft_len=args.spec_decode_max_draft_len, use_relaxed_acceptance_for_thinking=args. @@ -318,7 +320,7 @@ def setup_llm(args, **kwargs): use_dynamic_tree=args.use_dynamic_tree, dynamic_tree_max_topK=args.dynamic_tree_max_topK, max_total_draft_tokens=args.max_total_draft_tokens, - speculative_model=args.model_dir) + speculative_model=speculative_model) elif spec_decode_algo == "EAGLE3": spec_config = Eagle3DecodingConfig( max_draft_len=args.spec_decode_max_draft_len, diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md index 1315755e85b4..721c160aadeb 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,49 @@ curl http://localhost:8000/v1/chat/completions \ The `/v1/chat/completions` endpoint applies the Gemma 4 chat template automatically. +### MTP speculative decoding + +Gemma 4 supports Multi-Token Prediction (MTP) speculative decoding through the +PyTorch execution path. The target loads its matching assistant checkpoint, +and the Q-only assistant reads the target model's KV cache. Create a server +configuration for the target/assistant pair: + +```bash +cat > gemma4_mtp.yaml <<'EOF' +speculative_config: + decoding_type: MTP + max_draft_len: 3 + mtp_eagle_one_model: true + speculative_model: google/gemma-4-E4B-it-assistant +kv_cache_config: + enable_block_reuse: false +EOF + +trtllm-serve google/gemma-4-E4B-it \ + --host 0.0.0.0 \ + --port 8000 \ + --config gemma4_mtp.yaml +``` + +The assistant shares the target model's KV cache, so TensorRT-LLM does not +allocate a second full-size GPU KV cache for it. The current implementation +supports the `Gemma4AssistantForCausalLM` assistants for E2B, E4B, 26B-A4B, +and 31B. Gemma 4 12B uses the `Gemma4UnifiedForConditionalGeneration` and +`Gemma4UnifiedAssistantForCausalLM` architectures, which are not supported. + +For offline inference, pass both checkpoints to the advanced LLM API example: + +```bash +python3 examples/llm-api/quickstart_advanced.py \ + --model_dir google/gemma-4-E4B-it \ + --draft_model_dir google/gemma-4-E4B-it-assistant \ + --spec_decode_algo MTP \ + --spec_decode_max_draft_len 3 \ + --use_one_model \ + --disable_kv_cache_reuse \ + --apply_chat_template +``` + ### Accuracy evaluation with `trtllm-eval` `trtllm-eval` is the canonical entry point for accuracy benchmarks. Two tasks relevant to Gemma 4: diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 733d477009cb..d49ab084a128 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 @@ -165,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. @@ -267,6 +269,18 @@ class FlashInferAttentionMetadata(AttentionMetadata): _multi_item_params: Optional[FlashInferMultiItemParams] = field( init=False, default=None) + _draft_metadata: Optional["FlashInferAttentionMetadata"] = field( + init=False, default=None, 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: @@ -288,6 +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 or self._is_separate_kv_draft_view: + if result._backend != "trtllm-gen": + raise ValueError( + "FlashInfer draft metadata views require the trtllm-gen " + "decode backend.") + num_seqs = self.num_seqs + result._kv_lens_buffer[:num_seqs].copy_( + self._draft_kv_runtime_lens[:num_seqs]) return result def get_ragged_prefill_wrapper( @@ -578,6 +600,238 @@ def batch_indices(self) -> torch.Tensor: def positions(self) -> torch.Tensor: return self._positions[:self.num_tokens] + 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 = 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 + 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 = {} + 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 + 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 + # 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_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.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 + 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._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 + # 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, + 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._draft_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._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) + 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__() self._post_init_with_buffers(self.cuda_graph_buffers) @@ -632,6 +886,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._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') @@ -704,15 +965,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. @@ -765,6 +1017,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._draft_metadata = None # Post init again to make sure all tensors are allocated metadata.__post_init__() return metadata @@ -999,7 +1254,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:], @@ -1275,7 +1531,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 @@ -1331,7 +1588,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 @@ -1339,8 +1595,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 @@ -1408,6 +1663,11 @@ 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 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() @@ -1439,6 +1699,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 708e50e0ac5f..c5893c21bfae 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -1,7 +1,23 @@ +# 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_unified import ( +from tensorrt_llm._torch.configs.gemma4 import ( + Gemma4AssistantConfig, Gemma4UnifiedAudioConfig, Gemma4UnifiedConfig, Gemma4UnifiedTextConfig, @@ -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_unified.py b/tensorrt_llm/_torch/configs/gemma4.py similarity index 56% rename from tensorrt_llm/_torch/configs/gemma4_unified.py rename to tensorrt_llm/_torch/configs/gemma4.py index 561474ffa891..b0b2c17e42cb 100644 --- a/tensorrt_llm/_torch/configs/gemma4_unified.py +++ b/tensorrt_llm/_torch/configs/gemma4.py @@ -12,20 +12,88 @@ # 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). +"""Compatibility configs for Gemma 4 assistant and Unified checkpoints.""" -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. +from transformers import Gemma4TextConfig, PreTrainedConfig -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 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=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 class Gemma4UnifiedTextConfig(Gemma4TextConfig): @@ -39,7 +107,7 @@ class Gemma4UnifiedTextConfig(Gemma4TextConfig): model_type = "gemma4_unified_text" -class Gemma4UnifiedVisionConfig(PretrainedConfig): +class Gemma4UnifiedVisionConfig(PreTrainedConfig): """Sub-config for the encoder-free vision projector.""" model_type = "gemma4_unified_vision" @@ -63,7 +131,7 @@ def __init__( self.rms_norm_eps = rms_norm_eps -class Gemma4UnifiedAudioConfig(PretrainedConfig): +class Gemma4UnifiedAudioConfig(PreTrainedConfig): """Sub-config for the encoder-free audio projector. `output_proj_dims` and `hidden_size` alias `audio_embed_dim` (the raw audio @@ -90,7 +158,7 @@ def __init__( self.hidden_size = hidden_size if hidden_size is not None else audio_embed_dim -class Gemma4UnifiedConfig(PretrainedConfig): +class Gemma4UnifiedConfig(PreTrainedConfig): """Top-level config for Gemma 4 12B Unified (encoder-free multimodal). Parses `config.json` fields required by 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..24aa619b8e34 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -14,8 +14,9 @@ # limitations under the License. """TensorRT-LLM PyTorch backend implementation for Gemma4 text model.""" +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 @@ -40,6 +41,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,9 +56,14 @@ 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, _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( @@ -656,7 +663,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). @@ -1219,11 +1226,8 @@ def forward( return hidden_states -# --------------------------------------------------------------------------- -# 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 +1247,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: @@ -1371,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, @@ -1380,6 +1418,9 @@ 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, + resource_manager=None, + orig_input_ids: Optional[torch.IntTensor] = None, **kwargs, ) -> torch.Tensor: local_attention_mask_data = None @@ -1406,6 +1447,22 @@ 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] + + if self.spec_worker is not None: + return self._forward_speculative( + output, + input_ids, + orig_input_ids, + position_ids, + attn_metadata, + spec_metadata, + resource_manager, + ) + logits = self.logits_processor.forward( output, self.lm_head, @@ -1426,3 +1483,202 @@ 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 + 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, + 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), + ) + + @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) + + @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( + 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.""" + + shares_target_kv_cache = True + + 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 + ) + # 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: "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_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" + 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 + + @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) + + def forward_draft_step( + self, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + recurrent_hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + 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: + 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) + ) + 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 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_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/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 67622709c417..a5d08d4d06b8 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -732,6 +732,17 @@ def post_config(self): self.config = self.llm.config self.model_config.pretrained_config = self.llm.config + @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 +754,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 +783,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 ae3335b187ea..a991e266965d 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1890,6 +1890,12 @@ 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: + 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, model_config.pretrained_config.num_hidden_layers, @@ -1993,6 +1999,20 @@ def __init__(self, model_config.quant_config.kv_cache_quant_algo self.draft_config.extra_attrs = model_config.extra_attrs + elif spec_config._use_shared_kv_cache: + 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 + elif spec_config.spec_dec_mode.is_external_drafter(): self.draft_config = ModelConfig.from_pretrained( model_config.spec_config.speculative_model, @@ -2028,6 +2048,11 @@ 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)): + self.draft_model.load_weights_from_target_model(self) + def forward( self, attn_metadata: AttentionMetadata, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 5d7bd983c8df..2a6ea23076aa 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -416,13 +416,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) @@ -439,6 +432,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 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) + # 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 @@ -553,7 +555,8 @@ def load( loads_draft_weights = ( self.spec_config is not None - and self.spec_config.spec_dec_mode.need_load_draft_weights()) + 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, @@ -704,8 +707,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 loads_draft_weights: weights = checkpoint_loader.load_weights( self.spec_config.speculative_model, mapping=self.mapping) @@ -840,8 +842,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 loads_draft_weights: draft_weights = checkpoint_loader.load_weights( self.spec_config.speculative_model, mapping=self.mapping) @@ -968,8 +969,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 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 9f6c0e5cbef2..f6adba822097 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -461,14 +461,12 @@ def create_py_executor( ) llm_args.disable_overlap_scheduler = True - # Check FLASHINFER compatibility with one-engine speculative decoding - if llm_args.attn_backend == "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"backend instead by setting attn_backend='TRTLLM'.") + 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 diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 09ba663c744d..d2c6e037a85c 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 @@ -10,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 @@ -641,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 @@ -766,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( @@ -814,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_shared_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, @@ -846,7 +947,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) @@ -1008,9 +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: + 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: attn_metadata.num_contexts = 0 if hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( @@ -1268,6 +1374,31 @@ def prepare_1st_drafter_inputs( "spec_metadata": spec_metadata, } + def _prepare_shared_kv_draft_inputs( + self, + *, + 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 + class MTPEagleWorker(Eagle3OneModelWorker): """Backward-compatible alias for ``Eagle3OneModelWorker`` in MTP Eagle mode. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 721b3942ed04..bdb122ac4457 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -115,6 +115,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 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. if spec_config.spec_dec_mode.is_dspark(): @@ -2135,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 @@ -2156,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 @@ -2172,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/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index bffa8833058c..a8d92fed4f46 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 @@ -443,6 +448,8 @@ def get_spec_drafter(model_engine, def get_num_spec_layers(spec_config): + if getattr(spec_config, "_use_shared_kv_cache", False): + return 0 if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): return 1 if spec_config.spec_dec_mode.is_mtp_vanilla(): @@ -519,6 +526,8 @@ def get_num_extra_kv_tokens(spec_config): """ if spec_config is None: return 0 + 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 return 0 @@ -543,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/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d325b69f65d0..a87a25483bf1 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/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py index e512ca0196f7..fe00e3d69d07 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_attention.py +++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py @@ -66,6 +66,93 @@ 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") + 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( + 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 b66eea54cc95..5913dd3a518d 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,13 +27,16 @@ 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, @@ -105,6 +109,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.""" @@ -113,6 +138,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 +563,86 @@ 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(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) + self.assertEqual(config.text_config.num_kv_shared_layers, 4) + + with tempfile.TemporaryDirectory() as directory: + config.save_pretrained(directory) + restored = AutoConfig.from_pretrained(directory) + + self.assertEqual(restored.model_type, "gemma4_assistant") + self.assertEqual(restored.backbone_hidden_size, 256) + self.assertEqual(restored.text_config.num_kv_shared_layers, 4) + + def test_assistant_rejects_partial_kv_sharing(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_input_embeddings, target.model.embed_tokens) + self.assertIsNot(assistant.model.embed_tokens, target.model.embed_tokens) + + # --------------------------------------------------------------------------- # HF reference comparison tests (sub-module + full model) # --------------------------------------------------------------------------- @@ -758,7 +870,12 @@ 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, +): """Create KVCacheManagerV2 supporting Gemma4 per-layer head_dim / kv_heads. Mirrors ``Gemma4Attention``'s layout (global kv heads only for K=V layers) @@ -2466,14 +2583,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 @@ -2485,6 +2601,47 @@ 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(self) -> None: + """The draft view advances lengths without modifying target KV.""" + kv_cache_manager, layers, metadata, queries, _, _ = self._make_trtllm_gen_decode_case( + [3, 2] + ) + 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_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) + torch.testing.assert_close( + draft_metadata._draft_kv_runtime_lens[: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 @@ -2508,7 +2665,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), @@ -2537,11 +2694,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(), ) @@ -2553,13 +2711,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), 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 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", [