[TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers

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

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

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)

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