From 7c75198684cacc0ff9dc13b9430f04ae923f5bd2 Mon Sep 17 00:00:00 2001 From: varadrane1707 Date: Thu, 18 Jun 2026 06:35:38 +0000 Subject: [PATCH] [None][feat] add Gemma4 MTP speculative decoding support Implements Multi-Token Prediction (MTP) speculative decoding for Gemma4 models using the one-engine path where the assistant's Q-only attention layers share the backbone's KV cache via cache_layer_idx. Key changes: modeling_gemma4.py: - Promote Gemma4ForCausalLM to extend SpecDecOneEngineForCausalLM - Add Gemma4MTPHead, Gemma4MTPDecoderLayer, Gemma4MTP classes - Gemma4MTPDecoderLayer uses is_kv_shared=True (Q-only, no KV projections) - cache_layer_idx computed from backbone layer_types per attention type - Gemma4TextModel gains aux_stream_dict for MTPForCausalLM interface compat modeling_speculative.py: - Add gemma4_text case in MTPForCausalLM to create Gemma4MTP layers - Gemma4 uses max_draft_len (not num_nextn_predict_layers) - load_weights maps HF assistant checkpoint layout to TRT-LLM paths - load_weights_from_target_model shares lm_head/embed_tokens with backbone flashinfer.py: - Mirror spec-dec fields from TrtllmAttentionMetadata onto FlashInferAttentionMetadata so the MTP draft loop can use FlashInfer for Gemma4 (needed for head_dim 256/512 and VSWA on non-Blackwell) py_executor_creator.py: - Detect uses_shared_backbone_kv_for_mtp class attribute to route Gemma4 MTP through the correct KV-cache and backend path model_loader.py: - Return model_cls from load_config_and_apply_defaults for class-level capability checks in py_executor_creator _util.py: - _target_only_kv_layer_mask: restrict main KV cache to backbone layers when draft layers share backbone KV or use a separate draft KV manager speculative/utils.py: - Disable separate draft KV cache for Gemma4 MTP one-engine mode (assistant reads backbone KV via Q-only attention, no separate pool needed) modeling_gemma4mm.py: - Propagate flashinfer_supports_one_engine_spec_decode and uses_shared_backbone_kv_for_mtp from Gemma4ForCausalLM examples/auto_deploy/model_registry/configs/gemma4_dense.yaml: - Update world_size, batch sizes, and token limits for 4-GPU serving Signed-off-by: varadrane1707 --- .../model_registry/configs/gemma4_dense.yaml | 14 +- .../_torch/attention_backend/flashinfer.py | 263 ++++++++++ tensorrt_llm/_torch/models/modeling_gemma4.py | 480 ++++++++++++++++-- .../_torch/models/modeling_gemma4mm.py | 12 +- .../_torch/models/modeling_speculative.py | 122 ++++- tensorrt_llm/_torch/pyexecutor/_util.py | 33 +- .../_torch/pyexecutor/model_loader.py | 6 +- .../_torch/pyexecutor/py_executor_creator.py | 41 +- tensorrt_llm/_torch/speculative/utils.py | 12 + 9 files changed, 911 insertions(+), 72 deletions(-) diff --git a/examples/auto_deploy/model_registry/configs/gemma4_dense.yaml b/examples/auto_deploy/model_registry/configs/gemma4_dense.yaml index 0cba86f38348..6d38b15e10df 100644 --- a/examples/auto_deploy/model_registry/configs/gemma4_dense.yaml +++ b/examples/auto_deploy/model_registry/configs/gemma4_dense.yaml @@ -6,20 +6,24 @@ # paged KV cache, CUDA-graph-compatible, FlashDecoding for decode. model_factory: Gemma4ForConditionalGeneration tokenizer: google/gemma-4-31B-it +world_size: 4 attn_backend: triton compile_backend: torch-cudagraph cuda_graph_config: - batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] -max_num_tokens: 8192 -max_batch_size: 512 -max_seq_len: 8192 + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8] +max_num_tokens: 16000 +max_batch_size: 8 +max_seq_len: 16000 enable_chunked_prefill: true kv_cache_config: enable_block_reuse: false - free_gpu_memory_fraction: 0.8 + free_gpu_memory_fraction: 0.6 transforms: compile_model: piecewise_enabled: true + piecewise_num_tokens: [1, 2, 4, 8, 16, 32, 64, 128, 256, + 512, 768, 1024, 1280, 1536, 1792, 2048, + 2560, 3072, 4096, 5120, 6144, 7168, 8192] mlir_elementwise_fusion: enabled: true gather_logits_before_lm_head: diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index 7ced64dec19a..34b5b5706e57 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -157,6 +157,116 @@ class FlashInferAttentionMetadata(AttentionMetadata): _mla_kv_len_arr_buf: Optional[torch.Tensor] = field(init=False, default=None) + # One-engine speculative decoding (MTP/Eagle3). Mirror TrtllmAttentionMetadata + # so update_spec_dec_param() can allocate/fill buffers on Hopper, where TRTLLM + # MMHA cannot serve Gemma4 global layers (head_dim=512). + is_spec_decoding_enabled: bool = False + use_spec_decoding: bool = False + is_spec_dec_tree: bool = False + is_spec_dec_dynamic_tree: bool = False + max_total_draft_tokens: Optional[int] = None + spec_decoding_position_offsets: Optional[torch.Tensor] = None + spec_decoding_position_offsets_cpp: Optional[torch.Tensor] = None + position_offsets_stride: int = 0 + spec_decoding_packed_mask: Optional[torch.Tensor] = None + spec_decoding_generation_lengths: Optional[torch.Tensor] = None + spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor] = None + spec_decoding_bl_tree_mask: Optional[torch.Tensor] = None + spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None + # MTP/Eagle3 draft loop updates these in place (mirrors TrtllmAttentionMetadata). + host_request_types: torch.Tensor = field(init=False) + kv_lens_cuda: torch.Tensor = field(init=False) + + def is_sm_version_trtllm_gen_kernel(self, sm: int) -> bool: + from .trtllm import TrtllmAttention + return TrtllmAttention.is_sm_version_trtllm_gen_kernel(sm) + + def update_spec_dec_param(self, *args, **kwargs) -> None: + from .trtllm import TrtllmAttentionMetadata + return TrtllmAttentionMetadata.update_spec_dec_param(self, *args, + **kwargs) + + def update_position_offsets_for_cpp(self, query_len: int) -> None: + from .trtllm import TrtllmAttentionMetadata + return TrtllmAttentionMetadata.update_position_offsets_for_cpp( + self, query_len) + + def generate_spec_decoding_generation_length(self, + runtime_draft_len: int) -> None: + from .trtllm import TrtllmAttentionMetadata + return TrtllmAttentionMetadata.generate_spec_decoding_generation_length( + self, runtime_draft_len) + + def update_for_spec_dec(self) -> None: + """Refresh paged-KV views after MTP draft-loop in-place kv_lens_cuda edits.""" + n = self.num_contexts + self.num_generations + if n == 0: + return + kv_lens = self.kv_lens_cuda[:n] + self._cached_token_lens[:n].copy_( + kv_lens - self.seq_lens_kv_cuda[:n]) + num_blocks = ((kv_lens + self.page_size - 1) // self.page_size) + if getattr(self, "num_blocks", None) is not None: + for i in range(n): + self.num_blocks[i] = int(num_blocks[i].item()) + self.num_context_blocks = sum(self.num_blocks[:self.num_contexts]) + self.num_generation_blocks = sum( + self.num_blocks[self.num_contexts:]) + paged_kv_last_page_len = kv_lens - (num_blocks - 1) * self.page_size + self._paged_kv_last_page_len[:n].copy_(paged_kv_last_page_len) + if self.num_contexts == 0 and self.num_generations > 0: + paged_kv_indptr_decode = torch.cumsum( + torch.tensor([0] + self.num_blocks[self.num_contexts:], + dtype=torch.int32), + dim=0, + ) + self.paged_kv_indptr_decode[:paged_kv_indptr_decode.size( + 0)].copy_(paged_kv_indptr_decode, non_blocking=True) + self.paged_kv_indptr = self.paged_kv_indptr_decode[: + paged_kv_indptr_decode + .size(0)] + + def _triton_physical_kv_lens( + self, + start: int, + end: int, + use_spec_dec: bool, + ) -> torch.Tensor: + """Map TRTLLM spec-dec KV lengths to physical paged-KV lengths for Triton.""" + kv_lens = self.kv_lens_cuda[start:end] + if (use_spec_dec + and self.spec_decoding_generation_lengths is not None): + kv_lens = kv_lens - self.spec_decoding_generation_lengths[ + start:end] + return kv_lens + + def _triton_gen_paged_params( + self, + gen_start: int, + gen_end: int, + use_spec_dec: bool, + ): + """Rebuild decode page table views using physical KV lengths.""" + kv_lens = self._triton_physical_kv_lens(gen_start, gen_end, + use_spec_dec) + page_size = self.page_size + num_blocks = ((kv_lens + page_size - 1) // page_size).to(torch.int32) + last_page_len = kv_lens - (num_blocks - 1) * page_size + indptr = torch.zeros(gen_end - gen_start + 1, + dtype=torch.int32, + device=kv_lens.device) + indptr[1:] = torch.cumsum(num_blocks, dim=0) + base_indptr = self.paged_kv_indptr_decode[gen_start:gen_end + 1] + indices_parts = [] + for i in range(gen_end - gen_start): + block_count = int(num_blocks[i].item()) + indices_parts.append( + self._paged_kv_indices[base_indptr[i]:base_indptr[i] + + block_count]) + kv_indices = torch.cat(indices_parts) if indices_parts else ( + self._paged_kv_indices[:0]) + return kv_lens, last_page_len, indptr, kv_indices + def needs_plan(self, plan_params: PlanParams) -> bool: if plan_params not in self._plan_params_to_wrappers: return True @@ -503,6 +613,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.kv_lens_cuda = torch.empty((self.max_num_requests, ), + dtype=torch.int, + device='cuda') + self.host_request_types = torch.empty((self.max_num_requests, ), + dtype=torch.int, + device='cpu', + pin_memory=prefer_pinned()) self._batch_indices = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -780,6 +897,10 @@ def prepare(self) -> None: # number of tokens needed in the kv cache for each sequence after the next pass kv_lens = self.cached_token_lens + self.seq_lens_kv_cuda + n_seqs = self.num_contexts + self.num_generations + self.kv_lens_cuda[:n_seqs].copy_(kv_lens[:n_seqs], non_blocking=True) + self.host_request_types[:self.num_contexts].fill_(0) + self.host_request_types[self.num_contexts:n_seqs].fill_(1) # start and end indices of each sequence in the ragged key and value # for self attention it's the same as qo_indptr so avoid computing twice. @@ -1048,6 +1169,9 @@ def _plan_with_params(self, if not self.needs_plan(plan_params): return plan_params + if flashinfer_backend == "triton": + return plan_params + if self.is_cuda_graph and torch.cuda.is_current_stream_capturing(): raise ValueError( "Cannot plan() for flashinfer kernels while stream is capturing. " @@ -1600,6 +1724,121 @@ def _mla_forward_paged_context( out=output[:num_tokens].view(-1, self.num_heads, self.kv_lora_rank)) + def _forward_triton_paged( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + kv_cache: torch.Tensor, + metadata: FlashInferAttentionMetadata, + output: torch.Tensor, + attention_mask_data: Optional[torch.Tensor], + attention_window_size: Optional[int], + num_contexts: int, + num_generations: int, + num_ctx_tokens: int, + ) -> None: + """Paged attention via Triton kernels (Hopper and pre-Blackwell fallback).""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention.triton_attention import ( + triton_context, + triton_decode, + ) + + from .triton_prefill import triton_prefill_with_custom_mask + + sm_scale = 1 / (math.sqrt(self.head_dim) * self.q_scaling) + sliding_window = attention_window_size + window_left = attention_window_size if attention_window_size is not None else -1 + if metadata.is_cuda_graph: + raise RuntimeError( + "Gemma4 Triton attention (Hopper fallback) is incompatible with " + "decode CUDA graphs. Set cuda_graph_config: null in your serve " + "config, or use --backend _autodeploy with gemma4_dense.yaml." + ) + out_view = output.view(-1, self.num_heads, self.head_dim) + + if num_contexts > 0: + q_ctx = q[:num_ctx_tokens].view(-1, self.num_heads, self.head_dim) + kv_lens = metadata.kv_lens_cuda[:num_contexts] + if attention_mask_data is not None and k is not None: + k_ctx = k[:num_ctx_tokens].view(-1, self.num_kv_heads, self.head_dim) + v_ctx = v[:num_ctx_tokens].view(-1, self.num_kv_heads, self.head_dim) + triton_prefill_with_custom_mask( + q=q_ctx, + k=k_ctx, + v=v_ctx, + output=out_view[:num_ctx_tokens], + qo_indptr=metadata.qo_indptr[:num_contexts + 1], + kv_cache=kv_cache, + prefix_lens=metadata.cached_token_lens[:num_contexts].clone(), + page_table_indptr=metadata.paged_kv_indptr_prefill[:num_contexts + 1], + page_table_indices=metadata._paged_kv_indices[:metadata.num_context_blocks], + page_size=metadata.page_size, + custom_mask=attention_mask_data, + sm_scale=sm_scale, + window_left=window_left, + ) + else: + triton_context( + q=q_ctx, + kv_cache=kv_cache, + qo_indptr=metadata.qo_indptr[:num_contexts + 1], + kv_indptr=metadata.paged_kv_indptr_prefill[:num_contexts + 1], + kv_indices=metadata._paged_kv_indices[:metadata.num_context_blocks], + kv_last_page_len=metadata.paged_kv_last_page_len[:num_contexts], + seq_len_with_cache=kv_lens, + sm_scale=sm_scale, + sliding_window=sliding_window, + out=out_view[:num_ctx_tokens], + ) + + # Generation phase. Do not route head_dim=512 global-attention layers + # through FlashInfer fa2 on SM90 (see flashinfer PR #3652). + if num_generations > 0: + num_gen_tokens = q.shape[0] - num_ctx_tokens + gen_start = num_contexts + gen_end = num_contexts + num_generations + gen_kv_indices = metadata._paged_kv_indices[ + metadata.num_context_blocks: + metadata.num_context_blocks + metadata.num_generation_blocks + ] + # MTP/Eagle3 verification sends multiple Q tokens per generation + # sequence; triton_decode only supports one Q token per sequence. + use_spec_dec = getattr(metadata, 'use_spec_decoding', False) + multi_token_gen = num_gen_tokens > num_generations + if use_spec_dec or multi_token_gen: + q_gen = q[num_ctx_tokens:].view(-1, self.num_heads, self.head_dim) + gen_qo_indptr = metadata.qo_indptr[gen_start:gen_end + 1].clone() + gen_qo_indptr -= gen_qo_indptr[0].item() + kv_lens, last_page_len, kv_indptr, gen_kv_indices = ( + metadata._triton_gen_paged_params(gen_start, gen_end, + use_spec_dec)) + triton_context( + q=q_gen, + kv_cache=kv_cache, + qo_indptr=gen_qo_indptr, + kv_indptr=kv_indptr, + kv_indices=gen_kv_indices, + kv_last_page_len=last_page_len, + seq_len_with_cache=kv_lens, + sm_scale=sm_scale, + sliding_window=sliding_window, + out=out_view[num_ctx_tokens:], + ) + else: + q_dec = q[num_ctx_tokens:num_ctx_tokens + num_generations].view( + num_generations, self.num_heads, self.head_dim) + triton_decode( + q=q_dec, + kv_cache=kv_cache, + kv_indices=gen_kv_indices, + kv_indptr=metadata.paged_kv_indptr_decode[:num_generations + 1], + kv_last_page_len=metadata._paged_kv_last_page_len[gen_start:gen_end], + sm_scale=sm_scale, + sliding_window=sliding_window, + out=out_view[num_ctx_tokens:num_ctx_tokens + num_generations], + ) + def forward_impl( self, q: torch.Tensor, @@ -1748,6 +1987,30 @@ def forward_impl( num_contexts = metadata.num_contexts num_generations = metadata.num_generations num_ctx_tokens = metadata.num_ctx_tokens + use_spec_dec = getattr(metadata, 'use_spec_decoding', False) + num_gen_tokens = q.shape[0] - num_ctx_tokens + multi_token_gen = num_generations > 0 and num_gen_tokens > num_generations + + # Hopper Gemma4 routes head_dim=512 layers through Triton. Sliding + # layers (head_dim=256) use fa2 for single-token decode, but MTP/Eagle3 + # verification sends multiple Q tokens per generation sequence; fa2 + # batch decode rejects that, so fall back to Triton multi-token context. + if (self.flashinfer_backend == "triton" or use_spec_dec + or multi_token_gen): + self._forward_triton_paged( + q=q, + k=k, + v=v, + kv_cache=kv_cache, + metadata=metadata, + output=output, + attention_mask_data=attention_mask_data, + attention_window_size=attention_window_size, + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + ) + return def prefill_forward(plan_params: PlanParams, out: torch.Tensor): wrapper = metadata.get_prefill_wrapper(plan_params) diff --git a/tensorrt_llm/_torch/models/modeling_gemma4.py b/tensorrt_llm/_torch/models/modeling_gemma4.py index e5a1c1945476..44ba2269b1f2 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4.py @@ -15,7 +15,7 @@ """TensorRT-LLM PyTorch backend implementation for Gemma4 text model.""" import math -from typing import Dict, Optional, Tuple +from typing import Dict, List, Optional, Tuple import torch import torch.nn.functional as F @@ -23,6 +23,7 @@ from packaging.version import Version from torch import nn +from tensorrt_llm._utils import get_sm_version from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper from tensorrt_llm._torch.modules.fused_moe.create_moe import create_moe from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode @@ -46,7 +47,10 @@ from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ..modules.rms_norm import RMSNorm -from ..utils import ActivationType +from ..distributed import allgather +from ..speculative import SpecMetadata +from ..utils import ActivationType, create_lm_head_tp_mapping +from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0" @@ -60,6 +64,21 @@ from transformers import Gemma4TextConfig # noqa: E402 +def _gemma4_flashinfer_backend(head_dim: int) -> str: + """Pick the FlashInfer sub-backend for Gemma4 hybrid attention. + + trtllm-gen FMHA cubins require Blackwell (SM100/SM103). On Hopper, + FlashInfer fa2 handles sliding layers (head_dim=256); global layers + (head_dim=512) use the Triton paged-attention path because fa2 rejects + head_dim=512 on SM90 (flashinfer PR #3652). + """ + if get_sm_version() in (100, 103): + return "trtllm-gen" + if head_dim > 256: + return "triton" + return "fa2" + + # --------------------------------------------------------------------------- # Scaled embedding (reused from Gemma3 pattern) # --------------------------------------------------------------------------- @@ -260,14 +279,10 @@ def __init__( # the rotate split, matching HF's rotate_half(head_dim//2) pairing. self.rotary_emb.head_dim = layer_head_dim - # Use trtllm-gen for ALL layers. trtllm-gen has pre-compiled cubins - # for both H256+SWA and H512 across all supported dtypes. - # For FP8 KV cache (NVFP4), Q is also cast to FP8 in the FlashInfer - # backend so that QkvE4m3OBfloat16 context cubins can be used - # (context cubins require same Q/KV dtype; decode cubins support - # mixed dtypes natively). Uniform backend avoids workspace - # corruption between different wrapper types under CUDA graphs. - self.attn.flashinfer_backend = "trtllm-gen" + # trtllm-gen has pre-compiled cubins for H256+SWA and H512 on + # Blackwell (SM100/SM103). On Hopper and earlier, fall back to the + # Triton paged-attention path (head_dim=512, VSWA, CUDA-graph safe). + self.attn.flashinfer_backend = _gemma4_flashinfer_backend(layer_head_dim) # KV shared layers: use target layer's index for KV cache access # so the attention backend reads from the target layer's cache slot. @@ -723,6 +738,17 @@ def __init__(self, model_config: ModelConfig[Gemma4TextConfig]): pretrained = config.pretrained_config self.hidden_size = pretrained.hidden_size + # Empty aux_stream_dict for interface compatibility with MTPForCausalLM. + # Gemma4 does not use auxiliary CUDA streams within the backbone model. + from ..utils import AuxStreamType + self.aux_stream_dict = { + AuxStreamType.Attention: torch.cuda.Stream(), + AuxStreamType.MoeShared: torch.cuda.Stream(), + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + AuxStreamType.MoeBalancer: torch.cuda.Stream(), + AuxStreamType.MoeOutputMemset: torch.cuda.Stream(), + } + # Under AttentionDP, each rank runs the full sequence locally so the # embedding must be replicated (tp_size=1 effectively); otherwise # embed_tokens is COLUMN-sharded by vocab. This mirrors the @@ -908,7 +934,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], @@ -930,22 +956,38 @@ def __init__( super().__init__( Gemma4TextModel(model_config), - config=model_config, - hidden_size=model_config.pretrained_config.hidden_size, - vocab_size=model_config.pretrained_config.vocab_size, + model_config, ) + @classmethod + def flashinfer_supports_one_engine_spec_decode(cls) -> bool: + """Whether FLASHINFER can serve one-engine speculative decoding on this GPU. + + Native FlashInfer batch decode assumes one query token per generation + sequence. Gemma4 Hopper uses Triton multi-token context for spec-dec + verification on fa2 sliding layers and for all head_dim=512 layers. + """ + return get_sm_version() not in (100, 103) + + uses_shared_backbone_kv_for_mtp = True + @classmethod def get_model_defaults(cls, llm_args) -> dict: """Gemma4-specific defaults. FlashInfer backend is required for hybrid attention (per-layer - head_dim 256/512 with VSWA), trtllm-gen cubin dispatch, and - bidirectional attention masks for multimodal tokens. + head_dim 256/512 with VSWA) and bidirectional attention masks for + multimodal tokens. trtllm-gen cubins are used on Blackwell; Hopper + uses the Triton paged-attention sub-backend instead. + + On Hopper, decode CUDA graphs must stay disabled: the Triton prefill/ + decode fallback is not compatible with the FlashInfer-oriented decode + CUDA graph capture path and produces corrupted generation if enabled. """ - return { - "attn_backend": "FLASHINFER", - } + defaults = {"attn_backend": "FLASHINFER"} + if get_sm_version() not in (100, 103): + defaults["cuda_graph_config"] = None + return defaults def _get_token_type_mask(self, mm_token_type_ids: torch.Tensor): """Build bidirectional attention mask from mm_token_type_ids. @@ -1044,36 +1086,46 @@ def get_flashinfer_attention_mask( context_mask_list.append(mask_i.flatten()) return torch.cat(context_mask_list, dim=0).contiguous() - @torch.inference_mode() - def forward( + def _build_attention_masks( self, + mm_token_type_ids: Optional[torch.Tensor], 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, - mm_token_type_ids: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - local_attention_mask_data = None - global_attention_mask_data = None - # Only build bidirectional masks when use_bidirectional_attention is - # set to "vision" (26B, 31B). E2B/E4B have this as None and should - # use standard causal attention even for multimodal tokens. + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Build local and global bidirectional attention masks for multimodal tokens.""" + local_mask = None + global_mask = None use_bidir = getattr(self.config, "use_bidirectional_attention", None) if mm_token_type_ids is not None and use_bidir == "vision": - global_attention_mask_data = self.get_flashinfer_attention_mask( + global_mask = self.get_flashinfer_attention_mask( mm_token_type_ids=mm_token_type_ids, attn_metadata=attn_metadata, effective_sliding_window=None, ) - local_attention_mask_data = self.get_flashinfer_attention_mask( + local_mask = self.get_flashinfer_attention_mask( mm_token_type_ids=mm_token_type_ids, attn_metadata=attn_metadata, effective_sliding_window=self.config.sliding_window, ) + return local_mask, global_mask + + @torch.inference_mode() + 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, + mm_token_type_ids: Optional[torch.Tensor] = None, + spec_metadata=None, + resource_manager=None, + **kwargs, + ) -> torch.Tensor: + local_attention_mask_data, global_attention_mask_data = ( + self._build_attention_masks(mm_token_type_ids, attn_metadata) + ) - output = self.model( + hidden_states = self.model( input_ids=input_ids, attn_metadata=attn_metadata, position_ids=position_ids, @@ -1081,26 +1133,372 @@ def forward( local_attention_mask_data=local_attention_mask_data, global_attention_mask_data=global_attention_mask_data, ple_input_ids=kwargs.pop("ple_input_ids", None), + spec_metadata=spec_metadata, **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, hidden_states) + + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[:attn_metadata.num_tokens] + + if self.spec_worker is not None: + logits = self.logits_processor.forward( + hidden_states[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 + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if input_ids is not None: + spec_input_ids = input_ids[:attn_metadata.num_tokens] + if position_ids is not None: + from .modeling_speculative import _slice_spec_position_ids + 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=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + logits = self.logits_processor.forward( - output, + hidden_states, self.lm_head, attn_metadata, return_context_logits, ) - # Logit softcapping if self.config.final_logit_softcapping is not None: cap = self.config.final_logit_softcapping logits = torch.tanh(logits / cap) * cap return logits - def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper): - weights = weight_mapper.preprocess_weights(weights) - super().load_weights(weights, weight_mapper) + def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper = None, **kwargs): + if weight_mapper is not None: + weights = weight_mapper.preprocess_weights(weights) + super().load_weights(weights=weights, weight_mapper=weight_mapper, **kwargs) # 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() + + +# --------------------------------------------------------------------------- +# Gemma4 MTP (Multi-Token Prediction) support +# --------------------------------------------------------------------------- + +class Gemma4MTPHead(nn.Module): + """Final RMSNorm + lm_head projection for Gemma4 MTP output. + + Called by the MTP worker as ``mtp_layer.shared_head(hidden_states, lm_head, attn_metadata)`` + after ``mtp_layer.forward`` has already applied the norm and returned the normalised + hidden states. This forward therefore only selects the last-token states per sequence + and computes logits via the shared ``lm_head``. + """ + + def __init__(self, model_config: ModelConfig[Gemma4TextConfig]): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.mapping_lm_head_tp = None + + @torch.compile(options={"max-autotune": True}) + def get_last_token_states(self, 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, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: AttentionMetadata, + return_context_logits: bool = False) -> torch.Tensor: + if not return_context_logits: + if attn_metadata is not None: + hidden_states = self.get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1].unsqueeze(0) + + enable_attention_dp = self.model_config.mapping.enable_attention_dp + enable_lm_head_tp_in_adp = (enable_attention_dp + and self.model_config.mapping.enable_lm_head_tp_in_adp) + + if enable_lm_head_tp_in_adp: + self.mapping_lm_head_tp = create_lm_head_tp_mapping( + self.model_config.mapping, hidden_states.shape[0]) + hidden_states = allgather(hidden_states, self.mapping_lm_head_tp, dim=0) + + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = False + target_dtype = lm_head.weight.dtype + if hidden_states.dtype != target_dtype: + hidden_states = hidden_states.to(target_dtype) + logits = lm_head(hidden_states, + mapping_lm_head_tp=self.mapping_lm_head_tp, + is_spec_decoding_head=True) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = True + return logits + + +class Gemma4MTPDecoderLayer(DecoderLayer): + """Single decoder layer for Gemma4 MTP with Q-only (KV-shared) attention. + + Unlike ``Gemma4DecoderLayer``, this layer: + - Always uses ``is_kv_shared=True`` (no K/V projections) + - Accepts explicit ``is_sliding`` and ``cache_layer_idx`` parameters + - Skips PLE / MoE features (not present in the assistant model) + """ + + def __init__( + self, + model_config: ModelConfig[Gemma4TextConfig], + layer_idx: int, + is_sliding: bool, + cache_layer_idx: int, + ) -> None: + super().__init__() + config = model_config.pretrained_config + + self.self_attn = Gemma4Attention( + model_config, + layer_idx=layer_idx, + is_sliding=is_sliding, + is_kv_shared=True, + cache_layer_idx=cache_layer_idx, + ) + + mlp_tp_size = 1 if model_config.mapping.enable_attention_dp else None + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + bias=False, + activation=gelu_tanh, + dtype=config.torch_dtype, + config=model_config, + layer_idx=layer_idx, + overridden_tp_size=mlp_tp_size, + ) + + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.pre_feedforward_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.post_feedforward_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.register_buffer( + "layer_scalar", + torch.ones(1, dtype=config.torch_dtype), + ) + + @torch.inference_mode() + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + target_dtype = self.input_layernorm.weight.dtype + if hidden_states.dtype != target_dtype: + hidden_states = hidden_states.to(target_dtype) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + **kwargs, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = hidden_states * self.layer_scalar + return hidden_states + + +class Gemma4MTP(DecoderLayer): + """Gemma4 MTP (Multi-Token Prediction) layer for one-engine speculative decoding. + + Implements the Gemma4 assistant model architecture: + 1. Concatenate ``embed_tokens(input_ids)`` with backbone ``hidden_states`` + 2. Project to MTP hidden size via ``pre_projection`` + 3. Run through ``Gemma4MTPDecoderLayer`` blocks (Q-only, KV shared with backbone) + 4. Apply final ``norm`` + 5. Project back to backbone hidden size via ``post_projection`` + + The ``cache_layer_idx`` for each decoder sub-layer is computed from the backbone + config: it points to the last backbone layer with its own KV cache of the + matching attention type (sliding or full). + + Args: + model_config: Backbone model config (``Gemma4TextConfig``). + layer_idx: Absolute layer index starting at ``backbone.num_hidden_layers``. + aux_stream_dict: Auxiliary CUDA stream dict (unused; accepted for interface compat). + mtp_layer_types: Optional list of layer type strings for the MTP decoder sub-layers. + When ``None``, inferred by extending the backbone's ``layer_types`` pattern. + """ + + def __init__( + self, + model_config: ModelConfig[Gemma4TextConfig], + layer_idx: int, + aux_stream_dict: Dict, + mtp_layer_types: Optional[List[str]] = None, + ) -> None: + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + backbone_hidden_size = config.hidden_size + + # --- pre_projection: concat([embed, hidden]) -> mtp_hidden --- + # For Gemma4, mtp_hidden_size == backbone_hidden_size; the projection + # maps 2*H -> H so the concatenation is collapsed to backbone dim. + if model_config.mapping.enable_attention_dp: + self.pre_projection = Linear( + backbone_hidden_size * 2, + backbone_hidden_size, + bias=False, + dtype=config.torch_dtype, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + else: + self.pre_projection = Linear( + backbone_hidden_size * 2, + backbone_hidden_size, + bias=False, + dtype=config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + + # Compute which backbone layer each MTP sub-layer should share KV with. + num_kv_shared = getattr(config, "num_kv_shared_layers", 0) + first_kv_shared_idx = config.num_hidden_layers - num_kv_shared + + # Last backbone layer with own KV per attention type. + _last_own_kv: Dict[str, int] = {} + for i in range(first_kv_shared_idx - 1, -1, -1): + t = config.layer_types[i] + if t not in _last_own_kv: + _last_own_kv[t] = i + + # Infer MTP sub-layer types when not provided. + if mtp_layer_types is None: + # Continue the backbone alternating pattern beyond num_hidden_layers. + mtp_layer_types = [ + config.layer_types[layer_idx % config.num_hidden_layers] + ] + + # Build MTP decoder sub-layers. + self.mtp_layers = nn.ModuleList() + for sub_i, layer_type in enumerate(mtp_layer_types): + is_sliding = layer_type == "sliding_attention" + # Fall back to last backbone layer if no own-KV layer of this type exists. + cache_layer_idx = _last_own_kv.get( + layer_type, + first_kv_shared_idx - 1 if first_kv_shared_idx > 0 else 0, + ) + self.mtp_layers.append( + Gemma4MTPDecoderLayer( + model_config, + layer_idx=layer_idx + sub_i, + is_sliding=is_sliding, + cache_layer_idx=cache_layer_idx, + ) + ) + + self.shared_head = Gemma4MTPHead(model_config) + + @torch.inference_mode() + def forward( + self, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + embed_tokens, + attn_metadata: AttentionMetadata, + all_rank_num_tokens: Optional[List[int]] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + # 1. Embed input tokens (backbone's scaled embedding). + inputs_embeds = embed_tokens(input_ids) + + target_dtype = self.pre_projection.weight.dtype + if hidden_states.dtype != target_dtype: + hidden_states = hidden_states.to(target_dtype) + + # 2. Concatenate embeddings with backbone hidden states and project. + hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1) + + # Split for ROW-parallel pre_projection when TP>1 (no ADP). + tp_size = self.model_config.mapping.tp_size + tp_rank = self.model_config.mapping.tp_rank + if tp_size > 1 and not self.model_config.mapping.enable_attention_dp: + hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank] + + hidden_states = self.pre_projection(hidden_states) + + # 3. Run through MTP decoder sub-layers. + for mtp_layer in self.mtp_layers: + hidden_states = mtp_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + **kwargs, + ) + + # 4. Apply final norm. + hidden_states = self.shared_head.norm(hidden_states) + + # 5. Capture hidden states for the spec worker. + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(0, hidden_states, None) + + return hidden_states diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 38e33d11069a..288b0d5bc6d7 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -585,12 +585,18 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): - mm_token_type_ids-based bidirectional masking """ + @classmethod + def flashinfer_supports_one_engine_spec_decode(cls) -> bool: + from .modeling_gemma4 import Gemma4ForCausalLM + return Gemma4ForCausalLM.flashinfer_supports_one_engine_spec_decode() + + uses_shared_backbone_kv_for_mtp = True + @classmethod def get_model_defaults(cls, llm_args) -> dict: """Gemma4-specific defaults — see Gemma4ForCausalLM.get_model_defaults.""" - return { - "attn_backend": "FLASHINFER", - } + from .modeling_gemma4 import Gemma4ForCausalLM + return Gemma4ForCausalLM.get_model_defaults(llm_args) def _check_and_adjust_experts_implementation(self, *args, **kwargs): # transformers 5.x ``PreTrainedModel.__init__`` calls this with an diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 62683b3f62f2..fe0faea51d32 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1437,13 +1437,37 @@ def __init__( case "step3p7" | "step3p5": from .modeling_step3p7 import Step3p7MTP mtp_layer = Step3p7MTP + case "gemma4_text": + from .modeling_gemma4 import Gemma4MTP + mtp_layer = Gemma4MTP case _: raise ValueError( f"Model type {model_type} not supported for MTP") spec_dec_mode = model_config.spec_config.spec_dec_mode assert spec_dec_mode.is_mtp_one_model() - checkpoint_mtp_num_layers = model_config.pretrained_config.num_nextn_predict_layers + + # Gemma4 uses a separate assistant model checkpoint; num_nextn_predict_layers + # is not defined in its config. Use max_draft_len as the layer count and + # load the assistant config to determine internal decoder layer types. + mtp_kwargs: dict = {} + if model_type == "gemma4_text": + checkpoint_mtp_num_layers = model_config.spec_config.max_draft_len + _assistant_layer_types = None + _speculative_model = getattr(model_config.spec_config, "speculative_model", None) + if _speculative_model: + try: + from transformers import Gemma4AssistantConfig + _asst_cfg = Gemma4AssistantConfig.from_pretrained(_speculative_model) + if _asst_cfg.text_config is not None: + _assistant_layer_types = _asst_cfg.text_config.layer_types + except Exception: + pass + if _assistant_layer_types is not None: + mtp_kwargs["mtp_layer_types"] = _assistant_layer_types + else: + checkpoint_mtp_num_layers = model_config.pretrained_config.num_nextn_predict_layers + if spec_dec_mode.is_mtp_eagle_one_model(): mtp_num_layers = 1 mtp_repeat_count = model_config.spec_config.max_draft_len @@ -1456,11 +1480,105 @@ def __init__( self.mtp_layers = nn.ModuleList([ mtp_layer(model_config, layer_idx + start_layer_idx, - model.aux_stream_dict) + model.aux_stream_dict, **mtp_kwargs) for layer_idx in range(mtp_num_layers) ]) self.lm_head = lm_head self.embed_tokens = model.embed_tokens + self._model_type = model_type + + def load_weights(self, weights: Dict, weight_mapper=None): + """Load MTP weights from a draft/assistant model checkpoint. + + Currently only needed for Gemma4 whose assistant is a separate checkpoint. + For other models, MTP weights are embedded in the backbone checkpoint and + loaded as part of the backbone's ``load_weights``. + """ + if self._model_type != "gemma4_text": + return + + # Map from HF Gemma4 assistant checkpoint layout to TRT-LLM module paths. + # HF layout: + # pre_projection.weight + # post_projection.weight (not used directly; backbone lm_head shared) + # model.norm.weight + # model.layers.{i}.self_attn.q_proj.weight + # model.layers.{i}.self_attn.q_norm.weight + # model.layers.{i}.self_attn.o_proj.weight + # model.layers.{i}.mlp.* + # model.layers.{i}.input_layernorm.weight + # model.layers.{i}.post_attention_layernorm.weight + # model.layers.{i}.pre_feedforward_layernorm.weight + # model.layers.{i}.post_feedforward_layernorm.weight + # model.layers.{i}.layer_scalar + # TRT-LLM layout (for each mtp_layers[k]): + # mtp_layers.{k}.pre_projection.* + # mtp_layers.{k}.shared_head.norm.* + # mtp_layers.{k}.mtp_layers.{j}.self_attn.qkv_proj.* (Q-only) + # mtp_layers.{k}.mtp_layers.{j}.self_attn.q_norm.* + # mtp_layers.{k}.mtp_layers.{j}.self_attn.o_proj.* + # mtp_layers.{k}.mtp_layers.{j}.mlp.* + # mtp_layers.{k}.mtp_layers.{j}.{norm}.* + # mtp_layers.{k}.mtp_layers.{j}.layer_scalar + # + # Each MTPForCausalLM instance (mtp_layers[k]) shares the same assistant + # checkpoint weights (the assistant model is called repeatedly for each + # draft step, not a new checkpoint per step). + + # Build mapping: assistant checkpoint key -> TRT-LLM parameter path + remap: Dict[str, str] = {} + for k in range(len(self.mtp_layers)): + prefix = f"mtp_layers.{k}." + remap[f"pre_projection.weight"] = f"{prefix}pre_projection.weight" + remap[f"model.norm.weight"] = f"{prefix}shared_head.norm.weight" + # Decoder sub-layers + num_sub = len(self.mtp_layers[k].mtp_layers) + for j in range(num_sub): + sub_prefix = f"{prefix}mtp_layers.{j}." + hf_sub = f"model.layers.{j}." + for sfx in [ + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + "pre_feedforward_layernorm.weight", + "post_feedforward_layernorm.weight", + "layer_scalar", + ]: + remap[hf_sub + sfx] = sub_prefix + sfx + # Q-only projection: q_proj -> qkv_proj + remap[hf_sub + "self_attn.q_proj.weight"] = ( + sub_prefix + "self_attn.qkv_proj.weight" + ) + + # Load into module via state_dict update (allow missing/unexpected). + state_dict = self.state_dict() + updated: set = set() + for hf_key, tensor in weights.items(): + trt_key = remap.get(hf_key) + if trt_key is not None and trt_key in state_dict: + state_dict[trt_key] = tensor + updated.add(trt_key) + + missing = set(state_dict.keys()) - updated - { + k for k in state_dict if "lm_head" in k or "embed_tokens" in k + } + if missing: + from tensorrt_llm.logger import logger + logger.warning( + f"Gemma4 MTP: {len(missing)} parameters not found in assistant " + f"checkpoint (e.g. {next(iter(missing))}). " + "This is expected if parameters are shared with the backbone." + ) + self.load_state_dict(state_dict, strict=False) + + def load_weights_from_target_model(self, target_model: torch.nn.Module) -> None: + """Share lm_head and embed_tokens with the backbone model.""" + self.lm_head = target_model.lm_head + self.embed_tokens = target_model.model.embed_tokens class MTPDraftModel(nn.Module): diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f4189f22ed92..331c41f8ec7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -776,14 +776,10 @@ def _create_kv_cache_manager( kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( model_engine, kv_cache_config) - # When using separate draft KV cache in one-model speculative decoding, - # use layer_mask to include only target layers. The draft layers should - # only be in the separate draft KV cache manager. + # When draft layers share backbone KV or use a separate draft manager, + # restrict the main KV cache to backbone layers only. # We still pass spec_config so that num_extra_kv_tokens is calculated. - spec_dec_layer_mask = None - if self._should_create_separate_draft_kv_cache(): - num_target_layers = model_engine.model.model_config.pretrained_config.num_hidden_layers - spec_dec_layer_mask = [True] * num_target_layers + spec_dec_layer_mask = self._target_only_kv_layer_mask(model_engine) estimating_kv_cache = estimating_kv_cache and not self._skip_est kv_cache_manager = _create_kv_cache_manager( @@ -846,6 +842,29 @@ def _should_create_separate_draft_kv_cache(self) -> bool: return False return should_use_separate_draft_kv_cache(self._speculative_config) + def _uses_shared_backbone_kv_for_mtp(self, model_engine) -> bool: + if self._speculative_config is None: + return False + if not self._speculative_config.spec_dec_mode.is_mtp_eagle_one_model(): + return False + return getattr(type(model_engine.model), + "uses_shared_backbone_kv_for_mtp", False) + + def _target_only_kv_layer_mask( + self, model_engine: PyTorchModelEngine) -> Optional[List[bool]]: + """Layer mask covering only backbone layers in the main KV cache manager. + + Used when draft MTP layers share backbone KV (Gemma4) or live in a + separate draft KV cache manager. Prevents get_pp_layers() from + appending speculative layer slots that have no per-layer KV metadata. + """ + if not (self._should_create_separate_draft_kv_cache() + or self._uses_shared_backbone_kv_for_mtp(model_engine)): + return None + num_target_layers = ( + model_engine.model.model_config.pretrained_config.num_hidden_layers) + return [True] * num_target_layers + def _get_effective_draft_config(self) -> ModelConfig: """ Return the ModelConfig to use for draft KV cache creation. diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 6df540b88487..561ba62b9d34 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -307,10 +307,10 @@ def __init__(self, @staticmethod def load_config_and_apply_defaults( checkpoint_dir: str, llm_args: TorchLlmArgs, - checkpoint_loader: BaseCheckpointLoader) -> TorchLlmArgs: + checkpoint_loader: BaseCheckpointLoader) -> tuple[TorchLlmArgs, type | None]: """Load model config and apply model-specific defaults to llm_args.""" if checkpoint_loader is None: - return llm_args + return llm_args, None config_kwargs = { 'trust_remote_code': llm_args.trust_remote_code, @@ -344,7 +344,7 @@ def load_config_and_apply_defaults( f"Applied model defaults for {model_cls.__name__}: {applied_defaults}" ) - return llm_args + return llm_args, model_cls @staticmethod def _needs_source_identity(checkpoint_loader: BaseCheckpointLoader, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 797f2fd48666..e04c042c4f87 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -235,9 +235,9 @@ def _load_config_and_create_checkpoint_loader( mx_config=llm_args.mx_config, mx_model_name=llm_args.model, ) - llm_args = ModelLoader.load_config_and_apply_defaults( + llm_args, model_cls = ModelLoader.load_config_and_apply_defaults( checkpoint_dir, llm_args, checkpoint_loader) - return llm_args, checkpoint_loader + return llm_args, checkpoint_loader, model_cls def create_encoder_executor( @@ -259,7 +259,7 @@ def create_encoder_executor( """ from .encoder_executor import EncoderExecutor - llm_args, checkpoint_loader = _load_config_and_create_checkpoint_loader( + llm_args, checkpoint_loader, _ = _load_config_and_create_checkpoint_loader( llm_args, checkpoint_dir) mapping = _get_mapping(llm_args.parallel_config.to_mapping()) @@ -324,7 +324,7 @@ def create_py_executor( """ skip_est = os.environ.get("TRTLLM_SKIP_KV_CACHE_ESTIMATION", '0') == '1' - llm_args, checkpoint_loader = _load_config_and_create_checkpoint_loader( + llm_args, checkpoint_loader, model_cls = _load_config_and_create_checkpoint_loader( llm_args, checkpoint_dir) garbage_collection_gen0_threshold = llm_args.garbage_collection_gen0_threshold @@ -399,6 +399,12 @@ def create_py_executor( from tensorrt_llm._torch.speculative import suggest_spec_config spec_config = suggest_spec_config(max_batch_size) + if (spec_config is not None + and spec_config.spec_dec_mode.is_mtp_eagle_one_model() + and model_cls is not None + and getattr(model_cls, "uses_shared_backbone_kv_for_mtp", False)): + spec_config._allow_separate_draft_kv_cache = False + if not llm_args.disable_overlap_scheduler and spec_config is not None: if not spec_config.spec_dec_mode.support_overlap_scheduler(): logger.warning( @@ -406,14 +412,27 @@ def create_py_executor( ) llm_args.disable_overlap_scheduler = True - # Check FLASHINFER compatibility with one-engine speculative decoding - if llm_args.attn_backend == "FLASHINFER": + if spec_config is not None and spec_config.spec_dec_mode.use_one_engine(): + gemma4_shared_kv_mtp = ( + model_cls is not None + and getattr(model_cls, "uses_shared_backbone_kv_for_mtp", False)) + if (llm_args.attn_backend == "TRTLLM" and gemma4_shared_kv_mtp + and get_sm_version() not in (100, 103)): + raise ValueError( + "TRTLLM attention does not support Gemma4 global attention layers " + "(head_dim=512) on Hopper: MMHA rejects head_dim 512. Use " + "attn_backend='FLASHINFER' for Gemma4 MTP on H100.") + if llm_args.attn_backend == "FLASHINFER" and not gemma4_shared_kv_mtp: 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'.") + f"FLASHINFER attention backend is not supported with one-engine " + f"speculative decoding mode '{spec_config.spec_dec_mode.name}'. " + f"Please use attn_backend='TRTLLM'.") + elif llm_args.attn_backend == "FLASHINFER" and gemma4_shared_kv_mtp: + logger.info( + "Using FLASHINFER for Gemma4 one-engine MTP: Triton paged " + "attention supports head_dim=512 with spec-dec metadata on " + "FlashInferAttentionMetadata." + ) if mm_encoder_only: llm_args.mm_encoder_only = True diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 91f60243834c..f9ca231386b4 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -441,6 +441,18 @@ def update_spec_config_from_model_config(spec_config, model_config): spec_config.max_total_draft_tokens = spec_config.max_draft_len + # Gemma4 assistant MTP layers use Q-only attention and read backbone KV via + # cache_layer_idx. A separate draft KV pool is unnecessary and breaks + # hybrid per-layer head layout (full backbone layer_types vs 1 draft slot). + from ..pyexecutor.config_utils import is_gemma4_hybrid + if (is_gemma4_hybrid(model_config) + and spec_config.spec_dec_mode.is_mtp_eagle_one_model()): + spec_config._allow_separate_draft_kv_cache = False + logger.info( + "Gemma4 MTP: disabled separate draft KV cache (assistant shares " + "backbone KV via Q-only attention)." + ) + @dataclass class SpecDecodingTensor: