diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py index 9df138259b57..ba8577da9613 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py +++ b/tensorrt_llm/_torch/models/modeling_exaone_moe.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + import math import os from typing import Dict, List, Optional, Tuple diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 62b24956d077..ddc2c5687879 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -402,53 +402,65 @@ def forward( inputs_embeds: Optional[torch.FloatTensor] = None, spec_metadata: Optional[SpecMetadata] = None, hidden_states: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[List[int]] = None, ) -> torch.Tensor: - assert self.embed_tokens is not None - - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError( - "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" - ) + # When ``all_rank_num_tokens`` is supplied the caller wants this draft + # forward to run with a different attention-DP token distribution + # (e.g. the worker's per-step value); restore the original on exit so + # the next call sees the same attn_metadata it had on entry. + previous_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids).to(self.dtype) - - assert hidden_states is not None - # NOTE: If hidden states from the target model have to be concatenated, - # ideally, we expect that to happen outside the model definition. This - # helps us avoid data-dependent control flow and gives us better CUDA - # graph coverage. - if self._eh_proj_before_attn: - input_embeds = self.enorm(inputs_embeds) - hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) - hidden_states = self.eh_proj(hidden_states) + try: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) - residual = None - if self.num_layers > 1: - for layer in self.midlayer: - if residual is not None: - hidden_states = hidden_states + residual - hidden_states, residual = layer( + if inputs_embeds is None: + assert self.embed_tokens is not None + inputs_embeds = self.embed_tokens(input_ids).to(self.dtype) + + assert hidden_states is not None + # NOTE: If hidden states from the target model have to be concatenated, + # ideally, we expect that to happen outside the model definition. This + # helps us avoid data-dependent control flow and gives us better CUDA + # graph coverage. + if self._eh_proj_before_attn: + input_embeds = self.enorm(inputs_embeds) + hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) + hidden_states = self.eh_proj(hidden_states) + + residual = None + if self.num_layers > 1: + for layer in self.midlayer: + if residual is not None: + hidden_states = hidden_states + residual + hidden_states, residual = layer( + position_ids=position_ids, + embeds=inputs_embeds, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + ) + else: + hidden_states, residual = self.midlayer( position_ids=position_ids, embeds=inputs_embeds, hidden_states=hidden_states, attn_metadata=attn_metadata, spec_metadata=spec_metadata, ) - else: - hidden_states, residual = self.midlayer( - position_ids=position_ids, - embeds=inputs_embeds, - hidden_states=hidden_states, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata, - ) - hidden_states, hidden_states_to_save = self.norm( - hidden_states, residual) - if self._return_hidden_post_norm: - return hidden_states, hidden_states - return hidden_states, hidden_states_to_save + hidden_states, hidden_states_to_save = self.norm( + hidden_states, residual) + if self._return_hidden_post_norm: + return hidden_states, hidden_states + return hidden_states, hidden_states_to_save + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous_all_rank_num_tokens # We use Llama3 as the base architecture for EAGLE3 draft layers @@ -632,14 +644,13 @@ def forward( spec_metadata: SpecMetadata | None = None, hidden_states: torch.Tensor | None = None, ) -> torch.Tensor: - assert self.embed_tokens is not None - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) if inputs_embeds is None: + assert self.embed_tokens is not None inputs_embeds = self.embed_tokens(input_ids).to(self.dtype) assert hidden_states is not None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ae88b6d62e80..f5c54c1af25f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1867,6 +1867,19 @@ def _get_all_rank_ctx_requests(self, num_ctx_requests: int): return list(self.dist.tp_allgather(num_ctx_requests)) return None + def _set_spec_metadata_all_rank_num_tokens( + self, spec_metadata: SpecMetadata, + spec_all_rank_num_tokens: List[int], + all_rank_num_seqs: List[int]) -> None: + # Eagle3 / MTP-eagle one-model use subseq_all_rank_num_tokens for + # draft loop iterations i>0 (per-sequence counts, since each + # sequence contributes one token per iteration). + spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens + spec_metadata.all_rank_num_seqs = all_rank_num_seqs + if (spec_metadata.spec_dec_mode.is_mtp_eagle_one_model() + or spec_metadata.spec_dec_mode.is_eagle3_one_model()): + spec_metadata.subseq_all_rank_num_tokens = all_rank_num_seqs + def _get_padding_params( self, total_num_tokens: int, num_ctx_requests: int, attn_all_rank_num_tokens: Optional[List[int]] @@ -2075,12 +2088,9 @@ def _prepare_incremental_update_metadata( all_rank_num_tokens = self.dist.tp_cp_allgather( [spec_metadata.num_tokens, len(sequence_lengths)]) - spec_metadata.all_rank_num_tokens = [ - item[0] for item in all_rank_num_tokens - ] - spec_metadata.all_rank_num_seqs = [ - item[1] for item in all_rank_num_tokens - ] + self._set_spec_metadata_all_rank_num_tokens( + spec_metadata, [item[0] for item in all_rank_num_tokens], + [item[1] for item in all_rank_num_tokens]) # Set iteration states - batch dictionary updates self.iter_states.update({ @@ -3302,13 +3312,9 @@ def previous_seq_slots_device(): all_rank_num_tokens = self.dist.tp_cp_allgather( [spec_metadata.num_tokens, len(sequence_lengths)]) - - spec_all_rank_num_tokens = [ - item[0] for item in all_rank_num_tokens - ] - all_rank_num_seqs = [item[1] for item in all_rank_num_tokens] - spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens - spec_metadata.all_rank_num_seqs = all_rank_num_seqs + self._set_spec_metadata_all_rank_num_tokens( + spec_metadata, [item[0] for item in all_rank_num_tokens], + [item[1] for item in all_rank_num_tokens]) if mm_token_indices is not None: mask = torch.ones(total_num_tokens, dtype=torch.bool) @@ -3470,16 +3476,12 @@ def _prepare_tp_inputs_no_cache( attn_metadata.num_tokens, spec_metadata.num_tokens, len(sequence_lengths) ]) - attn_all_rank_num_tokens = [ + attn_metadata.all_rank_num_tokens = [ item[0] for item in all_rank_num_tokens ] - spec_all_rank_num_tokens = [ - item[1] for item in all_rank_num_tokens - ] - all_rank_num_seqs = [item[2] for item in all_rank_num_tokens] - attn_metadata.all_rank_num_tokens = attn_all_rank_num_tokens - spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens - spec_metadata.all_rank_num_seqs = all_rank_num_seqs + self._set_spec_metadata_all_rank_num_tokens( + spec_metadata, [item[1] for item in all_rank_num_tokens], + [item[2] for item in all_rank_num_tokens]) else: all_rank_num_tokens = self.dist.tp_cp_allgather( attn_metadata.num_tokens) diff --git a/tensorrt_llm/_torch/speculative/__init__.py b/tensorrt_llm/_torch/speculative/__init__.py index 0f16df6baffd..d1c1b2605283 100644 --- a/tensorrt_llm/_torch/speculative/__init__.py +++ b/tensorrt_llm/_torch/speculative/__init__.py @@ -2,12 +2,12 @@ from .dflash import DFlashSpecMetadata, DFlashWorker from .draft_target import (DraftTargetOneModelSpecMetadata, DraftTargetOneModelWorker) -from .eagle3 import Eagle3SpecMetadata +from .eagle3 import Eagle3SpecMetadata, MTPEagleWorker from .interface import (SpecMetadata, SpecWorkerBase, prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, should_use_separate_draft_kv_cache) -from .mtp import MTPEagleWorker, MTPSampler, MTPSpecMetadata, MTPWorker +from .mtp import MTPSampler, MTPSpecMetadata, MTPWorker from .ngram import NGramDrafter, NGramPoolManager from .pard import PARDSpecMetadata, PARDWorker from .sa_enhancer import SADraftEnhancer diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index b5b0b9e24877..6acef9ed348f 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set import torch +import torch.nn.functional as F from torch import nn from tensorrt_llm._torch.custom_ops import inplace_slice_copy @@ -9,12 +10,15 @@ from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata +from ..distributed.ops import allgather +from ..model_config import ModelConfig from ..pyexecutor.llm_request import LlmRequest +from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.sampler import TorchSampler from ..pyexecutor.scheduler import ScheduledRequests from .interface import SpecMetadata, SpecWorkerBase -from .mtp import MTPSampler +from .mtp import MTPSampler, _select_mtp_position_ids from .sa_enhancer import SADraftEnhancer from .spec_tree_manager import SpecTreeManager @@ -72,6 +76,15 @@ def __init__(self, ) # sequence length, only used for metadata preparation self.seq_lens = {i: 0 for i in range(slot_size)} + + # Per-request delta pool tracking whether the request is in the + # thinking phase; mirrors MTPHiddenStatesManager.mtp_relaxed_delta_pool. + self.use_relaxed_acceptance_for_thinking = getattr( + config, 'use_relaxed_acceptance_for_thinking', False) + if self.use_relaxed_acceptance_for_thinking: + self.relaxed_delta_pool = torch.zeros((slot_size, ), + dtype=torch.float, + device='cuda') # start indices of each slot self.start_indices = {i: 0 for i in range(slot_size)} # whether the next draft forward is the first @@ -98,6 +111,8 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): if req.is_first_context_chunk: slot_id = self.slot_manager.add_slot(req.request_id) self.slot_ids.append(slot_id) + if self.use_relaxed_acceptance_for_thinking: + self.relaxed_delta_pool[slot_id].fill_(0) # reset the flag before model forward self.is_first_draft = True @@ -108,6 +123,8 @@ def free_resources(self, request: LlmRequest): slot_id = self.slot_manager.get_slot(request.request_id) self.seq_lens[slot_id] = 0 self.start_indices[slot_id] = 0 + if self.use_relaxed_acceptance_for_thinking: + self.relaxed_delta_pool[slot_id].fill_(0) self.slot_manager.remove_slot(request.request_id) if self.sa_manager is not None: self.sa_manager.remove_request(request.request_id) @@ -365,15 +382,32 @@ class Eagle3OneModelSpecMetadata(SpecMetadata): dtype: torch.dtype = torch.bfloat16 # The index of the batch inputs batch_indices_cuda: Optional[torch.Tensor] = None - # Optional resource manager (used to access SA manager for EAGLE3+SA) + # Optional resource manager (used to access SA manager and relaxed-acceptance + # delta pool for Eagle3+SA / Eagle3+relaxed-thinking / MTP Eagle modes) spec_resource_manager: Optional[Eagle3ResourceManager] = None # Dynamic tree flags use_dynamic_tree: bool = False eagle_choices: Optional[List[List[int]]] = None + # Slot IDs for each request; populated in prepare() when spec_resource_manager + # is present (required for relaxed acceptance, mirrors MTPSpecMetadata.slot_ids). + slot_ids: Optional[torch.Tensor] = None + # One-model speculative decoding uses the first draft forward token counts + # for the first loop iteration and per-sequence token counts for + # subsequent iterations. + subseq_all_rank_num_tokens: Optional[List[int]] = None def __post_init__(self): if self.layers_to_capture is None: - if self.num_layers == 1: + if self.spec_dec_mode.is_mtp_eagle_one_model(): + # MTP Eagle one-model feeds the target model's hidden_states + # directly to the MTP layer (see Eagle3OneModelWorker + # prepare_1st_drafter_inputs / _run_draft_forward, both gated + # on self.is_mtp_eagle). It never reads spec_metadata.hidden_states, + # so leave layers_to_capture empty: this makes is_layer_capture() + # return False everywhere and avoids the post-MLP/MoE fusion + # disable side effect in modeling_deepseekv3 / glm / etc. + self.layers_to_capture = () + elif self.num_layers == 1: self.layers_to_capture = (self.num_layers - 1, ) else: if self.num_layers <= 5: @@ -385,8 +419,13 @@ def __post_init__(self): else: self.layers_to_capture = sorted(list(self.layers_to_capture)) self.num_capture_layers = len(self.layers_to_capture) - if (self.spec_resource_manager is not None - and self.spec_resource_manager.hidden_states is not None): + if self.num_capture_layers == 0: + # No layers to capture (MTP Eagle one-model). Skip buffer + # allocation entirely; nothing reads self.hidden_states on this + # path. + self.hidden_states = None + elif (self.spec_resource_manager is not None + and self.spec_resource_manager.hidden_states is not None): self.hidden_states = self.spec_resource_manager.hidden_states expected_cols = self.hidden_size * len(self.layers_to_capture) assert self.hidden_states.shape[1] == expected_cols, ( @@ -415,6 +454,13 @@ def __post_init__(self): dtype=torch.int, device='cuda', ) + # Pre-allocate slot_ids; filled in prepare() when spec_resource_manager + # is present. Mirrors MTPSpecMetadata.slot_ids allocation pattern. + self.slot_ids = torch.empty( + [self.max_num_requests], + dtype=torch.long, + device='cuda', + ) # Set tree flags based on config if self.use_dynamic_tree: @@ -441,11 +487,33 @@ def prepare(self): pin_memory=prefer_pinned()) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - if self.is_spec_dec_tree: - self.num_tokens -= ( - self.num_generations) * self.max_total_draft_tokens - else: - self.num_tokens -= (self.num_generations) * self.max_draft_len + # `num_tokens` here only feeds the attention-DP shape hint + # (allgathered in model_engine and overridden into + # `attn_metadata.all_rank_num_tokens` on the step-0 draft forward). + # Each mode uses a different convention: + # - MTP Eagle: keep the 1st-iter shape (matches input_ids). + # - Eagle3: subtract to the subseq shape. + if not self.spec_dec_mode.is_mtp_eagle_one_model(): + if self.is_spec_dec_tree: + self.num_tokens -= ( + self.num_generations) * self.max_total_draft_tokens + else: + self.num_tokens -= (self.num_generations) * self.max_draft_len + + if getattr(self.spec_resource_manager, "slot_manager", + None) is not None: + # Populate slot_ids for all requests in this batch. Used by relaxed + # acceptance (relaxed_delta_pool indexing), mirroring the pattern + # in MTPSpecMetadata.prepare(). + eagle_slot_ids = [ + self.spec_resource_manager.slot_manager.get_slot(rid) + for rid in self.request_ids + ] + eagle_slot_ids_tensor = torch.tensor(eagle_slot_ids, + dtype=torch.int, + pin_memory=prefer_pinned()) + self.slot_ids[:num_seqs].copy_(eagle_slot_ids_tensor, + non_blocking=True) sa_manager = getattr(self.spec_resource_manager, 'sa_manager', None) if sa_manager is not None: @@ -485,41 +553,51 @@ def _get_max_new_tokens(self, args: TorchSampler.Args, class Eagle3OneModelWorker(SpecWorkerBase): - """Eagle3 one-model worker for linear tree speculative decoding. + """Unified one-model worker for Eagle3 and MTP Eagle speculative decoding. - For dynamic tree mode, use Eagle3OneModelDynamicTreeWorker from - eagle3_dynamic_tree.py instead. + The operating mode is determined by ``spec_config.spec_dec_mode``: + - EAGLE3_ONE_MODEL: multi-layer hidden states from Eagle3, apply_eagle3_fc + projection, independent EAGLE draft model network. + - MTP_EAGLE_ONE_MODEL: single last-layer hidden states, MTP layer called + repeatedly, supports TP-aware sampling and Mamba hybrid cache. + + Where the two modes differ, ``self.is_mtp_eagle`` is used to branch. + For dynamic tree Eagle3, use ``Eagle3OneModelDynamicTreeWorker`` from + ``eagle3_dynamic_tree.py``. """ def __init__(self, spec_config: "EagleDecodingConfig", - mapping: Mapping, + mapping: Optional[Mapping] = None, + model_config: Optional[ModelConfig] = None, use_separate_draft_kv_cache: bool = False): super().__init__(use_separate_draft_kv_cache) self.spec_config = spec_config self.mapping = mapping + # model_config is required for MTP Eagle TP / ADP / Mamba support; the + # Eagle3 path can leave it as None. + self.model_config = model_config + + # Mode flag: True = MTP Eagle one-model, False = Eagle3 one-model. + self.is_mtp_eagle = spec_config.spec_dec_mode.is_mtp_eagle_one_model() + + # SA enhancer (common to both modes) self.sa_enhancer: Optional[SADraftEnhancer] = None 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.spec_tree_manager = None + # MTP Eagle: lazily-resolved flag for Mamba hybrid cache support + self._is_mamba_hybrid_cache = None + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - # Save kv_lens_cuda values separately instead of routing through - # prepare_for_spec_dec, which would clone the tensor and break the - # kv_lens_cuda_runtime view that TRTLLM attention reads from. batch_size = attn_metadata.num_seqs - if hasattr(attn_metadata, 'kv_lens_cuda'): - self._saved_kv_lens_cuda = attn_metadata.kv_lens_cuda[: - batch_size].clone( - ) - else: - self._saved_kv_lens_cuda = None # Save spec-dec params that the drafting loop will overwrite. # Without this, CUDA graph warmup's second iteration would run @@ -547,12 +625,6 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): def _restore_attn_metadata_from_spec_dec(self, attn_metadata): super()._restore_attn_metadata_from_spec_dec(attn_metadata) - if self._saved_kv_lens_cuda is not None: - batch_size = self._saved_kv_lens_cuda.shape[0] - attn_metadata.kv_lens_cuda[:batch_size].copy_( - self._saved_kv_lens_cuda) - self._saved_kv_lens_cuda = None - if self._saved_packed_mask is not None: batch_size = self._saved_packed_mask.shape[0] attn_metadata.spec_decoding_packed_mask[:batch_size].copy_( @@ -599,9 +671,24 @@ def forward(self, self._execute_guided_decoder_if_present(logits) - # Sample and accept tokens + # Sample and accept tokens. ``input_ids`` is required by the relaxed- + # acceptance path (scans for thinking-phase tokens); ignored otherwise. accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( - logits, attn_metadata, spec_metadata) + input_ids, logits, attn_metadata, spec_metadata) + + # Mamba hybrid models need state updates after token acceptance because + # the accepted token count affects which Mamba states are valid. The + # isinstance check below naturally no-ops on non-Mamba kv_cache_managers, + # so this is safe to run unconditionally regardless of spec mode (Eagle3 + # over a Mamba-style draft is plausible, even if no such draft exists today). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices) sa_manager = getattr(spec_metadata.spec_resource_manager, 'sa_manager', None) @@ -630,7 +717,9 @@ def forward(self, spec_metadata=spec_metadata, draft_model=draft_model) - # Predict draft tokens + # 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 @@ -679,28 +768,38 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, num_contexts, batch_size, num_accepted_tokens, original_all_rank_num_tokens): - """Original linear draft loop (1 token per layer).""" + """Linear draft loop, unified for Eagle3 and MTP Eagle.""" runtime_draft_len = spec_metadata.runtime_draft_len + num_gens = batch_size - num_contexts next_draft_tokens = [] draft_logits_list = [] - position_ids = inputs["position_ids"] + last_tokens_idx = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): for i in range(runtime_draft_len): + # Run draft model (mode-specific via helper). The helper + # passes ``all_rank_num_tokens`` as a kwarg so the draft model + # handles save/restore internally (Eagle3DraftModel.forward + # uses try/finally); attn_metadata is left untouched here. + hidden_states, hidden_states_to_save = self._run_draft_forward( + draft_model, inputs, spec_metadata, i) + + # Compute gather_ids: on the first draft step each generation + # request may have accepted multiple tokens, so we index into + # the flattened token sequence to find the last accepted one. + # From step 1 onwards every sequence has length 1, so + # ``batch_indices_cuda`` is sufficient. if i == 0: - num_gens = batch_size - num_contexts start_ids_gen = ( spec_metadata.batch_indices_cuda[:num_gens] * (runtime_draft_len + 1)).long() gather_ids_gen = (start_ids_gen + num_accepted_tokens[num_contexts:] - 1 + attn_metadata.num_ctx_tokens) - gather_ids = torch.concat([ - spec_metadata.gather_ids[:num_contexts], gather_ids_gen - ], - dim=0) + gather_ids = torch.concat( + [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) else: - # All of the seq_len are 1, use batch_indices_cuda as gather_ids gather_ids = spec_metadata.batch_indices_cuda[:batch_size] if self.guided_decoder is not None: @@ -709,59 +808,133 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, num_accepted_tokens, draft_step=i) - # Update attn_metadata.all_rank_num_tokens for attention DP - if original_all_rank_num_tokens is not None: - if i == 0: - attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens - elif spec_metadata.all_rank_num_seqs is not None: - attn_metadata.all_rank_num_tokens = spec_metadata.all_rank_num_seqs - - hidden_states, hidden_states_to_save = draft_model.model( - **inputs) - - # FIXME (jhaotingc): Currently we disable use_spec_decoding mode for Eagle engine nth steps except 1st step. - # Eagle engine takes in draft_len tokens from the previous step, run spec-dec mode with those tokens, - # then the following step can use regular decoding mode to generate 1 tokens per step. - # Currently the spec-dec mask for chained tree is not implemented yet. - # When token tree is supported, this can be removed and all steps may use spec-dec mode as well. - attn_metadata.use_spec_decoding = False - - logits = draft_model.logits_processor(hidden_states[gather_ids], - draft_model.lm_head, - attn_metadata, True) + # Compute logits. + # MTP Eagle: shared_head of the MTP layer, with optional + # ADP+LM-head-TP padding to ``max_num_requests`` so every TP + # rank produces logits of the same shape. + # Eagle3: logits_processor of the EAGLE draft model. + use_lm_head_tp_in_adp = ( + self.is_mtp_eagle and self.model_config is not None + and self.model_config.mapping.enable_attention_dp + and getattr(self.model_config.mapping, + 'enable_lm_head_tp_in_adp', False)) + if self.is_mtp_eagle: + if use_lm_head_tp_in_adp: + hidden_states_gathered = hidden_states[gather_ids] + token_count = hidden_states_gathered.view( + -1, hidden_states_gathered.shape[-1]).shape[0] + max_num_requests = spec_metadata.max_num_requests + pad_len = max_num_requests - token_count + if pad_len > 0: + padded_hidden_states = F.pad( + hidden_states_gathered.view( + -1, hidden_states_gathered.shape[-1]), + (0, 0, 0, pad_len), + mode="constant", + value=0) + elif pad_len == 0: + padded_hidden_states = hidden_states_gathered.view( + -1, hidden_states_gathered.shape[-1]) + else: + raise ValueError( + "Eagle3OneModelWorker (MTP Eagle mode): " + "token_count > max_num_requests, which is not supported" + ) + logits = draft_model.mtp_layers[0].shared_head( + padded_hidden_states, draft_model.lm_head, + attn_metadata, True) + else: + logits = draft_model.mtp_layers[0].shared_head( + hidden_states[gather_ids], draft_model.lm_head, + attn_metadata, True) + else: + logits = draft_model.logits_processor( + hidden_states[gather_ids], draft_model.lm_head, + attn_metadata, True) + if self.guided_decoder is not None: + if self.is_mtp_eagle: + self.guided_decoder.execute_draft_batch(logits, + draft_step=i) + else: + d2t = getattr(draft_model.model, "d2t", None) + self.guided_decoder.execute_draft_batch(logits, + d2t, + draft_step=i) + + # Sample the next draft token. + # MTP Eagle: TP-aware sampler; when ADP+LM-head-TP is active + # logits are padded to max_num_requests across TP ranks, so + # the result must be trimmed back to token_count. + # Eagle3: simple greedy sampling; d2t remaps vocab indices when + # the draft model uses a compressed vocabulary. + if self.is_mtp_eagle: + if use_lm_head_tp_in_adp: + mapping_lm_head_tp = draft_model.mtp_layers[ + 0].shared_head.mapping_lm_head_tp + new_draft_token = self.draft_sampler( + logits, mapping_lm_head_tp) + new_draft_token = new_draft_token[:token_count] + else: + new_draft_token = self.draft_sampler(logits) + else: d2t = getattr(draft_model.model, "d2t", None) - self.guided_decoder.execute_draft_batch(logits, - d2t, - draft_step=i) + new_draft_token = self._draft_sampler_greedy(logits, d2t) - if spec_metadata.use_rejection_sampling: + # Stash unpadded Eagle3 draft logits for rejection sampling on + # the next iteration. MTP Eagle's logits may be ADP-padded to + # max_num_requests, so we skip them here. + if not self.is_mtp_eagle and spec_metadata.use_rejection_sampling: draft_logits_list.append(logits.clone()) - new_draft_token = self.draft_decoder(logits, draft_model) next_draft_tokens.append(new_draft_token) - # update inputs - hidden_states = hidden_states_to_save[gather_ids] - position_ids = inputs["position_ids"][gather_ids] + 1 - # update attn_metadata + + # Update hidden states for the next iteration. + # MTP Eagle: the MTP layer returns a single tensor; slice by + # gather_ids to get one hidden state per request. + # Eagle3: the EAGLE draft model returns a secondary + # ``hidden_states_to_save`` specifically for this purpose. + if self.is_mtp_eagle: + hidden_states = hidden_states[gather_ids] + else: + hidden_states = hidden_states_to_save[gather_ids] + position_ids = (_select_mtp_position_ids( + inputs["position_ids"], gather_ids) + 1) + + # Update attn_metadata for the next iteration. if i == 0: attn_metadata._seq_lens[:batch_size].fill_(1) attn_metadata._seq_lens_cuda[:batch_size].fill_(1) attn_metadata.on_update() - # cannot run generation if there is no kv cache - if inputs["attn_metadata"].kv_cache_manager is not None: + has_kv_cache = inputs[ + "attn_metadata"].kv_cache_manager is not None + if has_kv_cache: attn_metadata.host_request_types[:attn_metadata. num_contexts].fill_(1) attn_metadata.num_contexts = 0 - # update kv_lens_cuda if hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( runtime_draft_len - num_accepted_tokens[num_contexts:]) attn_metadata.kv_lens_cuda[:num_contexts] += 1 - elif hasattr(attn_metadata, 'kv_lens_cuda'): - attn_metadata.kv_lens_cuda[:batch_size] += 1 - # support attention dp + + if has_kv_cache: + self._prepare_flash_mla_generation_layout( + attn_metadata, num_contexts, batch_size) + if hasattr(attn_metadata, 'kv_lens_cuda'): + attn_metadata.update_for_spec_dec() + + # Both Eagle3 and MTP Eagle drafters take ``draft_len + 1`` + # tokens in the first draft step (attention runs in spec-dec + # mode), then 1 token per step in subsequent iterations. + # Disable spec_decoding here so the masks/positions stay + # correct on subsequent iters. + attn_metadata.use_spec_decoding = False + else: + if hasattr(attn_metadata, 'kv_lens_cuda'): + attn_metadata.kv_lens_cuda[:batch_size] += 1 + attn_metadata.update_for_spec_dec() + inputs = { "input_ids": new_draft_token, "position_ids": position_ids, @@ -787,53 +960,220 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, return next_draft_tokens + def _get_step_all_rank_num_tokens(self, spec_metadata, step_idx: int): + """Pick the right ``all_rank_num_tokens`` for this draft iteration. + + Step 0 uses ``spec_metadata.all_rank_num_tokens``; subsequent steps + use ``spec_metadata.subseq_all_rank_num_tokens`` since every sequence + contributes a single token per iteration. + """ + return (spec_metadata.all_rank_num_tokens + if step_idx == 0 else spec_metadata.subseq_all_rank_num_tokens) + + def _run_draft_forward(self, draft_model, inputs, spec_metadata, + step_idx: int): + """Invoke the draft model for one iteration, branching on mode. + + ``all_rank_num_tokens`` is passed as a kwarg in both modes. For MTP + Eagle it goes directly to ``mtp_layers[0]``; for Eagle3 it goes to + ``Eagle3DraftModel.forward`` which guards it with a try/finally so + attn_metadata sees the original value on return. + """ + all_rank_num_tokens = self._get_step_all_rank_num_tokens( + spec_metadata, step_idx) + + if self.is_mtp_eagle: + hidden_states = draft_model.mtp_layers[0]( + embed_tokens=draft_model.embed_tokens, + all_rank_num_tokens=all_rank_num_tokens, + **inputs) + return hidden_states, None + + inputs["all_rank_num_tokens"] = all_rank_num_tokens + hidden_states, hidden_states_to_save = draft_model.model(**inputs) + return hidden_states, hidden_states_to_save + + def _prepare_flash_mla_generation_layout(self, attn_metadata, num_contexts, + batch_size): + """Reorder ``kv_block_ids_per_seq`` so gen requests precede context. + + Flash MLA on first-step expects the layout used during normal + generation; both Eagle3 and MTP Eagle hit this when context requests + share the batch with gen requests. + """ + if num_contexts <= 0 or not attn_metadata.enable_flash_mla: + return + reorder_block_ids_per_seq = torch.cat([ + attn_metadata.kv_block_ids_per_seq[num_contexts:batch_size], + attn_metadata.kv_block_ids_per_seq[:num_contexts] + ]) + attn_metadata.block_ids_per_seq[:batch_size, :].copy_( + reorder_block_ids_per_seq, non_blocking=True) + + @torch.compile(options={"max-autotune": True}) + def _get_local_max_and_combined(self, logits, mapping_lm_tp=None): + local_max_values, local_argmax = torch.max(logits, dim=-1, keepdim=True) + vocab_per_rank = logits.shape[-1] + mapping_lm_tp = mapping_lm_tp if mapping_lm_tp is not None else \ + self.model_config.mapping + max_index_per_rank = local_argmax.type( + torch.int32) + (mapping_lm_tp.tp_rank * vocab_per_rank) + max_index_per_rank_float = max_index_per_rank.float() + local_max_values_float32 = local_max_values.float() + combined = torch.stack( + [max_index_per_rank_float, local_max_values_float32], + dim=-1).flatten(-2) + return combined + + @torch.compile(options={"max-autotune": True}) + def _get_draft_tokens_from_gathered(self, gathered): + gathered_indices_float = gathered[..., 0::2] + gathered_values_float = gathered[..., 1::2] + max_indices = torch.argmax(gathered_values_float, dim=-1, keepdim=True) + draft_tokens = torch.gather(gathered_indices_float, -1, + max_indices).squeeze(-1).type(torch.int32) + return draft_tokens + + def draft_sampler( + self, + logits: torch.Tensor, + mapping_lm_head_tp=None, + ): + """TP-aware greedy draft token sampler (MTP Eagle path). + + Falls back to simple argmax when no tensor parallelism is active or + when only attention DP is enabled without LM-head TP. + """ + if (self.model_config is not None + and hasattr(self.model_config, 'mapping') + and self.model_config.mapping.tp_size > 1 + and not self.model_config.mapping.enable_attention_dp): + combined = self._get_local_max_and_combined(logits) + gathered = allgather(combined, self.model_config.mapping, dim=-1) + return self._get_draft_tokens_from_gathered(gathered) + elif (self.model_config is not None + and hasattr(self.model_config, 'mapping') + and self.model_config.mapping.tp_size > 1 + and self.model_config.mapping.enable_lm_head_tp_in_adp): + combined = self._get_local_max_and_combined(logits, + mapping_lm_head_tp) + gathered = allgather(combined, mapping_lm_head_tp, dim=-1) + batch_size = logits.shape[0] + local_batch_size = batch_size // mapping_lm_head_tp.tp_size + gathered = gathered.view(mapping_lm_head_tp.tp_size, + local_batch_size, -1) + sliced_gathered = gathered[mapping_lm_head_tp.tp_rank] + return self._get_draft_tokens_from_gathered(sliced_gathered) + else: + return self._draft_sampler_greedy(logits) + + @torch.compile(options={"max-autotune": True}) + def _topk_kernel(self, gen_logprobs, num_gens, mtp_num_modules, + spec_metadata): + topk_value, topk_indices = torch.topk(gen_logprobs, + k=self.spec_config.relaxed_topk, + dim=-1) + topk_indices = topk_indices.reshape(num_gens, mtp_num_modules + 1, + self.spec_config.relaxed_topk) + topk_value = topk_value.reshape(num_gens, mtp_num_modules + 1, + self.spec_config.relaxed_topk) + draft_tokens = spec_metadata.draft_tokens.reshape( + num_gens, mtp_num_modules) + return topk_value, topk_indices, draft_tokens + + @torch.compile(options={"max-autotune": True}) + def _process_generation_logits(self, logits, num_contexts): + gen_logits = logits[num_contexts:] + gen_logprobs = torch.softmax(gen_logits, dim=-1) + return gen_logprobs + def sample_and_accept_draft_tokens( self, + input_ids: torch.IntTensor, logits: torch.Tensor, attn_metadata: AttentionMetadata, spec_metadata: Eagle3OneModelSpecMetadata, ): + """Sample the golden token and verify previously proposed draft tokens. + + ``input_ids`` is scanned for thinking-phase tokens when relaxed + acceptance is enabled (both Eagle3 and MTP Eagle); ignored otherwise. + """ batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts - # Linear mode: reshape draft tokens for base implementation + runtime_draft_len = spec_metadata.runtime_draft_len + + if getattr(self.spec_config, 'use_relaxed_acceptance_for_thinking', + False): + # Relaxed acceptance — common path for Eagle3 and MTP Eagle. + # Accepts draft tokens that fall within the top-K candidates of the + # target distribution during the thinking phase. + if logits.dim() == 1: + logits = logits.unsqueeze(0) + + accepted_tokens = torch.ones((batch_size, runtime_draft_len + 1), + dtype=torch.int, + device=logits.device) + num_accepted_tokens = torch.ones(batch_size, + dtype=torch.int, + device=logits.device) + + resource_manager = spec_metadata.spec_resource_manager + relaxed_delta_pool = resource_manager.relaxed_delta_pool + + # Context phase: detect thinking tokens and update the delta pool + con_logits = logits[:num_contexts] + con_target_tokens = torch.argmax(con_logits, dim=-1) + accepted_tokens[:num_contexts, 0] = con_target_tokens[:num_contexts] + last_tokens_idx_for_thinking = torch.cumsum( + attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1 + ctx_input_ids = input_ids[:attn_metadata.num_ctx_tokens] + ctx_is_think = (ctx_input_ids == + self.spec_config.begin_thinking_phase_token).int() + ctx_is_think_cumsum = torch.cumsum(ctx_is_think, dim=0) + ctx_last_cumsum = ctx_is_think_cumsum[ + last_tokens_idx_for_thinking[:num_contexts]] + ctx_think_tokens_num = torch.diff( + ctx_last_cumsum, + dim=0, + prepend=torch.zeros(1, + dtype=torch.int, + device=ctx_last_cumsum.device)) + ctx_delta = (ctx_think_tokens_num + >= 1).int() * self.spec_config.relaxed_delta + ctx_slot_ids = spec_metadata.slot_ids[:num_contexts] + relaxed_delta_pool.index_copy_(0, ctx_slot_ids, ctx_delta) + + # Generation phase: top-k logprobs + relaxed acceptance op + gen_logprobs = self._process_generation_logits(logits, num_contexts) + topk_value, topk_indices, draft_tokens = self._topk_kernel( + gen_logprobs, num_gens, runtime_draft_len, spec_metadata) + + accepted_tokens, num_accepted_tokens = torch.ops.trtllm.mtp_relaxed_acceptance_op( + spec_metadata.slot_ids, topk_value, topk_indices, draft_tokens, + relaxed_delta_pool, num_accepted_tokens, accepted_tokens, + runtime_draft_len, batch_size, num_contexts, + self.spec_config.relaxed_topk, self.spec_config.relaxed_delta, + self.spec_config.begin_thinking_phase_token, + self.spec_config.end_thinking_phase_token) + + num_accepted_tokens = self._apply_force_accepted_tokens( + num_accepted_tokens, num_contexts, runtime_draft_len) + + return accepted_tokens, num_accepted_tokens + + # Strict acceptance — common path for Eagle3 and MTP Eagle. Both modes + # use runtime_draft_len for dynamic draft length support. + if logits.dim() == 1: + logits = logits.unsqueeze(0) draft_tokens = spec_metadata.draft_tokens.reshape( - num_gens, - spec_metadata.runtime_draft_len) if num_gens > 0 else torch.empty( - 0, - spec_metadata.runtime_draft_len, - dtype=torch.int, - device=logits.device) + num_gens, runtime_draft_len) if num_gens > 0 else torch.empty( + 0, runtime_draft_len, dtype=torch.int, device=logits.device) return self._accept_draft_tokens(logits, draft_tokens, num_contexts, batch_size, spec_metadata) - def draft_decoder( - self, - logits: torch.Tensor, - draft_model: nn.Module, - ): - ''' - Sampling draft tokens with support for non-greedy sampling. - - Args: - logits: torch.Tensor - [num_tokens, vocab_size] - Logits produced by the draft model. - draft_model: nn.Module - The draft model. - - Returns: - draft_tokens: torch.Tensor - [batch_size * max_draft_len] - Draft token ids. Flattened. - ''' - - d2t = getattr(draft_model.model, "d2t", None) - draft_tokens = self._draft_sampler_greedy(logits, d2t) - - return draft_tokens - def prepare_1st_drafter_inputs( self, input_ids: torch.LongTensor, @@ -844,15 +1184,24 @@ def prepare_1st_drafter_inputs( spec_metadata: Eagle3OneModelSpecMetadata, draft_model: nn.Module, ): + """Prepare inputs for the first draft model forward. + + Branching: + - Eagle3: applies ``apply_eagle3_fc`` on multi-layer concatenated + hidden states. + - MTP Eagle: uses ``hidden_states`` directly (single last layer); + no FC projection. + """ num_contexts = attn_metadata.num_contexts num_tokens = input_ids.shape[0] - # prepare hidden states - hidden_size_up = spec_metadata.hidden_size * len( - spec_metadata.layers_to_capture) - hidden_states = spec_metadata.hidden_states[:num_tokens, : - hidden_size_up] - hidden_states = draft_model.apply_eagle3_fc(hidden_states) + if not self.is_mtp_eagle: + # Eagle3: project the multi-layer concatenated hidden states. + hidden_size_up = spec_metadata.hidden_size * len( + spec_metadata.layers_to_capture) + hidden_states = spec_metadata.hidden_states[:num_tokens, : + hidden_size_up] + hidden_states = draft_model.apply_eagle3_fc(hidden_states) # context input_ids_ctx = self._prepare_context_input_ids( @@ -873,3 +1222,25 @@ def prepare_1st_drafter_inputs( "attn_metadata": attn_metadata, "spec_metadata": spec_metadata, } + + +class MTPEagleWorker(Eagle3OneModelWorker): + """Backward-compatible alias for ``Eagle3OneModelWorker`` in MTP Eagle mode. + + The constructor matches the historical positional signature + ``(spec_config, model_config, use_separate_draft_kv_cache)`` so callers + that import ``MTPEagleWorker`` from ``mtp.py`` or instantiate it directly + keep working. All logic is inherited from :class:`Eagle3OneModelWorker`. + """ + + def __init__(self, + spec_config, + model_config: Optional[ModelConfig] = None, + use_separate_draft_kv_cache: bool = False): + super().__init__( + spec_config, + mapping=None, + model_config=model_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) + # Preserved for callers/tests that still expect this attribute. + self.is_thop = False diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 90afb921d0a8..47376001166d 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -180,7 +180,19 @@ def __init__( self, spec_config: "EagleDecodingConfig", mapping, use_separate_draft_kv_cache: bool = False ): """Initialize dynamic-tree specific buffers and helper ops.""" - super().__init__(spec_config, mapping, use_separate_draft_kv_cache) + super().__init__( + spec_config, + mapping=mapping, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + if ( + getattr(spec_config, "use_relaxed_acceptance_for_thinking", False) + or getattr(spec_config, "sa_config", None) is not None + ): + raise ValueError( + "Dynamic tree mode does not support relaxed acceptance or " + "suffix-automaton enhancement." + ) assert self.use_dynamic_tree, ( "Eagle3OneModelDynamicTreeWorker requires use_dynamic_tree=True" ) @@ -453,8 +465,13 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): ) @nvtx_range("eagle3_dyn.sample_and_accept_draft_tokens") - def sample_and_accept_draft_tokens(self, logits, attn_metadata, spec_metadata): - """Override to handle dynamic tree verification.""" + def sample_and_accept_draft_tokens(self, input_ids, logits, attn_metadata, spec_metadata): + """Override to handle dynamic tree verification. + + ``input_ids`` is unused here (relaxed acceptance is not supported in + dynamic-tree mode); accepted to match the base class signature. + """ + del input_ids batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index c7a8358124cf..c62111f0f511 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + import copy import os from abc import ABC, abstractmethod @@ -220,7 +235,10 @@ class SpeculativeDecodingMode(IntEnum): AUTO = auto() def is_mtp_one_model(self): - return self == SpeculativeDecodingMode.MTP or self == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL + # Union: covers vanilla MTP and MTP_EAGLE_ONE_MODEL. Use is_mtp_vanilla() + # when only the vanilla MTP variant should match. + return (self == SpeculativeDecodingMode.MTP + or self == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) def is_mtp_eagle_one_model(self): return self == SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL @@ -296,7 +314,7 @@ def support_capturable_guided_decoder(self): def support_dynamic_draft_len(self): # TODO: expand to all one-model algorithms - return self.is_eagle3_one_model() + return self.is_eagle3_one_model() or self.is_mtp_eagle_one_model() def has_draft_model(self): return self.is_eagle3() or self.is_draft_target() or self.is_mtp_eagle() @@ -716,8 +734,15 @@ def skip_forward( attn_metadata, spec_metadata, draft_model, + resource_manager=None, ): - """Skip spec dec for non-last rank (PP). Returns placeholder outputs.""" + """Skip spec dec for non-last rank (PP). Returns placeholder outputs. + + ``resource_manager`` is accepted but unused; it appears in the + ``forward()`` signature of one-model workers (Eagle3 / MTP-Eagle) and + the caller in ``modeling_speculative.py`` forwards it unconditionally, + so the skip path must accept it as well. + """ batch_size = attn_metadata.num_seqs accepted_tokens = torch.empty((batch_size, (self.max_draft_len + 1)), dtype=torch.int, diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 20181bf29179..256def91104a 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -3,16 +3,12 @@ from typing import TYPE_CHECKING, List, Optional import torch -import torch.nn.functional as F -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \ - MambaHybridCacheManager from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.mapping import Mapping from ..attention_backend import AttentionMetadata from ..distributed.ops import allgather -from ..model_config import ModelConfig from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager from ..pyexecutor.sampler import TorchSampler @@ -211,7 +207,7 @@ def prepare(self): mtp_slot_ids.append(slot_id) # MTP Vanilla: Update mtp hidden states and past tokens - if self.spec_dec_mode.is_mtp_one_model(): + if self.spec_dec_mode.is_mtp_vanilla(): mtp_hidden_states_ptrs = [] mtp_past_tokens_ptrs = [] for slot_id in mtp_slot_ids: @@ -1139,277 +1135,3 @@ def draft_sampler( draft_tokens = self._draft_sampler_greedy(logits) return draft_tokens - - -class MTPEagleWorker(MTPWorker): - - def __init__(self, - spec_config: "MTPDecodingConfig", - model_config: Optional[ModelConfig] = None, - use_separate_draft_kv_cache: bool = False): - super().__init__(spec_config, model_config, use_separate_draft_kv_cache) - self.model_config = model_config - self.mtp_num_modules = spec_config.max_draft_len - self._is_mamba_hybrid_cache = None - - @torch.compile(options={"max-autotune": True}) - def update_draft_tokens(self, next_draft_tokens, new_draft_token, - hidden_states, gather_ids, inputs): - next_draft_tokens.append(new_draft_token) - # update inputs - hidden_states = hidden_states[gather_ids] - position_ids = ( - _select_mtp_position_ids(inputs["position_ids"], gather_ids) + 1) - return hidden_states, position_ids - - @torch.compile(options={"max-autotune": True}) - def prepare_position_ids_and_last_tokens(self, position_ids, seq_lens_cuda): - position_ids = position_ids.squeeze(0) - last_tokens_idx = torch.cumsum(seq_lens_cuda, dim=0, - dtype=torch.long) - 1 - return position_ids, last_tokens_idx - - def forward( - self, - input_ids, - position_ids, - hidden_states, - logits, - attn_metadata, - spec_metadata, - draft_model, - resource_manager=None, - ): - - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - - raw_logits = logits - - self._execute_guided_decoder_if_present(logits) - - # Sample and verify draft tokens - accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( - input_ids, logits, spec_metadata, attn_metadata) - - if self._is_mamba_hybrid_cache is None: - self._is_mamba_hybrid_cache = isinstance( - attn_metadata.kv_cache_manager, MambaHybridCacheManager) - if num_gens > 0 and self._is_mamba_hybrid_cache: - attn_metadata.kv_cache_manager.update_mamba_states( - attn_metadata=attn_metadata, - num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices) - - # Save the old attn_metadata and spec_metadata - self._prepare_attn_metadata_for_spec_dec(attn_metadata) - - position_ids, last_tokens_idx = self.prepare_position_ids_and_last_tokens( - position_ids, attn_metadata.seq_lens_cuda) - inputs = self.prepare_drafter_inputs(input_ids=input_ids, - position_ids=position_ids, - last_tokens_idx=last_tokens_idx, - hidden_states=hidden_states, - accepted_tokens=accepted_tokens, - attn_metadata=attn_metadata, - spec_metadata=spec_metadata) - - # Get the draft KV cache manager if using separate layouts - draft_kv_cache_manager = self.get_draft_kv_cache_manager( - resource_manager) - - # Predict draft tokens - next_draft_tokens = [] - with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager): - for i in range(self.mtp_num_modules): - if i == 0: - hidden_states = draft_model.mtp_layers[0]( - embed_tokens=draft_model.embed_tokens, - all_rank_num_tokens=spec_metadata.all_rank_num_tokens, - **inputs) - - start_ids_gen = ( - spec_metadata.batch_indices_cuda[:num_gens] * - (self.mtp_num_modules + 1)).long() - gather_ids_gen = (start_ids_gen + - num_accepted_tokens[num_contexts:] - 1 + - attn_metadata.num_ctx_tokens) - gather_ids = torch.concat( - [last_tokens_idx[:num_contexts], gather_ids_gen], dim=0) - else: - hidden_states = draft_model.mtp_layers[0]( - embed_tokens=draft_model.embed_tokens, - all_rank_num_tokens=spec_metadata. - subseq_all_rank_num_tokens, - **inputs) - - # All of the seq_len are 1, use batch_indices_cuda as gather_ids - gather_ids = spec_metadata.batch_indices_cuda[:batch_size] - - if self.guided_decoder is not None: - new_tokens = inputs["input_ids"][gather_ids] - self.guided_decoder.add_draft_batch(new_tokens, - num_accepted_tokens, - draft_step=i) - if self.model_config.mapping.enable_attention_dp and \ - getattr(self.model_config.mapping, 'enable_lm_head_tp_in_adp', False): - hidden_states_gathered = hidden_states[gather_ids] - token_count = hidden_states_gathered.view( - -1, hidden_states_gathered.shape[-1]).shape[0] - max_num_requests = spec_metadata.max_num_requests - pad_len = max_num_requests - token_count - if pad_len > 0: - padded_hidden_states = F.pad( - hidden_states_gathered.view( - -1, hidden_states_gathered.shape[-1]), - (0, 0, 0, pad_len), - mode="constant", - value=0) - elif pad_len == 0: - padded_hidden_states = hidden_states_gathered.view( - -1, hidden_states_gathered.shape[-1]) - else: - raise ValueError( - "In MTPEagleWorker.forward(), token_count > max_num_requests, which is not supported" - ) - logits = draft_model.mtp_layers[0].shared_head( - padded_hidden_states, draft_model.lm_head, - attn_metadata, True) - else: - logits = draft_model.mtp_layers[0].shared_head( - hidden_states[gather_ids], draft_model.lm_head, - attn_metadata, True) - if self.guided_decoder is not None: - self.guided_decoder.execute_draft_batch(logits, - draft_step=i) - - if self.model_config.mapping.enable_attention_dp and \ - getattr(self.model_config.mapping, 'enable_lm_head_tp_in_adp', False): - mapping_lm_head_tp = draft_model.mtp_layers[ - 0].shared_head.mapping_lm_head_tp - new_draft_token = self.draft_sampler( - logits, mapping_lm_head_tp) - new_draft_token = new_draft_token[:token_count] - else: - new_draft_token = self.draft_sampler(logits) - - hidden_states, position_ids = self.update_draft_tokens( - next_draft_tokens, new_draft_token, hidden_states, - gather_ids, inputs) - # update attn_metadata - if i == 0: - attn_metadata._seq_lens[:batch_size].fill_(1) - attn_metadata._seq_lens_cuda[:batch_size].fill_(1) - attn_metadata.on_update() - # cannot run generation if there is no kv cache - has_kv_cache = inputs[ - "attn_metadata"].kv_cache_manager is not None - if has_kv_cache: - attn_metadata.host_request_types[:attn_metadata. - num_contexts].fill_(1) - attn_metadata.num_contexts = 0 - # update kv_lens_cuda - if hasattr(attn_metadata, 'kv_lens_cuda'): - attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= ( - self.mtp_num_modules - - num_accepted_tokens[num_contexts:]) - attn_metadata.kv_lens_cuda[:num_contexts] += 1 - # update metadata for flash mla - if has_kv_cache and num_contexts > 0 and attn_metadata.enable_flash_mla: - reorder_block_ids_per_seq = torch.cat([ - attn_metadata. - kv_block_ids_per_seq[num_contexts:batch_size], - attn_metadata.kv_block_ids_per_seq[:num_contexts] - ]) - attn_metadata.block_ids_per_seq[:batch_size, :].copy_( - reorder_block_ids_per_seq, non_blocking=True) - # update metadata - # some attention metadata needs to be updated when changing seq_lens/kv_lens - attn_metadata.update_for_spec_dec() - # Disable spec-dec mode for subsequent iterations (i>0) - # as draft model only infer 1 token for the subsequent inference. - attn_metadata.use_spec_decoding = False - elif hasattr(attn_metadata, 'kv_lens_cuda'): - # update kv_lens_cuda - attn_metadata.kv_lens_cuda[:batch_size] += 1 - - # update metadata - # some attention metadata needs to be updated when changing kv_lens - attn_metadata.update_for_spec_dec() - inputs = { - "input_ids": new_draft_token, - "position_ids": position_ids, - "hidden_states": hidden_states, - "attn_metadata": attn_metadata, - } - - # restore attn_metadata to support cuda graph - self._restore_attn_metadata_from_spec_dec(attn_metadata) - attn_metadata.use_spec_decoding = True - - # Override with SA draft tokens after all MTP layers have run, - # so that MTP layers never see SA tokens in their inputs. - # Must happen before stacking since next_draft_tokens is still a list. - if self.sa_enhancer is not None: - stacked = torch.stack(next_draft_tokens, dim=1) - gen_draft_tokens = stacked[num_contexts:] - gen_draft_tokens = self.sa_enhancer.maybe_override_all_draft_tokens( - gen_draft_tokens) - stacked[num_contexts:] = gen_draft_tokens - next_draft_tokens = [stacked[:, i] for i in range(stacked.shape[1])] - - next_draft_tokens, next_new_tokens = self._prepare_next_tokens( - next_draft_tokens, accepted_tokens, spec_metadata, batch_size, - num_accepted_tokens) - - return { - 'logits': raw_logits, - 'new_tokens': accepted_tokens, - 'new_tokens_lens': num_accepted_tokens, - 'next_draft_tokens': next_draft_tokens, - 'next_new_tokens': next_new_tokens - } - - @torch.compile(options={"max-autotune": True}) - def _prepare_next_tokens(self, next_draft_tokens, accepted_tokens, - spec_metadata, batch_size, num_accepted_tokens): - """ - Stack draft tokens and prepare next_new_tokens for overlap scheduler. - """ - next_draft_tokens = torch.stack(next_draft_tokens, dim=1) - next_new_tokens = self._prepare_next_new_tokens( - accepted_tokens, next_draft_tokens, - spec_metadata.batch_indices_cuda, batch_size, num_accepted_tokens) - return next_draft_tokens, next_new_tokens - - @torch.compile(options={"max-autotune": True}) - def prepare_drafter_inputs( - self, - input_ids: torch.IntTensor, - position_ids: torch.IntTensor, - last_tokens_idx: torch.LongTensor, - hidden_states: torch.Tensor, - accepted_tokens: torch.Tensor, - attn_metadata: AttentionMetadata, - spec_metadata: MTPSpecMetadata, - ): - num_contexts = attn_metadata.num_contexts - - # context - input_ids_ctx = self._prepare_context_input_ids( - input_ids, attn_metadata.num_ctx_tokens, last_tokens_idx, - accepted_tokens, num_contexts) - - # generation - input_ids_gen = accepted_tokens[num_contexts:, :].flatten() - - # get draft inputs - input_ids = torch.concat([input_ids_ctx, input_ids_gen], dim=0) - - return { - "input_ids": input_ids, - "position_ids": position_ids, - "hidden_states": hidden_states, - "attn_metadata": attn_metadata, - } diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 68ef347b0533..8bcf20d35b9d 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -20,11 +20,10 @@ from .eagle3 import (Eagle3OneModelDynamicTreeResourceManager, Eagle3OneModelSampler, Eagle3OneModelSpecMetadata, Eagle3OneModelWorker, Eagle3ResourceManager, - Eagle3SpecMetadata) + Eagle3SpecMetadata, MTPEagleWorker) from .eagle3_dynamic_tree import Eagle3OneModelDynamicTreeWorker from .model_drafter import ModelDrafter -from .mtp import (MTPEagleWorker, MTPHiddenStatesManager, MTPSampler, - MTPSpecMetadata, MTPWorker) +from .mtp import MTPHiddenStatesManager, MTPSampler, MTPSpecMetadata, MTPWorker from .ngram import NGramDrafter, NGramPoolManager from .pard import PARDSpecMetadata, PARDWorker from .sa_worker import SASampler, SASpecMetadata, SAWorker @@ -43,7 +42,28 @@ def get_spec_metadata(spec_config, use_rejection_sampling = getattr(spec_config, "use_rejection_sampling", False) vocab_size = getattr(model_config, "vocab_size", 0) - if spec_config.spec_dec_mode.is_mtp_one_model(): + if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): + # MTP Eagle one-model reuses Eagle3 one-model metadata for the + # unified worker/sampler/slot_ids plumbing, but skips per-layer + # hidden-state capture: the worker feeds the target model's + # hidden_states directly into the MTP layer, so we leave + # layers_to_capture unset and let Eagle3OneModelSpecMetadata default + # it to an empty tuple. This also keeps post-MLP/MoE fusion enabled + # on models that gate it on is_layer_capture(). + return Eagle3OneModelSpecMetadata( + max_draft_len=spec_config.max_draft_len, + max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, + spec_dec_mode=spec_config.spec_dec_mode, + max_num_requests=max_num_requests, + num_layers=model_config.num_hidden_layers, + hidden_size=model_config.hidden_size, + max_num_tokens=max_num_tokens, + allow_advanced_sampling=spec_config.allow_advanced_sampling, + use_rejection_sampling=use_rejection_sampling, + vocab_size=vocab_size, + spec_resource_manager=spec_resource_manager, + ) + if spec_config.spec_dec_mode.is_mtp_vanilla(): return MTPSpecMetadata( max_draft_len=spec_config.max_draft_len, max_total_draft_tokens=spec_config.tokens_per_gen_step - 1, @@ -185,16 +205,21 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): sa_manager = SuffixAutomatonManager(sa_cfg, max_num_requests, max_seq_len) if spec_config.use_relaxed_acceptance_for_thinking or sa_manager is not None: - return MTPHiddenStatesManager( + # Unified resource manager: the unified worker reads + # ``relaxed_delta_pool`` from ``Eagle3ResourceManager`` (mirrors the + # pool ``MTPHiddenStatesManager`` used to provide). + return Eagle3ResourceManager( spec_config, model_config.torch_dtype, model_config.hidden_size, max_num_requests, + max_seq_len, + max_num_tokens, sa_manager=sa_manager, ) else: return None - if spec_dec_mode.is_mtp_one_model(): + if spec_dec_mode.is_mtp_vanilla(): sa_manager = None sa_cfg = getattr(spec_config, 'sa_config', None) if sa_cfg is not None: @@ -263,7 +288,10 @@ def get_spec_decoder( sampler_args: TorchSampler.Args, spec_config: "DecodingBaseConfig", ): - if spec_config.spec_dec_mode.is_mtp_one_model(): + if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): + # MTP Eagle one-model now uses the same sampler as Eagle3 one-model. + return Eagle3OneModelSampler(sampler_args, spec_config=spec_config) + if spec_config.spec_dec_mode.is_mtp_vanilla(): return MTPSampler(sampler_args, nextn=spec_config.max_draft_len) if spec_config.spec_dec_mode.is_eagle3( ) or spec_config.spec_dec_mode.is_mtp_eagle(): @@ -314,7 +342,9 @@ def get_spec_drafter(model_engine, def get_num_spec_layers(spec_config): - if spec_config.spec_dec_mode.is_mtp_one_model(): + if spec_config.spec_dec_mode.is_mtp_eagle_one_model(): + return 1 + if spec_config.spec_dec_mode.is_mtp_vanilla(): return spec_config.num_nextn_predict_layers if spec_config.spec_dec_mode.is_eagle3_one_model(): num_eagle_layers = spec_config.num_eagle_layers @@ -336,8 +366,10 @@ def get_spec_worker(spec_config, if getattr(spec_config, 'use_dynamic_tree', False): return Eagle3OneModelDynamicTreeWorker(spec_config, mapping, use_separate_draft_kv_cache) - return Eagle3OneModelWorker(spec_config, mapping, - use_separate_draft_kv_cache) + return Eagle3OneModelWorker( + spec_config, + mapping=mapping, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) if spec_dec_mode.is_pard(): return PARDWorker(spec_config, mapping, use_separate_draft_kv_cache) if spec_dec_mode.is_dflash(): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2c3aad405890..87b5c89ea34e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + import ast import functools import json @@ -1588,13 +1603,13 @@ class MTPDecodingConfig(DecodingBaseConfig): description= "Enable relaxed acceptance during thinking phase for reasoning models. Accepts draft tokens matching any top-K candidate instead of exact top-1." ) - relaxed_topk: int = Field( + relaxed_topk: PositiveInt = Field( default=1, description= "Number of top candidate tokens to consider for relaxed acceptance. Draft token is accepted if it matches any of these." ) - relaxed_delta: float = Field( - default=0., + relaxed_delta: NonNegativeFloat = Field( + default=0.0, description= "Probability threshold for relaxed acceptance. Only candidates with prob >= (top-1 prob - delta) are kept." ) @@ -1625,12 +1640,12 @@ class MTPDecodingConfig(DecodingBaseConfig): "Auto-populated from the model's pretrained config. Do not set manually." ) - begin_thinking_phase_token: int = Field( + begin_thinking_phase_token: NonNegativeInt = Field( default=128798, description= "Token ID marking start of thinking phase. Relaxed acceptance only applies within this phase." ) - end_thinking_phase_token: int = Field( + end_thinking_phase_token: NonNegativeInt = Field( default=128799, description= "Token ID marking end of thinking phase. Strict acceptance resumes after this." @@ -1676,6 +1691,14 @@ def supports_backend(self, backend: str) -> bool: @property def num_capture_layers(self) -> int: + # MTP_EAGLE (two-model) feeds captured target hidden states into the + # separate draft engine, so the shared Eagle3ResourceManager must + # allocate a hidden_states buffer for it. MTP_EAGLE_ONE_MODEL passes + # the target model's hidden_states straight to the MTP layer + # (see Eagle3OneModelWorker.prepare_1st_drafter_inputs / _run_draft_forward, + # both gated on self.is_mtp_eagle), so no capture buffer is needed + # and we should skip allocation to avoid disabling post-MLP/MoE + # fusion via the layer-capture hook. return 1 if self.spec_dec_mode.is_mtp_eagle() else 0 @property