[None][feat] Eagle, use last hidden post norm (NVIDIA#7546)

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

Author: IzzyPutterman

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -410,6 +410,7 @@ def __init__(self, model: TModel, model_config: ModelConfig[TConfig]):
                     assert key in ('attn_layers', 'mla_layers')
                     assert key in model_config.extra_attrs
                     model_config.extra_attrs[key].update(value)
+        self.layer_idx = -1
 
     def forward(
         self,
@@ -430,6 +431,10 @@ def forward(
             **kwargs,
         )
 
+        if spec_metadata is not None and spec_metadata.is_layer_capture(
+                self.layer_idx):
+            spec_metadata.maybe_capture_hidden_states(self.layer_idx,
+                                                      hidden_states)
         if attn_metadata.padded_num_tokens is not None:
             hidden_states = hidden_states[:attn_metadata.num_tokens]
 

tensorrt_llm/_torch/speculative/eagle3.py

@@ -94,8 +94,10 @@ class Eagle3SpecMetadata(SpecMetadata):
     eagle3_resource_manager: Optional[Eagle3ResourceManager] = None
 
     def __post_init__(self):
-        if self.layers_to_capture is None:
-            if self.is_draft_model or self.num_layers == 1:
+        if self.is_draft_model:
+            self.layers_to_capture = (self.num_layers - 1, )
+        elif self.layers_to_capture is None:
+            if self.num_layers == 1:
                 self.layers_to_capture = (self.num_layers - 1, )
             else:
                 if self.num_layers <= 5: