[None][feat] Eagle 3.1 -- Support post-norm and per-aux fc_norm for Eagle3 draft models (#14988)

Signed-off-by: Doğaç Eldenk <dogacel@gmail.com>
Co-authored-by: Laikh Tewari <ltewari@nvidia.com>

Author: Dogacel

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -307,7 +307,8 @@ def __init__(
             self.hidden_size_in = config.hidden_size
 
         self._return_hidden_post_norm = eagle_config.get(
-            "return_hidden_post_norm", False)
+            "return_hidden_post_norm", False) or getattr(
+                config, "norm_output", False)
 
         # Create auxiliary CUDA stream for MLA operations (only needed for MLA)
         self.aux_stream = torch.cuda.Stream() if use_mla else None
@@ -330,6 +331,18 @@ def __init__(
         else:
             self.input_norm = None
 
+        self._use_fc_norm = getattr(config, "fc_norm", False)
+        if self._use_fc_norm:
+            self.fc_norm = nn.ModuleList([
+                RMSNorm(
+                    hidden_size=self.hidden_size_in,
+                    eps=config.rms_norm_eps,
+                    dtype=config.torch_dtype,
+                ) for _ in range(self.spec_config.num_capture_layers)
+            ])
+        else:
+            self.fc_norm = None
+
         if self.num_layers > 1:
             self.midlayer = nn.ModuleList([
                 Eagle3DecoderLayer(
@@ -590,7 +603,14 @@ def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         expected_hidden_size = self.model.hidden_size
         if hidden_states.shape[-1] != expected_hidden_size:
-            if self.model._norm_before_fc:
+            if self.model.fc_norm is not None:
+                chunks = hidden_states.chunk(len(self.model.fc_norm), dim=-1)
+                hidden_states = torch.cat([
+                    norm(chunk)
+                    for norm, chunk in zip(self.model.fc_norm, chunks)
+                ],
+                                          dim=-1)
+            elif self.model._norm_before_fc:
                 hidden_states = self.model.input_norm(hidden_states)
             hidden_states = self.model.fc(hidden_states)