Additional small changes

Signed-off-by: Izzy Putterman <[email protected]>

Author: IzzyPutterman

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1049,6 +1049,7 @@ def forward(
         position_ids: Optional[torch.IntTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         spec_metadata: Optional[SpecMetadata] = None,
+        **kwargs,
     ) -> torch.Tensor:
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError(
@@ -1145,6 +1146,7 @@ def forward(
         )
 
         if spec_metadata and spec_metadata.spec_dec_mode.is_mtp():
+            # TODO Merge API with EagleWorker in modeling_speculative.py
             # get logits
             logits = self.logits_processor.forward(
                 hidden_states[spec_metadata.gather_ids],

tensorrt_llm/_torch/speculative/eagle3.py

@@ -186,6 +186,8 @@ def get_hidden_states(self):
 class Eagle3OneModelSpecMetadata(SpecMetadata):
     # The hidden states
     hidden_states: Optional[torch.Tensor] = None
+    # The number of layers to be captured
+    num_capture_layers: int = 3
     # The layers to be captured
     layers_to_capture: Tuple[int, ...] = field(init=False)
     # The hidden size of the hidden states
@@ -198,8 +200,8 @@ class Eagle3OneModelSpecMetadata(SpecMetadata):
     batch_indices_cuda: Optional[torch.Tensor] = None
 
     def __post_init__(self):
-        if self.num_layers == 1:
-            self.layers_to_capture = (1, )
+        if self.num_layers == 1 or self.num_capture_layers == 1:
+            self.layers_to_capture = (self.num_layers - 1, )
         else:
             if self.num_layers <= 5:
                 raise ValueError("Not enough hidden layers for EAGLE")