From 3ef8a4639b198b4036ce00255b032ccbaa2ec665 Mon Sep 17 00:00:00 2001 From: Izzy Putterman Date: Mon, 2 Feb 2026 07:26:25 -0800 Subject: [PATCH] [None][feat] Nemotron H: Eagle3 support (#11131) Signed-off-by: Izzy Putterman --- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 249cab1547..d1a4974e74 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -361,6 +361,10 @@ class NemotronHLayer(DecoderLayer): spec_metadata=spec_metadata, **kwargs) hidden_states = torch.add(hidden_states, residual) + 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, None) return hidden_states