From 8a60e9de900e258923280cfa0b329d94303a911d Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 19 Nov 2025 16:12:09 -0800
Subject: [PATCH 01/18] Init commit for EagleDecodingConfig support in
 AutoDeploy. Adds: file for running AD with Eagle1, modifies existing TRTLLM
 example to use Eagle1 instead of Eagle3 for comparison. Eagle1 is chosen
 (last hidden layer only) to resemble MTPEagle

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/config/default.yaml    |   5 +
 tensorrt_llm/_torch/auto_deploy/llm_args.py   |  47 +++-
 .../_torch/auto_deploy/shim/ad_executor.py    | 173 +++++++++++-
 .../transform/library/hidden_states.py        | 259 ++++++++++++++++++
 tensorrt_llm/_torch/pyexecutor/py_executor.py |   8 +
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  17 +-
 .../examples/test_ad_speculative_decoding.py  |  65 +++--
 .../test_lists/test-db/l0_h100.yml            |   4 +-
 8 files changed, 535 insertions(+), 43 deletions(-)
 create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py

diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
index 68c5086226b8..e718ed5c9d97 100644
--- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml
+++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
@@ -75,6 +75,8 @@ transforms:
     stage: pattern_matcher
   quantize_mxfp4_moe:
     stage: pattern_matcher
+  detect_hidden_states_for_capture:
+    stage: pattern_matcher
   detect_sharding:
     stage: sharding
     simple_shard_only: false
@@ -163,6 +165,9 @@ transforms:
   insert_cached_delta_rule:
     stage: cache_init
     backend: fla_delta
+  insert_cached_residual_add:
+    stage: cache_init
+    backend: cached_residual_add
   initialize_cache:
     stage: cache_init
     run_per_gm: false
diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py
index ddaa64c3e2ad..ac9440f774e5 100644
--- a/tensorrt_llm/_torch/auto_deploy/llm_args.py
+++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py
@@ -1,6 +1,6 @@
 from importlib.resources import files
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Type, Union
+from typing import Any, Dict, List, Literal, Optional, Set, Type, Union
 
 import torch
 from pydantic import Field, PrivateAttr, ValidationInfo, field_validator, model_validator
@@ -8,7 +8,14 @@
 
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
-from ...llmapi.llm_args import BaseLlmArgs, BuildConfig, KvCacheConfig, SamplerType, _ParallelConfig
+from ...llmapi.llm_args import (
+    BaseLlmArgs,
+    BuildConfig,
+    EagleDecodingConfig,
+    KvCacheConfig,
+    SamplerType,
+    _ParallelConfig,
+)
 from .models import ModelFactory, ModelFactoryRegistry
 from .utils._config import DynamicYamlMixInForSettings
 from .utils.logger import ad_logger
@@ -38,6 +45,12 @@ def _check_for_default_value_only(
     return value
 
 
+def default_eagle3_layers_to_capture(num_hidden_layers: int) -> Set[int]:
+    if num_hidden_layers <= 5:
+        raise ValueError("Not enough hidden layers for default EAGLE3 capture")
+    return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
+
+
 _TRANSFORMS_SHORTCUT_LOOKUP = {
     "attn_backend": ("insert_cached_attention.backend", "transformers_replace_cached_attn.backend"),
     "free_mem_ratio": ("resize_kv_cache.free_mem_ratio",),
@@ -150,6 +163,11 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings):
 
     enable_chunked_prefill: bool = Field(default=False, description="Enable chunked prefill.")
 
+    draft_checkpoint_loader: Optional[object] = Field(
+        default=None,
+        description="The checkpoint loader to use for the draft model when using speculative decoding with two models.",
+    )
+
     ### INFERENCE OPTIMIZER CONFIG #################################################################
     mode: Literal["graph", "transformers"] = Field(
         default="graph",
@@ -190,11 +208,6 @@ class AutoDeployConfig(DynamicYamlMixInForSettings, BaseSettings):
         ),
     )
 
-    draft_checkpoint_loader: Optional[object] = Field(
-        default=None,
-        description="The checkpoint loader to use for the draft model when using speculative decoding with two models.",
-    )
-
     ### SEQUENCE INTERFACE CONFIG ##################################################################
     max_input_len: int = Field(default=1024, description="The maximum input length.")
     max_num_tokens: Optional[int] = Field(default=None, description="The maximum number of tokens.")
@@ -420,6 +433,26 @@ def ensure_no_custom_parallel_config(cls, value: Any, info: ValidationInfo) -> A
         msg = "AutoDeploy only supports parallelization via the `world_size` argument."
         return _check_for_default_value_only(cls, value, info, msg)
 
+    @model_validator(mode="after")
+    def default_eagle3_layers_to_capture(self):
+        if self.speculative_config is None or not isinstance(
+            self.speculative_config, EagleDecodingConfig
+        ):
+            return self
+
+        if self.speculative_config.eagle3_layers_to_capture is None:
+            num_hidden_layers = self.create_factory()._get_model_config()[0].num_hidden_layers
+            self.speculative_config.eagle3_layers_to_capture = default_eagle3_layers_to_capture(
+                num_hidden_layers
+            )
+
+        # insert the layers to capture into the transforms config.
+        self.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] = (
+            self.speculative_config.eagle3_layers_to_capture
+        )
+
+        return self
+
     @model_validator(mode="after")
     def validate_parallel_config(self):
         """Setup parallel config according to world_size.
diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 9ecf76405b90..d01e7e1def0b 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -13,10 +13,11 @@
 import types
 from collections import defaultdict
 from dataclasses import dataclass
-from types import SimpleNamespace
+from types import MethodType, SimpleNamespace
 from typing import Dict, List, Optional, Tuple
 
 import torch
+import torch.nn.functional as F
 from strenum import StrEnum
 from torch._prims_common import DeviceLikeType
 
@@ -32,9 +33,11 @@
 from tensorrt_llm._torch.pyexecutor.py_executor_creator import get_guided_decoding_config
 from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager
 from tensorrt_llm._torch.speculative import get_spec_drafter
+from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager
 from tensorrt_llm._utils import nvtx_range
 from tensorrt_llm.llmapi.llm_args import (
     ContextChunkingPolicy,
+    EagleDecodingConfig,
     LoadFormat,
     SamplerType,
     TorchLlmArgs,
@@ -57,6 +60,7 @@
 from ...pyexecutor.scheduler import (
     BindCapacityScheduler,
     BindMicroBatchScheduler,
+    RequestList,
     ScheduledRequests,
     SimpleScheduler,
 )
@@ -113,6 +117,90 @@ def calculate_max_num_blocks(
         return self.num_blocks, 0
 
 
+class ADHiddenStateManager(Eagle3ResourceManager):
+    def __init__(
+        self,
+        cache_seq_interface: CachedSequenceInterface,
+        config: EagleDecodingConfig,
+        max_num_requests: int,
+        max_seq_len: int,
+        max_num_tokens: int,
+    ):
+        hidden_state_buffer = self._get_hidden_state_buffers(cache_seq_interface)[0]
+        dtype = hidden_state_buffer.dtype
+        hidden_size = hidden_state_buffer.shape[1]
+
+        super().__init__(config, dtype, hidden_size, max_num_requests, max_seq_len, max_num_tokens)
+
+        self.hidden_state_write_indices: torch.Tensor = torch.empty(
+            max_num_tokens, dtype=torch.long, device="cuda"
+        )
+
+    def _get_hidden_state_buffers(
+        self, cache_seq_interface: CachedSequenceInterface
+    ) -> List[torch.Tensor]:
+        hidden_state_buffers = []
+        for name, tensor in cache_seq_interface.named_args.items():
+            if "hidden_states_cache" in name:
+                hidden_state_buffers.append(tensor)
+
+        if not hidden_state_buffers:
+            raise ValueError(
+                "No hidden_state_buffers found in cache_seq_interface. Check if we are actually running Eagle3."
+            )
+        return hidden_state_buffers
+
+    def prepare_hidden_states_capture(
+        self, ordered_requests: RequestList, cache_seq_interface: CachedSequenceInterface
+    ) -> None:
+        """Prepare the hidden states for capture by establishing indices that the hidden states will be written to."""
+        seq_lens = cache_seq_interface.info.seq_len
+        num_tokens = sum(seq_lens)
+
+        start_idx = 0
+        hidden_states_write_indices = []
+        for request, seq_len in zip(ordered_requests, seq_lens):
+            request_id = request.request_id
+            slot_id = self.slot_manager.get_slot(request_id)
+            self.start_indices[slot_id] = start_idx
+            hidden_states_write_indices.extend(range(start_idx, start_idx + seq_len))
+            start_idx += max(seq_len, self.max_total_draft_tokens + 1)
+            assert start_idx < self.hidden_states.shape[0], (
+                f"start_idx {start_idx} exceeds hidden_states capacity {self.hidden_states.shape[0]}"
+            )
+
+        if len(hidden_states_write_indices) != num_tokens:
+            raise ValueError(
+                f"len(hidden_state_write_indices) ({len(hidden_states_write_indices)}) != num_tokens \
+                ({num_tokens}). Check whether ordered_requests matches up with seq_lens."
+            )
+
+        hidden_state_write_indices_host = torch.tensor(
+            hidden_states_write_indices, dtype=torch.long
+        )
+
+        self.hidden_state_write_indices[:num_tokens].copy_(
+            hidden_state_write_indices_host, non_blocking=True
+        )
+
+    def capture_hidden_states(self, cache_seq_interface: CachedSequenceInterface) -> None:
+        """Capture configured hidden states that have been written by the model,
+        in a format that can be used by the draft model.
+        """
+        full_hidden_states = self._get_hidden_state_buffers(cache_seq_interface)
+        if not full_hidden_states:
+            return
+
+        num_tokens = sum(cache_seq_interface.info.seq_len)
+
+        hidden_states = [hidden_state[:num_tokens] for hidden_state in full_hidden_states]
+        hidden_states = torch.cat(hidden_states, dim=1) if hidden_states else None
+        hidden_states = hidden_states.to(dtype=self.dtype)
+
+        token_idx = self.hidden_state_write_indices[:num_tokens]
+        self.hidden_states[:, : hidden_states.shape[1]].index_copy_(0, token_idx, hidden_states)
+
+
 def construct_draft_llm_args(
     ad_config: LlmArgs,
 ) -> TorchLlmArgs:
@@ -461,6 +549,10 @@ def _prepare_inputs(
         kv_cache_manager = resource_manager.get_resource_manager(
             ResourceManagerType.KV_CACHE_MANAGER
         )
+        # resource manager for hidden state capture
+        spec_resource_manager = resource_manager.get_resource_manager(
+            ResourceManagerType.SPEC_RESOURCE_MANAGER
+        )
 
         # requests in order of context, generate
         context_requests = scheduled_requests.context_requests
@@ -471,6 +563,7 @@ def _prepare_inputs(
             r for r in scheduled_requests.generation_requests if get_draft_token_length(r) == 0
         ]
         gen_requests = extend_requests + generation_requests
+        ordered_requests = context_requests + gen_requests
         # info to be extracted
         input_ids: List[List[int]] = []
         position_ids: List[List[int]] = []
@@ -670,17 +763,32 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]:
 
         self.cache_seq_interface.info.run_host_prepare_for_attention_forward()
 
+        if spec_resource_manager is not None and isinstance(
+            spec_resource_manager, ADHiddenStateManager
+        ):
+            spec_resource_manager.prepare_hidden_states_capture(
+                ordered_requests, self.cache_seq_interface
+            )
+
         self.iter_states["num_ctx_requests"] = num_ctx_requests
         self.iter_states["num_ctx_tokens"] = num_ctx_tokens
         # TODO: handle extend requests and draft requests for specdec
         self.iter_states["num_generation_tokens"] = num_generation_tokens
 
     @nvtx_range("ad_compute_logits")
-    def _compute_logits(self) -> List[torch.Tensor]:
+    def _compute_logits(self, resource_manager: ResourceManager) -> List[torch.Tensor]:
         # run the model
         logits: torch.Tensor = self.model(**self.cache_seq_interface.named_args)[0]
         logits = self.cache_seq_interface.info.maybe_gather_and_squeeze_logits(logits)
 
+        spec_resource_manager = resource_manager.get_resource_manager(
+            ResourceManagerType.SPEC_RESOURCE_MANAGER
+        )
+        if spec_resource_manager is not None and isinstance(
+            spec_resource_manager, ADHiddenStateManager
+        ):
+            spec_resource_manager.capture_hidden_states(self.cache_seq_interface)
+
         # TRTLLMSampler expects float32 logits. PyTorchModelEngine always casts to float32 regardless.
         return logits.float()
 
@@ -708,7 +816,7 @@ def forward(
         self.iter_counter += 1
 
         outputs = {
-            "logits": self._compute_logits(),
+            "logits": self._compute_logits(resource_manager),
         }
         if self.mapping is not None:
             self._execute_logit_post_processors(scheduled_requests, outputs)
@@ -716,8 +824,30 @@ def forward(
         return outputs
 
 
+def share_embedding_weights(
+    target_model_engine: "ADEngine", draft_model_engine: PyTorchModelEngine
+):
+    # This function is necessary for supporting Eagle and other speculative decoding methods that
+    # copy the embed_tokens submodule. It is not necessary for MTP and other speculative decoding methods that
+    # use the draft model engine directly.
+
+    submodule = target_model_engine.model.model.embed_tokens
+
+    world_size = mpi_world_size()
+    assert world_size <= 1, f"This code assumes tp<=1. World size: {world_size}"
+
+    # Note: This simple forward function implementation assumes tp=1.
+    # TODO(govind): Handle the tp>1 case.
+    def new_embedding_forward(self, input_ids):
+        return F.embedding(input_ids, self.weight)
+
+    submodule.forward = MethodType(new_embedding_forward, submodule)
+
+    draft_model_engine.load_weights_from_target_model(target_model_engine.model)
+
+
 def create_draft_model_engine_maybe(
-    ad_config: LlmArgs, engine, dist_mapping: Mapping, mpi_dist: MPIDist
+    ad_config: LlmArgs, target_engine: ADEngine, dist_mapping: Mapping, mpi_dist: MPIDist
 ) -> Optional[PyTorchModelEngine]:
     """Create a draft model engine for speculative decoding.
 
@@ -745,7 +875,7 @@ def create_draft_model_engine_maybe(
         chunked_prefill=ad_config.enable_chunked_prefill,
         cache_reuse=kv_cache_config.enable_block_reuse,
         has_speculative_draft_tokens=has_spec_drafter,
-        chunk_size=engine.llm_args.max_num_tokens,
+        chunk_size=target_engine.llm_args.max_num_tokens,
     )
 
     # Construct TorchLlmArgs for the draft model
@@ -753,6 +883,10 @@ def create_draft_model_engine_maybe(
         ad_config=ad_config,
     )
 
+    # chain drafter is not supported currently for AutoDeploy.
+    # TODO(govind): Do this when we want to optimize 2-model spec dec performance.
+    drafting_loop_wrapper = None
+
     draft_model_engine = PyTorchModelEngine(
         model_path=draft_spec_config.speculative_model_dir,
         llm_args=draft_llm_args,
@@ -761,7 +895,11 @@ def create_draft_model_engine_maybe(
         dist=mpi_dist,
         spec_config=draft_spec_config,
         is_draft_model=True,
-        drafting_loop_wrapper=None,
+        drafting_loop_wrapper=drafting_loop_wrapper,
+    )
+
+    share_embedding_weights(
+        target_model_engine=target_engine, draft_model_engine=draft_model_engine
     )
 
     draft_model_engine.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER
@@ -855,9 +993,11 @@ def create_autodeploy_executor(ad_config: LlmArgs, tokenizer: Optional[Tokenizer
     engine = ADEngine.build_from_config(ad_config=ad_config, mapping=dist_mapping)
 
     spec_config = ad_config.speculative_config
-    if spec_config is not None and not spec_config.spec_dec_mode.is_draft_target():
+    if spec_config is not None and not (
+        spec_config.spec_dec_mode.is_draft_target() or spec_config.spec_dec_mode.is_eagle3()
+    ):
         raise ValueError(
-            "Currently, AutoDeploy only supports speculative decoding in draft target mode."
+            "Currently, AutoDeploy only supports speculative decoding in draft target or eagle3 mode."
         )
 
     if spec_config is not None and ad_config.guided_decoding_backend is not None:
@@ -865,11 +1005,20 @@ def create_autodeploy_executor(ad_config: LlmArgs, tokenizer: Optional[Tokenizer
             "Guided decoding is not currently supported for speculative decoding in AutoDeploy."
         )
 
-    # Speculative resource manager not needed for DraftTargetDecoding.
-    spec_resource_manager = None
-
     draft_model_engine = create_draft_model_engine_maybe(
-        ad_config=ad_config, engine=engine, dist_mapping=dist_mapping, mpi_dist=mpi_dist
+        ad_config=ad_config, target_engine=engine, dist_mapping=dist_mapping, mpi_dist=mpi_dist
+    )
+
+    spec_resource_manager = (
+        ADHiddenStateManager(
+            cache_seq_interface=engine.cache_seq_interface,
+            config=spec_config,
+            max_num_requests=ad_config.max_batch_size,
+            max_seq_len=engine.llm_args.max_seq_len,
+            max_num_tokens=engine.llm_args.max_num_tokens,
+        )
+        if isinstance(spec_config, EagleDecodingConfig)
+        else None
     )
 
     # check kvcache config for partial block reuse
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
new file mode 100644
index 000000000000..b7e698d83a0b
--- /dev/null
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -0,0 +1,259 @@
+"""The transform passes to capture the hidden states of the target model."""
+
+from typing import Dict, List, Optional, Set, Tuple, Type
+
+import torch
+from torch._ops import OpOverloadPacket
+from torch.fx import GraphModule, Node
+
+from ...custom_ops.attention_interface import (
+    AttentionDescriptor,
+    AttentionLayout,
+    AttentionRegistry,
+    BufferInitializerDict,
+    CacheConfig,
+    CacheInitializerDict,
+    Constant,
+    MHACallable,
+    PrepareMetadataCallable,
+    SequenceInfo,
+)
+from ...models.factory import ModelFactory
+from ...shim.interface import CachedSequenceInterface
+from ...utils.node_utils import get_all_layer_subgraphs, is_op
+from ..interface import (
+    BaseTransform,
+    SharedConfig,
+    TransformConfig,
+    TransformInfo,
+    TransformRegistry,
+)
+from .kvcache import InsertCachedAttention
+
+
+@torch.library.custom_op("auto_deploy::residual_add_for_capture", mutates_args=())
+def residual_add_for_capture(t1: torch.Tensor, t2: torch.Tensor) -> torch.Tensor:
+    return torch.ops.aten.add(t1, t2)
+
+
+@residual_add_for_capture.register_fake
+def residual_add_for_capture_fake(t1: torch.Tensor, t2: torch.Tensor) -> torch.Tensor:
+    return torch.ops.aten.add(t1, t2)
+
+
+@torch.library.custom_op("auto_deploy::cached_residual_add", mutates_args=())
+def cached_residual_add(
+    t1: torch.Tensor, t2: torch.Tensor, hidden_states_cache: torch.Tensor
+) -> torch.Tensor:
+    ret = torch.ops.aten.add(t1, t2)
+    b, s, _ = ret.shape
+    print(f"In cached residual add. Ret shape: {ret.shape}")
+    print(f"Shape of hidden_states_cache: {hidden_states_cache.shape}")
+    num_tokens = b * s
+    print(f"Num tokens: {num_tokens}")
+
+    # TODO(govind): do some of these correspond to padding tokens when there are varying sequence lengths?
+    # Might need to extract the actual sequence lengths from somewhere to get the appropriate indices to copy.
+    hidden_states_cache[:num_tokens].copy_(ret.view(num_tokens, -1), non_blocking=True)
+    return ret
+
+
+@cached_residual_add.register_fake
+def cached_residual_add_fake(
+    t1: torch.Tensor, t2: torch.Tensor, hidden_states_cache: torch.Tensor
+) -> torch.Tensor:
+    return torch.ops.aten.add(t1, t2)
+
+
+@torch.library.custom_op("auto_deploy::cached_residual_add_prepare_metadata", mutates_args=())
+def cached_residual_add_prepare_metadata(
+    position_ids: torch.Tensor,
+    seq_len: torch.Tensor,
+    input_pos: torch.Tensor,
+    cache_loc: torch.Tensor,
+    pages_per_seq: torch.Tensor,
+    slot_idx: torch.Tensor,
+    page_size: int,
+    chunk_size: int,
+) -> List[torch.Tensor]:
+    return [
+        position_ids,
+        seq_len,
+        input_pos,
+        cache_loc,
+        pages_per_seq,
+        slot_idx,
+        page_size,
+        chunk_size,
+    ]
+
+
+@cached_residual_add_prepare_metadata.register_fake
+def cached_residual_add_prepare_metadata_fake(
+    position_ids: torch.Tensor,
+    seq_len: torch.Tensor,
+    input_pos: torch.Tensor,
+    cache_loc: torch.Tensor,
+    pages_per_seq: torch.Tensor,
+    slot_idx: torch.Tensor,
+    page_size: int,
+    chunk_size: int,
+) -> List[torch.Tensor]:
+    return [
+        position_ids,
+        seq_len,
+        input_pos,
+        cache_loc,
+        pages_per_seq,
+        slot_idx,
+        page_size,
+        chunk_size,
+    ]
+
+
+class DetectHiddenStatesForCaptureConfig(TransformConfig):
+    """Configuration for the hidden states detection transform."""
+
+    # TODO: figure out how to get layers to capture.
+    # Right now default is None and EagleSpecMetadata has a heuristic to extract layer indices to capture.
+    # This seems fragile.
+    # We should consider if we can use the layer indices stored in eagle checkpoints, e.g.
+    # https://huggingface.co/nvidia/gpt-oss-120b-Eagle3/blob/main/config.json#L9-L14
+    eagle3_layers_to_capture: Optional[Set[int]] = None  # Default: Do not capture any layers
+
+
+@TransformRegistry.register("detect_hidden_states_for_capture")
+class DetectHiddenStatesForCapture(BaseTransform):
+    """Detect the hidden states we should capture in the graph."""
+
+    config: DetectHiddenStatesForCaptureConfig
+
+    @classmethod
+    def get_config_class(cls) -> Type[TransformConfig]:
+        return DetectHiddenStatesForCaptureConfig
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        if not self.config.eagle3_layers_to_capture:
+            info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
+            return gm, info
+
+        def _get_layer_number(lin_node: Node) -> Optional[int]:
+            weight = lin_node.args[1]
+            print(f"Calling _get_layer_number() with lin_node: {lin_node}")
+            if weight.op == "get_attr":
+                subnames = weight.target.split(".")
+                for subname in subnames:
+                    if subname.isdigit():
+                        print(f"Found layer number: {int(subname)}")
+                        return int(subname)
+
+            print("No layer number found")
+            return None
+
+        # find last closing linear node of each layer
+        # from there we will find the residual add node for that layer
+        layer_subgraphs, unprocessed_linear_nodes = get_all_layer_subgraphs(gm)
+        residual_add_nodes: Dict[int, Node] = {}
+        for _, _, lin_node_closing in layer_subgraphs:
+            # need layer number to correctly identify the residual add node
+            layer_number = _get_layer_number(lin_node_closing)
+            if layer_number is None or layer_number not in self.config.eagle3_layers_to_capture:
+                continue
+
+            # Conditions to identify as the hidden states after the residual
+            # The first node after the linear closing node that satisfies:
+            # 1. is an add node with > 1 users (hidden states before the last are used directly by next layer
+            # as well as having a residual add to the next hidden state). Stopping here prevents us from
+            # using a future residual add node for the next layer.
+            # 2. is the last add node in a 1 user chain (for last layer or layers with no following residual add)
+            # This stops us before we go to the next layer.
+            res_node = lin_node_closing
+            while len(res_node.users) == 1:
+                user_node = list(res_node.users)[0]
+                if not is_op(user_node, torch.ops.aten.add):
+                    break
+                res_node = user_node
+
+            if is_op(res_node, torch.ops.aten.add):
+                # this stores the last residual add node encountered for each layer
+                residual_add_nodes[layer_number] = res_node
+
+        assert residual_add_nodes.keys() == self.config.eagle3_layers_to_capture, (
+            f"Unable to find residual add nodes for layers. Expected: {self.config.eagle3_layers_to_capture}, \
+            Found: {residual_add_nodes.keys()}"
+        )
+
+        # replace residual add nodes with special placeholder nodes
+        for layer_number, res_node in residual_add_nodes.items():
+            with gm.graph.inserting_before(res_node):
+                new_node = gm.graph.call_function(
+                    torch.ops.auto_deploy.residual_add_for_capture.default,
+                    args=res_node.args,
+                    kwargs=res_node.kwargs,
+                )
+            res_node.replace_all_uses_with(new_node)
+            gm.graph.erase_node(res_node)
+
+        cnt = len(residual_add_nodes)
+        info = TransformInfo(
+            skipped=False, num_matches=cnt, is_clean=(cnt == 0), has_valid_shapes=(cnt == 0)
+        )
+        return gm, info
+
+
+@AttentionRegistry.register("cached_residual_add")
+class CachedResidualAdd(AttentionDescriptor):
+    @classmethod
+    def is_paged(cls) -> bool:
+        return True
+
+    @classmethod
+    def get_attention_layout(cls) -> AttentionLayout:
+        return "bsnd"
+
+    @classmethod
+    def get_num_qkv_args(cls) -> int:
+        return 2
+
+    @classmethod
+    def get_source_attention_op(cls) -> OpOverloadPacket:
+        return torch.ops.auto_deploy.residual_add_for_capture
+
+    @classmethod
+    def get_cached_attention_op(cls) -> MHACallable:
+        return torch.ops.auto_deploy.cached_residual_add
+
+    @classmethod
+    def get_prepare_metadata_op(cls) -> Tuple[PrepareMetadataCallable, int]:
+        return torch.ops.auto_deploy.cached_residual_add_prepare_metadata, 0
+
+    @classmethod
+    def get_cache_initializers(
+        cls, source_attn_node: Node, cache_config: CacheConfig
+    ) -> CacheInitializerDict:
+        hidden_size = source_attn_node.meta["val"].shape[-1]
+        hidden_type = source_attn_node.meta["val"].dtype
+
+        def _get_hidden_states_cache(si: SequenceInfo):
+            return torch.empty(si.max_num_tokens, hidden_size, device=si.device, dtype=hidden_type)
+
+        return {"hidden_states_cache": _get_hidden_states_cache}
+
+    @classmethod
+    def get_global_buffer_initializers(cls, source_attn_node: Node) -> BufferInitializerDict:
+        return {}
+
+    @classmethod
+    def get_constants(cls, source_attn_node: Node) -> List[Constant]:
+        return []
+
+
+@TransformRegistry.register("insert_cached_residual_add")
+class InsertCachedResidualAdd(InsertCachedAttention):
+    """A transform to handle residual add cache operations."""
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 6ccae36fdcd6..9203d21229ed 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -1757,6 +1757,14 @@ def _accept_draft_tokens(
             batch_indices = torch.arange(batch_size, device=device)
             new_tokens[0, :, 0] = target_tokens[num_accepted_tokens,
                                                 batch_indices]
+
+            # Print speculative acceptance counts each iteration
+            accepted = num_accepted_tokens.detach().cpu().tolist()
+            print("[spec-decode] iter=%s batch=%d ctx=%d draft_len=%d "
+                  "accepted_draft_tokens=%s (draft+1 is the target token)" %
+                  (getattr(self, "iter_counter", None), batch_size,
+                   len(scheduled_batch.context_requests), max_draft_len,
+                   accepted))
         else:
             # No draft tokens to accept, just use the first (and only) sampled token
             batch_indices = torch.arange(batch_size, device=device)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index c0bb0acf785b..0b239f470c0b 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -1476,12 +1476,25 @@ def process_draft_tokens(
                 num_accepted = self._process_draft_tokens_greedy(
                     request, new_tokens=new_tokens_list, finish_reasons=finish_reasons
                 )
-            return num_accepted
         else:
-            return self._process_draft_tokens_rejection_sampling(
+            num_accepted = self._process_draft_tokens_rejection_sampling(
                 request, new_tokens_list=new_tokens_list, new_tokens_tensor=new_tokens_tensor
             )
 
+        # Print draft acceptance info each iteration (non-overlap scheduler path)
+        draft_len = get_draft_token_length(request)
+        print(
+            "[spec-decode] req=%s iter=%s draft_len=%d accepted=%d"
+            % (
+                getattr(request, "py_request_id", None),
+                getattr(request, "py_decoding_iter", None),
+                draft_len,
+                num_accepted,
+            )
+        )
+
+        return num_accepted
+
     def _get_logprobs_from_request(self, request: LlmRequest) -> tuple[torch.Tensor, torch.Tensor]:
         """Extract the logprobs from the request
 
diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index 1c328863ac52..2374abec8198 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -19,7 +19,7 @@
 from build_and_run_ad import ExperimentConfig, main
 from defs.conftest import llm_models_root
 
-from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig
+from tensorrt_llm.llmapi import DraftTargetDecodingConfig, EagleDecodingConfig, KvCacheConfig
 
 prompts = [
     "What is the capital of France?",
@@ -28,31 +28,51 @@
     "What is the highest mountain in the world?",
 ]
 
+EAGLE_MODEL_SUBPATH = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
+LLAMA_BASE_SUBPATH = "llama-3.1-model/Llama-3.1-8B-Instruct"
+DRAFT_TARGET_MAX_DRAFT_LEN = 3
+EAGLE_MAX_DRAFT_LEN = 3
+
 
 def get_model_paths():
     """Get model paths using llm_models_root()."""
     models_root = llm_models_root()
-    base_model = os.path.join(
-        models_root,
-        "llama-3.1-model/Llama-3.1-8B-Instruct",
-    )
-    speculative_model = os.path.join(
+    base_model = os.path.join(models_root, LLAMA_BASE_SUBPATH)
+    draft_target_model = os.path.join(
         models_root,
         "llama-models-v2/TinyLlama-1.1B-Chat-v1.0",
     )
+    eagle_model = os.path.join(models_root, EAGLE_MODEL_SUBPATH)
 
     print(f"Base model path: {base_model}")
-    print(f"Speculative model path: {speculative_model}")
-    return base_model, speculative_model
+    print(f"DraftTarget draft model path: {draft_target_model}")
+    print(f"EAGLE model path: {eagle_model}")
+    return base_model, draft_target_model, eagle_model
+
+
+def make_spec_config(spec_dec_mode: str, spec_model_path: str):
+    if spec_dec_mode == "draft_target":
+        return DraftTargetDecodingConfig(
+            max_draft_len=DRAFT_TARGET_MAX_DRAFT_LEN, speculative_model_dir=spec_model_path
+        )
+    if spec_dec_mode == "eagle":
+        return EagleDecodingConfig(
+            max_draft_len=EAGLE_MAX_DRAFT_LEN,
+            speculative_model_dir=spec_model_path,
+            eagle3_one_model=False,
+            eagle3_layers_to_capture=None,
+        )
+    raise ValueError(f"Unknown speculative mode: {spec_dec_mode}")
 
 
-def run_with_autodeploy(model, speculative_model_dir, batch_size):
+def run_with_autodeploy(model, speculative_model_dir, batch_size, spec_dec_mode: str | None):
     """Run AutoDeploy with or without speculative decoding.
 
     Args:
         model: Path to the base model
         speculative_model_dir: Path to the speculative model (None for baseline mode)
         batch_size: Number of prompts to process
+        spec_dec_mode: Speculative decoding mode
 
     Returns:
         List of (prompt, output) tuples from prompts_and_outputs
@@ -62,10 +82,8 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
 
     # Configure speculative decoding if speculative_model_dir is provided
     spec_config = None
-    if speculative_model_dir is not None:
-        spec_config = DraftTargetDecodingConfig(
-            max_draft_len=3, speculative_model_dir=speculative_model_dir
-        )
+    if speculative_model_dir is not None and spec_dec_mode is not None:
+        spec_config = make_spec_config(spec_dec_mode, speculative_model_dir)
 
     # Configure KV cache
     kv_cache_config = KvCacheConfig(
@@ -76,7 +94,6 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     llm_args = {
         "model": model,
         "skip_loading_weights": False,
-        "speculative_config": spec_config,
         "runtime": "trtllm",
         "world_size": 1,
         "kv_cache_config": kv_cache_config,
@@ -100,6 +117,10 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     # Create ExperimentConfig
     cfg = ExperimentConfig(**experiment_config)
 
+    cfg.args.speculative_config = (
+        spec_config  # Add here to avoid Pydantic validation error for eagle3_layers_to_capture
+    )
+
     # Add sampling parameters (deterministic with temperature=0.0 and fixed seed)
     cfg.prompt.sp_kwargs = {
         "max_tokens": 50,
@@ -116,8 +137,8 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     return result["prompts_and_outputs"]
 
 
-@pytest.mark.parametrize("batch_size", [1, 4])
-def test_autodeploy_spec_dec(batch_size):
+@pytest.mark.parametrize("batch_size, spec_dec_mode", [(1, "draft_target"), (4, "eagle")])
+def test_autodeploy_spec_dec(batch_size, spec_dec_mode):
     """Test AutoDeploy speculative decoding with different batch sizes.
 
     Runs with and without speculative decoding and verifies outputs are identical.
@@ -126,23 +147,27 @@ def test_autodeploy_spec_dec(batch_size):
     print(f"Testing AutoDeploy Speculative Decoding - Batch Size {batch_size}")
     print("=" * 80)
 
-    base_model, speculative_model = get_model_paths()
+    base_model, draft_target_model, eagle_model = get_model_paths()
 
     print(f"\nBase Model: {base_model}")
-    print(f"Speculative Model: {speculative_model}")
+    spec_model_path = draft_target_model if spec_dec_mode == "draft_target" else eagle_model
+    print(f"Speculative Model: {spec_model_path}")
     print(f"Batch Size: {batch_size}")
 
     # Run with speculative decoding
     print("\n[1/2] Running with speculative decoding enabled...")
     spec_outputs = run_with_autodeploy(
-        model=base_model, speculative_model_dir=speculative_model, batch_size=batch_size
+        model=base_model,
+        speculative_model_dir=spec_model_path,
+        batch_size=batch_size,
+        spec_dec_mode=spec_dec_mode,
     )
     print(f"Generated {len(spec_outputs)} outputs with speculative decoding")
 
     # Run without speculative decoding (baseline)
     print("\n[2/2] Running without speculative decoding (baseline)...")
     baseline_outputs = run_with_autodeploy(
-        model=base_model, speculative_model_dir=None, batch_size=batch_size
+        model=base_model, speculative_model_dir=None, batch_size=batch_size, spec_dec_mode=None
     )
     print(f"Generated {len(baseline_outputs)} outputs in baseline mode")
 
diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml
index 26767235acee..55c5b60d4ea7 100644
--- a/tests/integration/test_lists/test-db/l0_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_h100.yml
@@ -116,8 +116,8 @@ l0_h100:
   - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[True]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_bf16
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[1]
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[4]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[1-draft_target]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[4-eagle]
 - condition:
     ranges:
       system_gpu_count:

From d4d16fb4bc481047b709d50e650e0ff82ea15c3f Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 10 Dec 2025 18:04:38 -0800
Subject: [PATCH 02/18] removing print statements for num accepted tokens

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  8 --------
 tensorrt_llm/_torch/pyexecutor/sampler.py     | 17 ++---------------
 2 files changed, 2 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 9203d21229ed..6ccae36fdcd6 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -1757,14 +1757,6 @@ def _accept_draft_tokens(
             batch_indices = torch.arange(batch_size, device=device)
             new_tokens[0, :, 0] = target_tokens[num_accepted_tokens,
                                                 batch_indices]
-
-            # Print speculative acceptance counts each iteration
-            accepted = num_accepted_tokens.detach().cpu().tolist()
-            print("[spec-decode] iter=%s batch=%d ctx=%d draft_len=%d "
-                  "accepted_draft_tokens=%s (draft+1 is the target token)" %
-                  (getattr(self, "iter_counter", None), batch_size,
-                   len(scheduled_batch.context_requests), max_draft_len,
-                   accepted))
         else:
             # No draft tokens to accept, just use the first (and only) sampled token
             batch_indices = torch.arange(batch_size, device=device)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index 0b239f470c0b..c0bb0acf785b 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -1476,25 +1476,12 @@ def process_draft_tokens(
                 num_accepted = self._process_draft_tokens_greedy(
                     request, new_tokens=new_tokens_list, finish_reasons=finish_reasons
                 )
+            return num_accepted
         else:
-            num_accepted = self._process_draft_tokens_rejection_sampling(
+            return self._process_draft_tokens_rejection_sampling(
                 request, new_tokens_list=new_tokens_list, new_tokens_tensor=new_tokens_tensor
             )
 
-        # Print draft acceptance info each iteration (non-overlap scheduler path)
-        draft_len = get_draft_token_length(request)
-        print(
-            "[spec-decode] req=%s iter=%s draft_len=%d accepted=%d"
-            % (
-                getattr(request, "py_request_id", None),
-                getattr(request, "py_decoding_iter", None),
-                draft_len,
-                num_accepted,
-            )
-        )
-
-        return num_accepted
-
     def _get_logprobs_from_request(self, request: LlmRequest) -> tuple[torch.Tensor, torch.Tensor]:
         """Extract the logprobs from the request
 

From 16bfce38dac6ca5200d00ede01cc4c06661db1c8 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 10 Dec 2025 18:36:05 -0800
Subject: [PATCH 03/18] removed some prints and added license

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/shim/ad_executor.py    |  2 +-
 .../transform/library/hidden_states.py        | 25 +++++++++++--------
 2 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index d01e7e1def0b..903fb2baa75c 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -194,7 +194,7 @@ def capture_hidden_states(self, cache_seq_interface: CachedSequenceInterface) ->
         num_tokens = sum(cache_seq_interface.info.seq_len)
 
         hidden_states = [hidden_state[:num_tokens] for hidden_state in full_hidden_states]
-        hidden_states = torch.cat(hidden_states, dim=1) if hidden_states else None
+        hidden_states = torch.cat(hidden_states, dim=1)
         hidden_states = hidden_states.to(dtype=self.dtype)
 
         token_idx = self.hidden_state_write_indices[:num_tokens]
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index b7e698d83a0b..e2e7dd05979a 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
 """The transform passes to capture the hidden states of the target model."""
 
 from typing import Dict, List, Optional, Set, Tuple, Type
@@ -47,13 +62,8 @@ def cached_residual_add(
 ) -> torch.Tensor:
     ret = torch.ops.aten.add(t1, t2)
     b, s, _ = ret.shape
-    print(f"In cached residual add. Ret shape: {ret.shape}")
-    print(f"Shape of hidden_states_cache: {hidden_states_cache.shape}")
     num_tokens = b * s
-    print(f"Num tokens: {num_tokens}")
 
-    # TODO(govind): do some of these correspond to padding tokens when there are varying sequence lengths?
-    # Might need to extract the actual sequence lengths from somewhere to get the appropriate indices to copy.
     hidden_states_cache[:num_tokens].copy_(ret.view(num_tokens, -1), non_blocking=True)
     return ret
 
@@ -115,8 +125,6 @@ class DetectHiddenStatesForCaptureConfig(TransformConfig):
     """Configuration for the hidden states detection transform."""
 
     # TODO: figure out how to get layers to capture.
-    # Right now default is None and EagleSpecMetadata has a heuristic to extract layer indices to capture.
-    # This seems fragile.
     # We should consider if we can use the layer indices stored in eagle checkpoints, e.g.
     # https://huggingface.co/nvidia/gpt-oss-120b-Eagle3/blob/main/config.json#L9-L14
     eagle3_layers_to_capture: Optional[Set[int]] = None  # Default: Do not capture any layers
@@ -145,15 +153,12 @@ def _apply(
 
         def _get_layer_number(lin_node: Node) -> Optional[int]:
             weight = lin_node.args[1]
-            print(f"Calling _get_layer_number() with lin_node: {lin_node}")
             if weight.op == "get_attr":
                 subnames = weight.target.split(".")
                 for subname in subnames:
                     if subname.isdigit():
-                        print(f"Found layer number: {int(subname)}")
                         return int(subname)
 
-            print("No layer number found")
             return None
 
         # find last closing linear node of each layer

From 7c8b8610dd0f0bf3a016f108599819487092fd6b Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 10 Dec 2025 18:47:32 -0800
Subject: [PATCH 04/18] fixing comments from AI

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/llm_args.py            | 10 ++++++++--
 .../defs/examples/test_ad_speculative_decoding.py      |  5 ++++-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py
index ac9440f774e5..50835bff5689 100644
--- a/tensorrt_llm/_torch/auto_deploy/llm_args.py
+++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py
@@ -46,7 +46,7 @@ def _check_for_default_value_only(
 
 
 def default_eagle3_layers_to_capture(num_hidden_layers: int) -> Set[int]:
-    if num_hidden_layers <= 5:
+    if num_hidden_layers <= 6:
         raise ValueError("Not enough hidden layers for default EAGLE3 capture")
     return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
 
@@ -434,7 +434,7 @@ def ensure_no_custom_parallel_config(cls, value: Any, info: ValidationInfo) -> A
         return _check_for_default_value_only(cls, value, info, msg)
 
     @model_validator(mode="after")
-    def default_eagle3_layers_to_capture(self):
+    def set_eagle3_layers_to_capture(self):
         if self.speculative_config is None or not isinstance(
             self.speculative_config, EagleDecodingConfig
         ):
@@ -447,6 +447,12 @@ def default_eagle3_layers_to_capture(self):
             )
 
         # insert the layers to capture into the transforms config.
+        if self.transforms is None:
+            self.transforms = {}
+
+        if "detect_hidden_states_for_capture" not in self.transforms:
+            self.transforms["detect_hidden_states_for_capture"] = {}
+
         self.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] = (
             self.speculative_config.eagle3_layers_to_capture
         )
diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index 2374abec8198..2ce40e88d7b5 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import os
+from typing import Optional
 
 import pytest
 from build_and_run_ad import ExperimentConfig, main
@@ -65,7 +66,9 @@ def make_spec_config(spec_dec_mode: str, spec_model_path: str):
     raise ValueError(f"Unknown speculative mode: {spec_dec_mode}")
 
 
-def run_with_autodeploy(model, speculative_model_dir, batch_size, spec_dec_mode: str | None):
+def run_with_autodeploy(
+    model, speculative_model_dir, batch_size, spec_dec_mode: Optional[str] = None
+):
     """Run AutoDeploy with or without speculative decoding.
 
     Args:

From a6142d9a7275617bbf58d4b96ec4dc50819681d3 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Thu, 11 Dec 2025 11:34:23 -0800
Subject: [PATCH 05/18] revert _compute_logits() args

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/shim/ad_executor.py    | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 903fb2baa75c..879e51e2a2a6 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -776,19 +776,11 @@ def _build_input_ids(request) -> Tuple[List[int], List[int], bool]:
         self.iter_states["num_generation_tokens"] = num_generation_tokens
 
     @nvtx_range("ad_compute_logits")
-    def _compute_logits(self, resource_manager: ResourceManager) -> List[torch.Tensor]:
+    def _compute_logits(self) -> List[torch.Tensor]:
         # run the model
         logits: torch.Tensor = self.model(**self.cache_seq_interface.named_args)[0]
         logits = self.cache_seq_interface.info.maybe_gather_and_squeeze_logits(logits)
 
-        spec_resource_manager = resource_manager.get_resource_manager(
-            ResourceManagerType.SPEC_RESOURCE_MANAGER
-        )
-        if spec_resource_manager is not None and isinstance(
-            spec_resource_manager, ADHiddenStateManager
-        ):
-            spec_resource_manager.capture_hidden_states(self.cache_seq_interface)
-
         # TRTLLMSampler expects float32 logits. PyTorchModelEngine always casts to float32 regardless.
         return logits.float()
 
@@ -816,8 +808,18 @@ def forward(
         self.iter_counter += 1
 
         outputs = {
-            "logits": self._compute_logits(resource_manager),
+            "logits": self._compute_logits(),
         }
+
+        # save hidden states after running model.forward() in _compute_logits()
+        spec_resource_manager = resource_manager.get_resource_manager(
+            ResourceManagerType.SPEC_RESOURCE_MANAGER
+        )
+        if spec_resource_manager is not None and isinstance(
+            spec_resource_manager, ADHiddenStateManager
+        ):
+            spec_resource_manager.capture_hidden_states(self.cache_seq_interface)
+
         if self.mapping is not None:
             self._execute_logit_post_processors(scheduled_requests, outputs)
 

From db5e8ca0234cf51d7592698ad108dcc13b947e81 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Mon, 15 Dec 2025 10:11:40 -0800
Subject: [PATCH 06/18] point integration test to existing eagle3 directory

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tests/integration/defs/examples/test_ad_speculative_decoding.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index 2ce40e88d7b5..0bf8a5cece0b 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -29,7 +29,7 @@
     "What is the highest mountain in the world?",
 ]
 
-EAGLE_MODEL_SUBPATH = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B"
+EAGLE_MODEL_SUBPATH = "EAGLE3-LLaMA3.1-Instruct-8B"
 LLAMA_BASE_SUBPATH = "llama-3.1-model/Llama-3.1-8B-Instruct"
 DRAFT_TARGET_MAX_DRAFT_LEN = 3
 EAGLE_MAX_DRAFT_LEN = 3

From 5dbdfc0fe92ff7912195c3502b16539841315761 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Tue, 16 Dec 2025 17:14:04 -0800
Subject: [PATCH 07/18] fix error with hidden state transform on rebase -
 implement abstract class method

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/transform/library/hidden_states.py    | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index e2e7dd05979a..c6e6d00ac532 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -258,6 +258,11 @@ def get_global_buffer_initializers(cls, source_attn_node: Node) -> BufferInitial
     def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         return []
 
+    @classmethod
+    def get_standard_metadata_args(cls) -> List[str]:
+        # unused, I think?
+        return []
+
 
 @TransformRegistry.register("insert_cached_residual_add")
 class InsertCachedResidualAdd(InsertCachedAttention):

From f09d2dc38de46adf93bfb65d411094116e8769e3 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 17 Dec 2025 10:26:21 -0800
Subject: [PATCH 08/18] remove triton backend workaround from test

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../integration/defs/examples/test_ad_speculative_decoding.py  | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index 0bf8a5cece0b..f972ff12a803 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -101,9 +101,6 @@ def run_with_autodeploy(
         "world_size": 1,
         "kv_cache_config": kv_cache_config,
         "disable_overlap_scheduler": True,
-        "transforms": {
-            "fuse_rmsnorm": {"rmsnorm_backend": "triton"},
-        },
         "max_num_tokens": 64,
     }
 

From bf3d4518d4d6531d1a389a0ee917d82da62c022a Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Fri, 19 Dec 2025 11:58:33 -0800
Subject: [PATCH 09/18] slimming down cached attention op for hidden states

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../transform/library/hidden_states.py        | 62 -------------------
 1 file changed, 62 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index c6e6d00ac532..113544846b27 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -25,12 +25,9 @@
     AttentionDescriptor,
     AttentionLayout,
     AttentionRegistry,
-    BufferInitializerDict,
     CacheConfig,
     CacheInitializerDict,
-    Constant,
     MHACallable,
-    PrepareMetadataCallable,
     SequenceInfo,
 )
 from ...models.factory import ModelFactory
@@ -75,52 +72,6 @@ def cached_residual_add_fake(
     return torch.ops.aten.add(t1, t2)
 
 
-@torch.library.custom_op("auto_deploy::cached_residual_add_prepare_metadata", mutates_args=())
-def cached_residual_add_prepare_metadata(
-    position_ids: torch.Tensor,
-    seq_len: torch.Tensor,
-    input_pos: torch.Tensor,
-    cache_loc: torch.Tensor,
-    pages_per_seq: torch.Tensor,
-    slot_idx: torch.Tensor,
-    page_size: int,
-    chunk_size: int,
-) -> List[torch.Tensor]:
-    return [
-        position_ids,
-        seq_len,
-        input_pos,
-        cache_loc,
-        pages_per_seq,
-        slot_idx,
-        page_size,
-        chunk_size,
-    ]
-
-
-@cached_residual_add_prepare_metadata.register_fake
-def cached_residual_add_prepare_metadata_fake(
-    position_ids: torch.Tensor,
-    seq_len: torch.Tensor,
-    input_pos: torch.Tensor,
-    cache_loc: torch.Tensor,
-    pages_per_seq: torch.Tensor,
-    slot_idx: torch.Tensor,
-    page_size: int,
-    chunk_size: int,
-) -> List[torch.Tensor]:
-    return [
-        position_ids,
-        seq_len,
-        input_pos,
-        cache_loc,
-        pages_per_seq,
-        slot_idx,
-        page_size,
-        chunk_size,
-    ]
-
-
 class DetectHiddenStatesForCaptureConfig(TransformConfig):
     """Configuration for the hidden states detection transform."""
 
@@ -234,10 +185,6 @@ def get_source_attention_op(cls) -> OpOverloadPacket:
     def get_cached_attention_op(cls) -> MHACallable:
         return torch.ops.auto_deploy.cached_residual_add
 
-    @classmethod
-    def get_prepare_metadata_op(cls) -> Tuple[PrepareMetadataCallable, int]:
-        return torch.ops.auto_deploy.cached_residual_add_prepare_metadata, 0
-
     @classmethod
     def get_cache_initializers(
         cls, source_attn_node: Node, cache_config: CacheConfig
@@ -250,17 +197,8 @@ def _get_hidden_states_cache(si: SequenceInfo):
 
         return {"hidden_states_cache": _get_hidden_states_cache}
 
-    @classmethod
-    def get_global_buffer_initializers(cls, source_attn_node: Node) -> BufferInitializerDict:
-        return {}
-
-    @classmethod
-    def get_constants(cls, source_attn_node: Node) -> List[Constant]:
-        return []
-
     @classmethod
     def get_standard_metadata_args(cls) -> List[str]:
-        # unused, I think?
         return []
 
 

From c5c9d9b9c17022a1674d0364d1baddfc8901d30f Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Fri, 19 Dec 2025 18:15:48 -0800
Subject: [PATCH 10/18] more robust sharing of embedding and lm_head weights
 with PyTorchModelEngine drafter

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/shim/ad_executor.py    | 66 ++++++++++++-----
 .../_torch/auto_deploy/utils/_graph.py        | 70 ++++++++++++++++++-
 .../_torch/auto_deploy/utils/node_utils.py    |  9 ++-
 3 files changed, 127 insertions(+), 18 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 879e51e2a2a6..9bdb2217d420 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -22,6 +22,9 @@
 from torch._prims_common import DeviceLikeType
 
 from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures
+from tensorrt_llm._torch.auto_deploy.utils._graph import find_embedding_node, find_lm_head_node
+from tensorrt_llm._torch.auto_deploy.utils.node_utils import get_weight_tensor
+from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM
 from tensorrt_llm._torch.pyexecutor._util import (
     _create_kv_cache_manager,
     get_decoding_mode,
@@ -826,26 +829,56 @@ def forward(
         return outputs
 
 
-def share_embedding_weights(
+def share_target_weights_with_draft(
     target_model_engine: "ADEngine", draft_model_engine: PyTorchModelEngine
 ):
-    # This function is necessary for supporting Eagle and other speculative decoding methods that
-    # copy the embed_tokens submodule. It is not necessary for MTP and other speculative decoding methods that
-    # use the draft model engine directly.
+    """
+    Certain speculative decoding methods (e.g. Eagle3) require sharing the target model's embedding and lm_head weights
+    with the draft model. This function does this sharing if necessary.
+    """
 
-    submodule = target_model_engine.model.model.embed_tokens
+    assert isinstance(draft_model_engine.model, Eagle3ForCausalLM), (
+        f"Expected draft_model_engine.model to be Eagle3ForCausalLM, got {type(draft_model_engine.model)}"
+    )
 
-    world_size = mpi_world_size()
-    assert world_size <= 1, f"This code assumes tp<=1. World size: {world_size}"
+    def share_embedding_weights_with_draft(
+        target_model_engine: "ADEngine", draft_model_engine: PyTorchModelEngine
+    ):
+        gm, embedding_node = find_embedding_node(target_model_engine.model)
+        embedding_weight = get_weight_tensor(gm, embedding_node)
+
+        world_size = mpi_world_size()
+        assert world_size <= 1, f"This code assumes tp<=1. World size: {world_size}"
+
+        # Note: This simple forward function implementation assumes tp=1.
+        # TODO(govind): Handle the tp>1 case.
+        def new_embedding_forward(self, input_ids):
+            return F.embedding(input_ids, self.weight)
 
-    # Note: This simple forward function implementation assumes tp=1.
-    # TODO(govind): Handle the tp>1 case.
-    def new_embedding_forward(self, input_ids):
-        return F.embedding(input_ids, self.weight)
+        if draft_model_engine.model.model.embed_tokens is None:
+            submodule = torch.nn.Module()
+            submodule.forward = MethodType(new_embedding_forward, submodule)
+            submodule.weight = embedding_weight
+            draft_model_engine.model.model.embed_tokens = submodule
 
-    submodule.forward = MethodType(new_embedding_forward, submodule)
+    def share_lm_head_weights_with_draft(
+        target_model_engine: "ADEngine", draft_model_engine: PyTorchModelEngine
+    ):
+        vocab_size = target_model_engine.cache_seq_interface.info.vocab_size_padded
+
+        gm, lm_head_node = find_lm_head_node(target_model_engine.model)
+        lm_head_weight = get_weight_tensor(gm, lm_head_node)
+
+        assert lm_head_weight.shape[0] == vocab_size, (
+            f"Expected lm_head weight first dimension to be vocab_size={vocab_size}, "
+            f"but got shape {lm_head_weight.shape}"
+        )
 
-    draft_model_engine.load_weights_from_target_model(target_model_engine.model)
+        if draft_model_engine.model.load_lm_head_from_target:
+            draft_model_engine.model.lm_head.weight = lm_head_weight
+
+    share_embedding_weights_with_draft(target_model_engine, draft_model_engine)
+    share_lm_head_weights_with_draft(target_model_engine, draft_model_engine)
 
 
 def create_draft_model_engine_maybe(
@@ -900,9 +933,10 @@ def create_draft_model_engine_maybe(
         drafting_loop_wrapper=drafting_loop_wrapper,
     )
 
-    share_embedding_weights(
-        target_model_engine=target_engine, draft_model_engine=draft_model_engine
-    )
+    if draft_spec_config.spec_dec_mode.is_eagle3():
+        share_target_weights_with_draft(
+            target_model_engine=target_engine, draft_model_engine=draft_model_engine
+        )
 
     draft_model_engine.kv_cache_manager_key = ResourceManagerType.DRAFT_KV_CACHE_MANAGER
 
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
index 2c536ec76922..6f97fa3b5b03 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
@@ -1,5 +1,6 @@
 """Graph-related utilities for transformations."""
 
+from collections import deque
 from contextlib import contextmanager
 from typing import Any, Dict, Iterator, Optional, Tuple, Union
 
@@ -17,7 +18,7 @@
 from torch.utils._pytree import _LEAF_SPEC
 
 from .logger import ad_logger
-from .node_utils import is_op
+from .node_utils import is_linear_op, is_op
 
 _NoValType = type("_NoValType", (), {})
 _NO_VAL = _NoValType()
@@ -344,3 +345,70 @@ def _is_meta_tensor(t) -> bool:
                 return True
 
     return False
+
+
+def find_embedding_node(model: nn.Module) -> tuple[GraphModule, Node]:
+    """Find the unique embedding node across all graph modules."""
+    embedding_nodes = []
+    for _, gm in named_graphmodules(model):
+        found_nodes = gm.graph.find_nodes(
+            op="call_function", target=torch.ops.aten.embedding.default
+        )
+        for node in found_nodes:
+            embedding_nodes.append((gm, node))
+
+    assert len(embedding_nodes) == 1, (
+        f"Expected exactly 1 aten.embedding.default node, but found {len(embedding_nodes)}."
+    )
+
+    return embedding_nodes[0]
+
+
+def find_output_node(model: nn.Module) -> tuple[GraphModule, Node]:
+    """Find the unique output node across all graph modules."""
+    output_nodes = []
+    for _, gm in named_graphmodules(model):
+        for node in gm.graph.nodes:
+            if node.op == "output":
+                output_nodes.append((gm, node))
+
+    assert len(output_nodes) == 1, f"Expected exactly 1 output node, but found {len(output_nodes)}."
+
+    return output_nodes[0]
+
+
+def find_lm_head_node(model: nn.Module) -> tuple[GraphModule, Node]:
+    """Find the lm_head node by traversing backwards from the output node."""
+    gm, output_node = find_output_node(model)
+
+    visited = set()
+    queue = deque()
+
+    for arg in output_node.args:
+        if isinstance(arg, Node):
+            queue.append(arg)
+        elif isinstance(arg, (list, tuple)):
+            for item in arg:
+                if isinstance(item, Node):
+                    queue.append(item)
+
+    lm_head_node = None
+    while queue:
+        node = queue.popleft()
+        if node in visited:
+            continue
+        visited.add(node)
+
+        if is_linear_op(node):
+            lm_head_node = node
+            break
+
+        for arg in node.args:
+            if isinstance(arg, Node) and arg not in visited:
+                queue.append(arg)
+
+    assert lm_head_node is not None, (
+        "Could not find lm_head linear op by traversing backwards from output node."
+    )
+
+    return gm, lm_head_node
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index d3446f5caf01..1a84aebf7d15 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -916,12 +916,19 @@ def filter_condition(node: Node, embd: Optional[int] = None, dim: Optional[int]
 def has_shape(node: Node) -> bool:
     return hasattr(node, "meta") and "val" in node.meta and hasattr(node.meta["val"], "shape")
 
-
 def shape(node: Node) -> Tuple[int, ...]:
     if not has_shape(node):
         return None
     return node.meta["val"].shape
 
+def get_weight_tensor(gm: GraphModule, node: Node) -> "torch.Tensor":
+    """Extract the weight tensor from a node within a GraphModule."""
+    weight_name = extract_param_names_from_node(node)[0]
+    weight_tensor = gm
+    for part in weight_name.split("."):
+        weight_tensor = getattr(weight_tensor, part)
+    return weight_tensor
+
 
 def draw_graph(gm: GraphModule, filename: str):
     """

From 589f689498bcba0a32971330404e6e402d9b8d95 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Fri, 19 Dec 2025 19:47:10 -0800
Subject: [PATCH 11/18] counting number of hidden layers inside transform

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/llm_args.py   |  9 +---
 .../transform/library/hidden_states.py        | 50 ++++++++++++++-----
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py
index 50835bff5689..2e16e4a76ef6 100644
--- a/tensorrt_llm/_torch/auto_deploy/llm_args.py
+++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py
@@ -434,18 +434,12 @@ def ensure_no_custom_parallel_config(cls, value: Any, info: ValidationInfo) -> A
         return _check_for_default_value_only(cls, value, info, msg)
 
     @model_validator(mode="after")
-    def set_eagle3_layers_to_capture(self):
+    def setup_hidden_state_capture(self):
         if self.speculative_config is None or not isinstance(
             self.speculative_config, EagleDecodingConfig
         ):
             return self
 
-        if self.speculative_config.eagle3_layers_to_capture is None:
-            num_hidden_layers = self.create_factory()._get_model_config()[0].num_hidden_layers
-            self.speculative_config.eagle3_layers_to_capture = default_eagle3_layers_to_capture(
-                num_hidden_layers
-            )
-
         # insert the layers to capture into the transforms config.
         if self.transforms is None:
             self.transforms = {}
@@ -453,6 +447,7 @@ def set_eagle3_layers_to_capture(self):
         if "detect_hidden_states_for_capture" not in self.transforms:
             self.transforms["detect_hidden_states_for_capture"] = {}
 
+        self.transforms["detect_hidden_states_for_capture"]["capture_hidden_states"] = True
         self.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] = (
             self.speculative_config.eagle3_layers_to_capture
         )
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index 113544846b27..3cfa1219fd3e 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -75,11 +75,20 @@ def cached_residual_add_fake(
 class DetectHiddenStatesForCaptureConfig(TransformConfig):
     """Configuration for the hidden states detection transform."""
 
+    # Whether to capture hidden states at all. If False we will not capture any layers.
+    capture_hidden_states: bool = False
+
     # TODO: figure out how to get layers to capture.
     # We should consider if we can use the layer indices stored in eagle checkpoints, e.g.
     # https://huggingface.co/nvidia/gpt-oss-120b-Eagle3/blob/main/config.json#L9-L14
     eagle3_layers_to_capture: Optional[Set[int]] = None  # Default: Do not capture any layers
 
+    @classmethod
+    def default_eagle3_layers_to_capture(cls, num_hidden_layers: int) -> Set[int]:
+        if num_hidden_layers <= 6:
+            raise ValueError("Not enough hidden layers for default EAGLE3 capture")
+        return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
+
 
 @TransformRegistry.register("detect_hidden_states_for_capture")
 class DetectHiddenStatesForCapture(BaseTransform):
@@ -91,17 +100,7 @@ class DetectHiddenStatesForCapture(BaseTransform):
     def get_config_class(cls) -> Type[TransformConfig]:
         return DetectHiddenStatesForCaptureConfig
 
-    def _apply(
-        self,
-        gm: GraphModule,
-        cm: CachedSequenceInterface,
-        factory: ModelFactory,
-        shared_config: SharedConfig,
-    ) -> Tuple[GraphModule, TransformInfo]:
-        if not self.config.eagle3_layers_to_capture:
-            info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
-            return gm, info
-
+    def collect_residual_add_nodes(self, gm: GraphModule) -> Dict[int, Node]:
         def _get_layer_number(lin_node: Node) -> Optional[int]:
             weight = lin_node.args[1]
             if weight.op == "get_attr":
@@ -119,7 +118,7 @@ def _get_layer_number(lin_node: Node) -> Optional[int]:
         for _, _, lin_node_closing in layer_subgraphs:
             # need layer number to correctly identify the residual add node
             layer_number = _get_layer_number(lin_node_closing)
-            if layer_number is None or layer_number not in self.config.eagle3_layers_to_capture:
+            if layer_number is None:
                 continue
 
             # Conditions to identify as the hidden states after the residual
@@ -140,6 +139,33 @@ def _get_layer_number(lin_node: Node) -> Optional[int]:
                 # this stores the last residual add node encountered for each layer
                 residual_add_nodes[layer_number] = res_node
 
+        return residual_add_nodes
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        if not self.config.capture_hidden_states:
+            info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
+            return gm, info
+
+        residual_add_nodes = self.collect_residual_add_nodes(gm)
+
+        if self.config.eagle3_layers_to_capture is None:
+            num_hidden_layers = len(residual_add_nodes)
+            self.config.eagle3_layers_to_capture = (
+                DetectHiddenStatesForCaptureConfig.default_eagle3_layers_to_capture(
+                    num_hidden_layers
+                )
+            )
+
+        residual_add_nodes = {
+            k: v for k, v in residual_add_nodes.items() if k in self.config.eagle3_layers_to_capture
+        }
+
         assert residual_add_nodes.keys() == self.config.eagle3_layers_to_capture, (
             f"Unable to find residual add nodes for layers. Expected: {self.config.eagle3_layers_to_capture}, \
             Found: {residual_add_nodes.keys()}"

From 54db13e64a8f132b09b0cfdb9cb9b4c67bd277f5 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Fri, 19 Dec 2025 19:56:56 -0800
Subject: [PATCH 12/18] set default layers to capture instead of classmethod

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../transform/library/hidden_states.py          | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index 3cfa1219fd3e..4382087957ea 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -81,13 +81,16 @@ class DetectHiddenStatesForCaptureConfig(TransformConfig):
     # TODO: figure out how to get layers to capture.
     # We should consider if we can use the layer indices stored in eagle checkpoints, e.g.
     # https://huggingface.co/nvidia/gpt-oss-120b-Eagle3/blob/main/config.json#L9-L14
-    eagle3_layers_to_capture: Optional[Set[int]] = None  # Default: Do not capture any layers
+    eagle3_layers_to_capture: Optional[Set[int]] = None
 
-    @classmethod
-    def default_eagle3_layers_to_capture(cls, num_hidden_layers: int) -> Set[int]:
+    def set_default_eagle3_layers_to_capture(self, num_hidden_layers: int):
+        """
+        Used to set default layers to capture when we want to capture hidden states, but
+        no layers to capture are provided.
+        """
         if num_hidden_layers <= 6:
             raise ValueError("Not enough hidden layers for default EAGLE3 capture")
-        return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
+        self.eagle3_layers_to_capture = {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
 
 
 @TransformRegistry.register("detect_hidden_states_for_capture")
@@ -156,11 +159,7 @@ def _apply(
 
         if self.config.eagle3_layers_to_capture is None:
             num_hidden_layers = len(residual_add_nodes)
-            self.config.eagle3_layers_to_capture = (
-                DetectHiddenStatesForCaptureConfig.default_eagle3_layers_to_capture(
-                    num_hidden_layers
-                )
-            )
+            self.config.set_default_eagle3_layers_to_capture(num_hidden_layers)
 
         residual_add_nodes = {
             k: v for k, v in residual_add_nodes.items() if k in self.config.eagle3_layers_to_capture

From 35e7c1e85d70a22a647c481f6d0c6cdb86218706 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Mon, 22 Dec 2025 12:18:17 -0800
Subject: [PATCH 13/18] adding review comments addressing again

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/llm_args.py   | 16 +--------
 .../_torch/auto_deploy/shim/ad_executor.py    |  5 ++-
 .../_torch/auto_deploy/utils/_graph.py        | 33 +++++++++++++++----
 .../_torch/auto_deploy/utils/node_utils.py    |  5 +--
 4 files changed, 30 insertions(+), 29 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/llm_args.py b/tensorrt_llm/_torch/auto_deploy/llm_args.py
index 2e16e4a76ef6..aa9f0147cd20 100644
--- a/tensorrt_llm/_torch/auto_deploy/llm_args.py
+++ b/tensorrt_llm/_torch/auto_deploy/llm_args.py
@@ -1,6 +1,6 @@
 from importlib.resources import files
 from pathlib import Path
-from typing import Any, Dict, List, Literal, Optional, Set, Type, Union
+from typing import Any, Dict, List, Literal, Optional, Type, Union
 
 import torch
 from pydantic import Field, PrivateAttr, ValidationInfo, field_validator, model_validator
@@ -45,12 +45,6 @@ def _check_for_default_value_only(
     return value
 
 
-def default_eagle3_layers_to_capture(num_hidden_layers: int) -> Set[int]:
-    if num_hidden_layers <= 6:
-        raise ValueError("Not enough hidden layers for default EAGLE3 capture")
-    return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
-
-
 _TRANSFORMS_SHORTCUT_LOOKUP = {
     "attn_backend": ("insert_cached_attention.backend", "transformers_replace_cached_attn.backend"),
     "free_mem_ratio": ("resize_kv_cache.free_mem_ratio",),
@@ -440,18 +434,10 @@ def setup_hidden_state_capture(self):
         ):
             return self
 
-        # insert the layers to capture into the transforms config.
-        if self.transforms is None:
-            self.transforms = {}
-
-        if "detect_hidden_states_for_capture" not in self.transforms:
-            self.transforms["detect_hidden_states_for_capture"] = {}
-
         self.transforms["detect_hidden_states_for_capture"]["capture_hidden_states"] = True
         self.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] = (
             self.speculative_config.eagle3_layers_to_capture
         )
-
         return self
 
     @model_validator(mode="after")
diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 9bdb2217d420..3e1055b70787 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -22,7 +22,7 @@
 from torch._prims_common import DeviceLikeType
 
 from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures
-from tensorrt_llm._torch.auto_deploy.utils._graph import find_embedding_node, find_lm_head_node
+from tensorrt_llm._torch.auto_deploy.utils._graph import find_lm_head_node, get_input_embeddings
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import get_weight_tensor
 from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM
 from tensorrt_llm._torch.pyexecutor._util import (
@@ -844,8 +844,7 @@ def share_target_weights_with_draft(
     def share_embedding_weights_with_draft(
         target_model_engine: "ADEngine", draft_model_engine: PyTorchModelEngine
     ):
-        gm, embedding_node = find_embedding_node(target_model_engine.model)
-        embedding_weight = get_weight_tensor(gm, embedding_node)
+        embedding_weight = get_input_embeddings(target_model_engine.model)
 
         world_size = mpi_world_size()
         assert world_size <= 1, f"This code assumes tp<=1. World size: {world_size}"
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
index 6f97fa3b5b03..22cf8562b423 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
@@ -18,7 +18,7 @@
 from torch.utils._pytree import _LEAF_SPEC
 
 from .logger import ad_logger
-from .node_utils import is_linear_op, is_op
+from .node_utils import get_weight_tensor, is_linear_op, is_op
 
 _NoValType = type("_NoValType", (), {})
 _NO_VAL = _NoValType()
@@ -347,21 +347,35 @@ def _is_meta_tensor(t) -> bool:
     return False
 
 
-def find_embedding_node(model: nn.Module) -> tuple[GraphModule, Node]:
+def get_input_embeddings(model: nn.Module) -> torch.Tensor:
     """Find the unique embedding node across all graph modules."""
-    embedding_nodes = []
+    embedding_weights = []
     for _, gm in named_graphmodules(model):
         found_nodes = gm.graph.find_nodes(
             op="call_function", target=torch.ops.aten.embedding.default
         )
         for node in found_nodes:
-            embedding_nodes.append((gm, node))
+            embedding_weights.append(get_weight_tensor(gm, node))
 
-    assert len(embedding_nodes) == 1, (
-        f"Expected exactly 1 aten.embedding.default node, but found {len(embedding_nodes)}."
+    if hasattr(model, "get_input_embeddings"):
+        embedding_weights.append(model.get_input_embeddings())
+
+    for _, gm in named_graphmodules(model):
+        if hasattr(gm, "get_input_embeddings"):
+            embedding_weights.append(gm.get_input_embeddings())
+
+    assert len(embedding_weights) > 0, "No embedding weights found"
+    unique_embedding_weights = [embedding_weights[0]]
+    for weight in embedding_weights:
+        if weight is not unique_embedding_weights[0]:
+            unique_embedding_weights.append(weight)
+
+    assert len(unique_embedding_weights) == 1, (
+        f"Expected exactly 1 unique embedding weight, but found {len(unique_embedding_weights)}."
     )
 
-    return embedding_nodes[0]
+    print(f"Unique embedding weights: {unique_embedding_weights}")
+    return unique_embedding_weights[0]
 
 
 def find_output_node(model: nn.Module) -> tuple[GraphModule, Node]:
@@ -379,8 +393,13 @@ def find_output_node(model: nn.Module) -> tuple[GraphModule, Node]:
 
 def find_lm_head_node(model: nn.Module) -> tuple[GraphModule, Node]:
     """Find the lm_head node by traversing backwards from the output node."""
+    print(
+        f"Finding lm_head node in model: {model}"
+    )  # Want to see if model is already a graph module.
     gm, output_node = find_output_node(model)
 
+    print(f"Output node: {output_node}")
+
     visited = set()
     queue = deque()
 
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index 1a84aebf7d15..3cf102f6ce57 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -924,10 +924,7 @@ def shape(node: Node) -> Tuple[int, ...]:
 def get_weight_tensor(gm: GraphModule, node: Node) -> "torch.Tensor":
     """Extract the weight tensor from a node within a GraphModule."""
     weight_name = extract_param_names_from_node(node)[0]
-    weight_tensor = gm
-    for part in weight_name.split("."):
-        weight_tensor = getattr(weight_tensor, part)
-    return weight_tensor
+    return gm.get_parameter(weight_name)
 
 
 def draw_graph(gm: GraphModule, filename: str):

From 6cc353ff9a18bcf24c7f294a0c3f1ab390cefd9d Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Mon, 22 Dec 2025 14:35:49 -0800
Subject: [PATCH 14/18] finished review comments

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../_torch/auto_deploy/shim/ad_executor.py    |  6 +-
 .../library/gather_logits_before_lm_head.py   |  8 +--
 .../transform/library/hidden_states.py        |  6 ++
 .../_torch/auto_deploy/utils/_graph.py        | 58 +++++--------------
 4 files changed, 26 insertions(+), 52 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
index 3e1055b70787..cf66991e9f4f 100644
--- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
+++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py
@@ -22,8 +22,7 @@
 from torch._prims_common import DeviceLikeType
 
 from tensorrt_llm._torch.attention_backend.interface import AttentionRuntimeFeatures
-from tensorrt_llm._torch.auto_deploy.utils._graph import find_lm_head_node, get_input_embeddings
-from tensorrt_llm._torch.auto_deploy.utils.node_utils import get_weight_tensor
+from tensorrt_llm._torch.auto_deploy.utils._graph import get_input_embeddings, get_lm_head_weights
 from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM
 from tensorrt_llm._torch.pyexecutor._util import (
     _create_kv_cache_manager,
@@ -865,8 +864,7 @@ def share_lm_head_weights_with_draft(
     ):
         vocab_size = target_model_engine.cache_seq_interface.info.vocab_size_padded
 
-        gm, lm_head_node = find_lm_head_node(target_model_engine.model)
-        lm_head_weight = get_weight_tensor(gm, lm_head_node)
+        lm_head_weight = get_lm_head_weights(target_model_engine.model)
 
         assert lm_head_weight.shape[0] == vocab_size, (
             f"Expected lm_head weight first dimension to be vocab_size={vocab_size}, "
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py b/tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py
index 3f2520971fcb..9af1d9aa1684 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/gather_logits_before_lm_head.py
@@ -25,7 +25,9 @@
 import torch
 from torch.fx import GraphModule
 
-from ...utils.node_utils import is_linear_op, is_op
+from tensorrt_llm._torch.auto_deploy.utils._graph import get_lm_head_node
+
+from ...utils.node_utils import is_linear_op
 from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
 
 
@@ -54,9 +56,7 @@ def _apply(
         self._log_info("Applying GatherLogitsBeforeLmHead transform...")
 
         # assume lm head node is the input to the output node
-        lm_head_node = gm.graph.find_nodes(op="output")[0].all_input_nodes[0]
-        if is_op(lm_head_node, torch.ops.aten.to):
-            lm_head_node = lm_head_node.all_input_nodes[0]
+        lm_head_node = get_lm_head_node(gm)
 
         if is_linear_op(lm_head_node):
             node_to_gather = lm_head_node.all_input_nodes[0]
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index 4382087957ea..74b455f96fc7 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -155,6 +155,12 @@ def _apply(
             info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
             return gm, info
 
+        if gm.graph.find_nodes(
+            op="call_function", target=torch.ops.auto_deploy.residual_add_for_capture.default
+        ):
+            info = TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
+            return gm, info
+
         residual_add_nodes = self.collect_residual_add_nodes(gm)
 
         if self.config.eagle3_layers_to_capture is None:
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
index 22cf8562b423..cd61bd52f1e7 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py
@@ -1,6 +1,5 @@
 """Graph-related utilities for transformations."""
 
-from collections import deque
 from contextlib import contextmanager
 from typing import Any, Dict, Iterator, Optional, Tuple, Union
 
@@ -18,7 +17,7 @@
 from torch.utils._pytree import _LEAF_SPEC
 
 from .logger import ad_logger
-from .node_utils import get_weight_tensor, is_linear_op, is_op
+from .node_utils import get_weight_tensor, is_op
 
 _NoValType = type("_NoValType", (), {})
 _NO_VAL = _NoValType()
@@ -374,60 +373,31 @@ def get_input_embeddings(model: nn.Module) -> torch.Tensor:
         f"Expected exactly 1 unique embedding weight, but found {len(unique_embedding_weights)}."
     )
 
-    print(f"Unique embedding weights: {unique_embedding_weights}")
     return unique_embedding_weights[0]
 
 
-def find_output_node(model: nn.Module) -> tuple[GraphModule, Node]:
+def get_output_node(model: nn.Module) -> tuple[GraphModule, Node]:
     """Find the unique output node across all graph modules."""
     output_nodes = []
     for _, gm in named_graphmodules(model):
-        for node in gm.graph.nodes:
-            if node.op == "output":
-                output_nodes.append((gm, node))
+        output_nodes.extend([(gm, node) for node in gm.graph.find_nodes(op="output")])
 
     assert len(output_nodes) == 1, f"Expected exactly 1 output node, but found {len(output_nodes)}."
-
     return output_nodes[0]
 
 
-def find_lm_head_node(model: nn.Module) -> tuple[GraphModule, Node]:
-    """Find the lm_head node by traversing backwards from the output node."""
-    print(
-        f"Finding lm_head node in model: {model}"
-    )  # Want to see if model is already a graph module.
-    gm, output_node = find_output_node(model)
-
-    print(f"Output node: {output_node}")
-
-    visited = set()
-    queue = deque()
+def get_lm_head_node(gm: GraphModule, output_node: Optional[Node] = None) -> Node:
+    if output_node is None:
+        output_node = gm.graph.find_nodes(op="output")[0]
 
-    for arg in output_node.args:
-        if isinstance(arg, Node):
-            queue.append(arg)
-        elif isinstance(arg, (list, tuple)):
-            for item in arg:
-                if isinstance(item, Node):
-                    queue.append(item)
+    lm_head_node = output_node.all_input_nodes[0]
+    if is_op(lm_head_node, torch.ops.aten.to):
+        lm_head_node = lm_head_node.all_input_nodes[0]
 
-    lm_head_node = None
-    while queue:
-        node = queue.popleft()
-        if node in visited:
-            continue
-        visited.add(node)
+    return lm_head_node
 
-        if is_linear_op(node):
-            lm_head_node = node
-            break
-
-        for arg in node.args:
-            if isinstance(arg, Node) and arg not in visited:
-                queue.append(arg)
-
-    assert lm_head_node is not None, (
-        "Could not find lm_head linear op by traversing backwards from output node."
-    )
 
-    return gm, lm_head_node
+def get_lm_head_weights(model: nn.Module) -> torch.Tensor:
+    gm, output_node = get_output_node(model)
+    lm_head_node = get_lm_head_node(gm, output_node)
+    return get_weight_tensor(gm, lm_head_node)

From 53d40c2560b79afe62cac3c4055218967de70926 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Tue, 23 Dec 2025 11:54:05 -0800
Subject: [PATCH 15/18] restoring original spec dec tests for draft-target.
 Adding test for acceptance rates for Eagle3 + Llama

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../examples/test_ad_speculative_decoding.py  | 124 ++++++++++++++----
 .../test_lists/test-db/l0_h100.yml            |   5 +-
 2 files changed, 100 insertions(+), 29 deletions(-)

diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index f972ff12a803..fb0961e142ec 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -14,12 +14,13 @@
 # limitations under the License.
 
 import os
-from typing import Optional
 
 import pytest
 from build_and_run_ad import ExperimentConfig, main
 from defs.conftest import llm_models_root
 
+from tensorrt_llm import SamplingParams
+from tensorrt_llm._torch.auto_deploy.llm import LLM
 from tensorrt_llm.llmapi import DraftTargetDecodingConfig, EagleDecodingConfig, KvCacheConfig
 
 prompts = [
@@ -51,31 +52,19 @@ def get_model_paths():
     return base_model, draft_target_model, eagle_model
 
 
-def make_spec_config(spec_dec_mode: str, spec_model_path: str):
-    if spec_dec_mode == "draft_target":
-        return DraftTargetDecodingConfig(
-            max_draft_len=DRAFT_TARGET_MAX_DRAFT_LEN, speculative_model_dir=spec_model_path
-        )
-    if spec_dec_mode == "eagle":
-        return EagleDecodingConfig(
-            max_draft_len=EAGLE_MAX_DRAFT_LEN,
-            speculative_model_dir=spec_model_path,
-            eagle3_one_model=False,
-            eagle3_layers_to_capture=None,
-        )
-    raise ValueError(f"Unknown speculative mode: {spec_dec_mode}")
+def make_draft_target_config(spec_model_path: str):
+    return DraftTargetDecodingConfig(
+        max_draft_len=DRAFT_TARGET_MAX_DRAFT_LEN, speculative_model_dir=spec_model_path
+    )
 
 
-def run_with_autodeploy(
-    model, speculative_model_dir, batch_size, spec_dec_mode: Optional[str] = None
-):
+def run_with_autodeploy(model, speculative_model_dir, batch_size):
     """Run AutoDeploy with or without speculative decoding.
 
     Args:
         model: Path to the base model
         speculative_model_dir: Path to the speculative model (None for baseline mode)
         batch_size: Number of prompts to process
-        spec_dec_mode: Speculative decoding mode
 
     Returns:
         List of (prompt, output) tuples from prompts_and_outputs
@@ -85,12 +74,12 @@ def run_with_autodeploy(
 
     # Configure speculative decoding if speculative_model_dir is provided
     spec_config = None
-    if speculative_model_dir is not None and spec_dec_mode is not None:
-        spec_config = make_spec_config(spec_dec_mode, speculative_model_dir)
+    if speculative_model_dir:
+        spec_config = make_draft_target_config(speculative_model_dir)
 
     # Configure KV cache
     kv_cache_config = KvCacheConfig(
-        free_gpu_memory_fraction=0.1,
+        free_gpu_memory_fraction=0.01,
     )
 
     # Configure AutoDeploy LLM arguments
@@ -137,8 +126,8 @@ def run_with_autodeploy(
     return result["prompts_and_outputs"]
 
 
-@pytest.mark.parametrize("batch_size, spec_dec_mode", [(1, "draft_target"), (4, "eagle")])
-def test_autodeploy_spec_dec(batch_size, spec_dec_mode):
+@pytest.mark.parametrize("batch_size", [1, 4])
+def test_autodeploy_spec_dec(batch_size):
     """Test AutoDeploy speculative decoding with different batch sizes.
 
     Runs with and without speculative decoding and verifies outputs are identical.
@@ -147,10 +136,10 @@ def test_autodeploy_spec_dec(batch_size, spec_dec_mode):
     print(f"Testing AutoDeploy Speculative Decoding - Batch Size {batch_size}")
     print("=" * 80)
 
-    base_model, draft_target_model, eagle_model = get_model_paths()
+    base_model, draft_target_model, _ = get_model_paths()
 
     print(f"\nBase Model: {base_model}")
-    spec_model_path = draft_target_model if spec_dec_mode == "draft_target" else eagle_model
+    spec_model_path = draft_target_model
     print(f"Speculative Model: {spec_model_path}")
     print(f"Batch Size: {batch_size}")
 
@@ -160,14 +149,13 @@ def test_autodeploy_spec_dec(batch_size, spec_dec_mode):
         model=base_model,
         speculative_model_dir=spec_model_path,
         batch_size=batch_size,
-        spec_dec_mode=spec_dec_mode,
     )
     print(f"Generated {len(spec_outputs)} outputs with speculative decoding")
 
     # Run without speculative decoding (baseline)
     print("\n[2/2] Running without speculative decoding (baseline)...")
     baseline_outputs = run_with_autodeploy(
-        model=base_model, speculative_model_dir=None, batch_size=batch_size, spec_dec_mode=None
+        model=base_model, speculative_model_dir=None, batch_size=batch_size
     )
     print(f"Generated {len(baseline_outputs)} outputs in baseline mode")
 
@@ -196,3 +184,85 @@ def test_autodeploy_spec_dec(batch_size, spec_dec_mode):
     print("\n" + "=" * 80)
     print("SUCCESS! All outputs are identical between spec-dec and baseline modes")
     print("=" * 80)
+
+
+def test_autodeploy_eagle3_acceptance_rate():
+    """Test Eagle3 acceptance rate with AutoDeploy engine.
+
+    Runs Eagle3 speculative decoding with streaming and verifies
+    that the acceptance rate is above a minimum threshold.
+    """
+    print("\n" + "=" * 80)
+    print("Testing AutoDeploy Eagle3 Acceptance Rate - Batch Size 1")
+    print("=" * 80)
+
+    base_model, _, eagle_model = get_model_paths()
+
+    print(f"\nBase Model: {base_model}")
+    print(f"Eagle3 Model: {eagle_model}")
+
+    max_draft_len = EAGLE_MAX_DRAFT_LEN
+
+    # Configure Eagle3 speculative decoding
+    speculative_config = EagleDecodingConfig(
+        max_draft_len=max_draft_len,
+        speculative_model_dir=eagle_model,
+        eagle3_one_model=False,
+        eagle3_layers_to_capture=None,
+    )
+
+    # Configure KV cache
+    kv_cache_config = KvCacheConfig(
+        free_gpu_memory_fraction=0.01,
+    )
+
+    # Create AutoDeploy LLM with Eagle3 speculative decoding
+    # We directly instantiate the LLM class instead of using the main() function
+    # so that we can stream the outputs to see acceptance rates without needing to
+    # collect them in the executor.
+    llm = LLM(
+        model=base_model,
+        skip_loading_weights=False,
+        runtime="trtllm",
+        world_size=1,
+        kv_cache_config=kv_cache_config,
+        speculative_config=speculative_config,
+        disable_overlap_scheduler=True,
+        max_num_tokens=64,
+    )
+
+    # Use a single prompt for batch size 1
+    prompt = prompts[0]
+    tok_ids = llm.tokenizer.encode(prompt)
+
+    sampling_params = SamplingParams(max_tokens=128, temperature=0, seed=42)
+
+    print("\nRunning Eagle3 speculative decoding with streaming...")
+    print(f"Prompt: {prompt}")
+
+    num_tokens = 0
+    num_drafted = 0
+    num_accepted = 0
+
+    for output in llm.generate_async(tok_ids, sampling_params, streaming=True):
+        new_tokens = output.outputs[0].token_ids
+        num_drafted += max_draft_len
+        num_accepted += len(new_tokens) - num_tokens - 1
+        num_tokens = len(new_tokens)
+
+    accept_rate = num_accepted / num_drafted
+
+    print("\nAcceptance Rate Statistics:")
+    print(f"  Total tokens drafted: {num_drafted}")
+    print(f"  Total tokens accepted: {num_accepted}")
+    print(f"  Acceptance rate: {accept_rate:.2%}")
+
+    # Verify acceptance rate is above minimum threshold (10%)
+    min_acceptance_rate = 0.10
+    assert accept_rate > min_acceptance_rate, (
+        f"Acceptance rate {accept_rate:.2%} is below minimum threshold {min_acceptance_rate:.0%}"
+    )
+
+    print("\n" + "=" * 80)
+    print(f"SUCCESS! Acceptance rate {accept_rate:.2%} > {min_acceptance_rate:.0%} threshold")
+    print("=" * 80)
diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml
index 55c5b60d4ea7..40878c5fe90f 100644
--- a/tests/integration/test_lists/test-db/l0_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_h100.yml
@@ -116,8 +116,9 @@ l0_h100:
   - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[True]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_bf16
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[1-draft_target]
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[4-eagle]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[1]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[4]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate
 - condition:
     ranges:
       system_gpu_count:

From 34db281bce6f08b8cd99c9e0b6b0f6d4f81f2a22 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Tue, 23 Dec 2025 12:00:53 -0800
Subject: [PATCH 16/18] some small cleanups to the integration test after
 reverting

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../defs/examples/test_ad_speculative_decoding.py    | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index fb0961e142ec..00d896126a88 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -86,6 +86,7 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     llm_args = {
         "model": model,
         "skip_loading_weights": False,
+        "speculative_config": spec_config,
         "runtime": "trtllm",
         "world_size": 1,
         "kv_cache_config": kv_cache_config,
@@ -106,10 +107,6 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     # Create ExperimentConfig
     cfg = ExperimentConfig(**experiment_config)
 
-    cfg.args.speculative_config = (
-        spec_config  # Add here to avoid Pydantic validation error for eagle3_layers_to_capture
-    )
-
     # Add sampling parameters (deterministic with temperature=0.0 and fixed seed)
     cfg.prompt.sp_kwargs = {
         "max_tokens": 50,
@@ -133,21 +130,20 @@ def test_autodeploy_spec_dec(batch_size):
     Runs with and without speculative decoding and verifies outputs are identical.
     """
     print("\n" + "=" * 80)
-    print(f"Testing AutoDeploy Speculative Decoding - Batch Size {batch_size}")
+    print(f"Testing AutoDeploy Speculative Decoding (Draft Target) - Batch Size {batch_size}")
     print("=" * 80)
 
     base_model, draft_target_model, _ = get_model_paths()
 
     print(f"\nBase Model: {base_model}")
-    spec_model_path = draft_target_model
-    print(f"Speculative Model: {spec_model_path}")
+    print(f"Speculative Model: {draft_target_model}")
     print(f"Batch Size: {batch_size}")
 
     # Run with speculative decoding
     print("\n[1/2] Running with speculative decoding enabled...")
     spec_outputs = run_with_autodeploy(
         model=base_model,
-        speculative_model_dir=spec_model_path,
+        speculative_model_dir=draft_target_model,
         batch_size=batch_size,
     )
     print(f"Generated {len(spec_outputs)} outputs with speculative decoding")

From d8a1cba1f563de19aa64254bf340f8c05a616474 Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Tue, 23 Dec 2025 15:09:10 -0800
Subject: [PATCH 17/18] fix precommit

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/utils/node_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index 3cf102f6ce57..286650c77fda 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -916,11 +916,13 @@ def filter_condition(node: Node, embd: Optional[int] = None, dim: Optional[int]
 def has_shape(node: Node) -> bool:
     return hasattr(node, "meta") and "val" in node.meta and hasattr(node.meta["val"], "shape")
 
+
 def shape(node: Node) -> Tuple[int, ...]:
     if not has_shape(node):
         return None
     return node.meta["val"].shape
 
+
 def get_weight_tensor(gm: GraphModule, node: Node) -> "torch.Tensor":
     """Extract the weight tensor from a node within a GraphModule."""
     weight_name = extract_param_names_from_node(node)[0]

From 131a58821f51974855545f8a5e60ab520cd37c8b Mon Sep 17 00:00:00 2001
From: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
Date: Wed, 24 Dec 2025 15:18:35 -0800
Subject: [PATCH 18/18] reverting batch_size > 1 test from before; it seems to
 fail due to IFB nondeterminism. Need to figure out how to meaningfully test
 batch_size > 1 in a deterministic way

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>
---
 .../transform/library/hidden_states.py        |   3 +-
 .../examples/test_ad_speculative_decoding.py  | 110 ++++++++++--------
 .../test_lists/test-db/l0_h100.yml            |   4 +-
 3 files changed, 66 insertions(+), 51 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
index 74b455f96fc7..e1c917160c72 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/hidden_states.py
@@ -118,7 +118,8 @@ def _get_layer_number(lin_node: Node) -> Optional[int]:
         # from there we will find the residual add node for that layer
         layer_subgraphs, unprocessed_linear_nodes = get_all_layer_subgraphs(gm)
         residual_add_nodes: Dict[int, Node] = {}
-        for _, _, lin_node_closing in layer_subgraphs:
+        for layer_subgraph in layer_subgraphs:
+            lin_node_closing = layer_subgraph.terminating_node
             # need layer number to correctly identify the residual add node
             layer_number = _get_layer_number(lin_node_closing)
             if layer_number is None:
diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py
index 00d896126a88..ddc785841e6e 100644
--- a/tests/integration/defs/examples/test_ad_speculative_decoding.py
+++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py
@@ -26,8 +26,6 @@
 prompts = [
     "What is the capital of France?",
     "Please explain the concept of gravity in simple words and a single sentence.",
-    "What is the capital of Norway?",
-    "What is the highest mountain in the world?",
 ]
 
 EAGLE_MODEL_SUBPATH = "EAGLE3-LLaMA3.1-Instruct-8B"
@@ -58,12 +56,21 @@ def make_draft_target_config(spec_model_path: str):
     )
 
 
-def run_with_autodeploy(model, speculative_model_dir, batch_size):
+def make_eagle3_config(spec_model_path: str):
+    return EagleDecodingConfig(
+        max_draft_len=EAGLE_MAX_DRAFT_LEN,
+        speculative_model_dir=spec_model_path,
+        eagle3_one_model=False,
+        eagle3_layers_to_capture=None,
+    )
+
+
+def run_with_autodeploy(model, speculative_config, batch_size):
     """Run AutoDeploy with or without speculative decoding.
 
     Args:
         model: Path to the base model
-        speculative_model_dir: Path to the speculative model (None for baseline mode)
+        speculative_config: Speculative decoding config (None for baseline mode)
         batch_size: Number of prompts to process
 
     Returns:
@@ -72,10 +79,7 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     # Select prompts based on batch size
     selected_prompts = prompts[:batch_size]
 
-    # Configure speculative decoding if speculative_model_dir is provided
-    spec_config = None
-    if speculative_model_dir:
-        spec_config = make_draft_target_config(speculative_model_dir)
+    spec_config = speculative_config
 
     # Configure KV cache
     kv_cache_config = KvCacheConfig(
@@ -123,36 +127,46 @@ def run_with_autodeploy(model, speculative_model_dir, batch_size):
     return result["prompts_and_outputs"]
 
 
-@pytest.mark.parametrize("batch_size", [1, 4])
-def test_autodeploy_spec_dec(batch_size):
-    """Test AutoDeploy speculative decoding with different batch sizes.
+# Note: This test tests exact equality of outputs between speculative and baseline modes.
+# This can fail for larger batch sizes due to nondeterminism with in flight batching.
+# TODO: Figure out a robust test for output correctness that can pass for larger batch sizes.
+@pytest.mark.parametrize("spec_dec_mode", ["draft_target", "eagle3"])
+def test_autodeploy_spec_dec_output(spec_dec_mode):
+    """Test AutoDeploy speculative decoding output correctness.
 
     Runs with and without speculative decoding and verifies outputs are identical.
     """
     print("\n" + "=" * 80)
-    print(f"Testing AutoDeploy Speculative Decoding (Draft Target) - Batch Size {batch_size}")
+    print(f"Testing AutoDeploy Speculative Decoding ({spec_dec_mode}) - Output Correctness")
     print("=" * 80)
 
-    base_model, draft_target_model, _ = get_model_paths()
+    base_model, draft_target_model, eagle_model = get_model_paths()
+
+    # Select model and config based on mode
+    if spec_dec_mode == "draft_target":
+        spec_model = draft_target_model
+        spec_config = make_draft_target_config(spec_model)
+    elif spec_dec_mode == "eagle3":  # eagle3
+        spec_model = eagle_model
+        spec_config = make_eagle3_config(spec_model)
+    else:
+        raise ValueError(f"Unsupported speculative decoding mode: {spec_dec_mode}")
 
     print(f"\nBase Model: {base_model}")
-    print(f"Speculative Model: {draft_target_model}")
-    print(f"Batch Size: {batch_size}")
+    print(f"Speculative Model ({spec_dec_mode}): {spec_model}")
 
     # Run with speculative decoding
     print("\n[1/2] Running with speculative decoding enabled...")
     spec_outputs = run_with_autodeploy(
         model=base_model,
-        speculative_model_dir=draft_target_model,
-        batch_size=batch_size,
+        speculative_config=spec_config,
+        batch_size=1,
     )
     print(f"Generated {len(spec_outputs)} outputs with speculative decoding")
 
     # Run without speculative decoding (baseline)
     print("\n[2/2] Running without speculative decoding (baseline)...")
-    baseline_outputs = run_with_autodeploy(
-        model=base_model, speculative_model_dir=None, batch_size=batch_size
-    )
+    baseline_outputs = run_with_autodeploy(model=base_model, speculative_config=None, batch_size=1)
     print(f"Generated {len(baseline_outputs)} outputs in baseline mode")
 
     # Verify outputs are identical
@@ -189,7 +203,7 @@ def test_autodeploy_eagle3_acceptance_rate():
     that the acceptance rate is above a minimum threshold.
     """
     print("\n" + "=" * 80)
-    print("Testing AutoDeploy Eagle3 Acceptance Rate - Batch Size 1")
+    print("Testing AutoDeploy Eagle3 Acceptance Rate")
     print("=" * 80)
 
     base_model, _, eagle_model = get_model_paths()
@@ -227,38 +241,38 @@ def test_autodeploy_eagle3_acceptance_rate():
         max_num_tokens=64,
     )
 
-    # Use a single prompt for batch size 1
-    prompt = prompts[0]
-    tok_ids = llm.tokenizer.encode(prompt)
+    # Tokenize 2 prompts to test multiple sequential requests
+    batch_tok_ids = [llm.tokenizer.encode(p) for p in prompts[:2]]
 
     sampling_params = SamplingParams(max_tokens=128, temperature=0, seed=42)
 
     print("\nRunning Eagle3 speculative decoding with streaming...")
-    print(f"Prompt: {prompt}")
 
-    num_tokens = 0
-    num_drafted = 0
-    num_accepted = 0
-
-    for output in llm.generate_async(tok_ids, sampling_params, streaming=True):
-        new_tokens = output.outputs[0].token_ids
-        num_drafted += max_draft_len
-        num_accepted += len(new_tokens) - num_tokens - 1
-        num_tokens = len(new_tokens)
-
-    accept_rate = num_accepted / num_drafted
-
-    print("\nAcceptance Rate Statistics:")
-    print(f"  Total tokens drafted: {num_drafted}")
-    print(f"  Total tokens accepted: {num_accepted}")
-    print(f"  Acceptance rate: {accept_rate:.2%}")
-
-    # Verify acceptance rate is above minimum threshold (10%)
-    min_acceptance_rate = 0.10
-    assert accept_rate > min_acceptance_rate, (
-        f"Acceptance rate {accept_rate:.2%} is below minimum threshold {min_acceptance_rate:.0%}"
-    )
+    # Process each request sequentially and verify acceptance rate
+    for i in range(len(batch_tok_ids)):
+        num_tokens = 0
+        num_drafted = 0
+        num_accepted = 0
+
+        for output in llm.generate_async(batch_tok_ids[i], sampling_params, streaming=True):
+            new_tokens = output.outputs[0].token_ids
+            num_drafted += max_draft_len
+            num_accepted += len(new_tokens) - num_tokens - 1
+            num_tokens = len(new_tokens)
+
+        accept_rate = num_accepted / num_drafted
+
+        print(f"\nRequest {i + 1} Acceptance Rate Statistics:")
+        print(f"  Total tokens drafted: {num_drafted}")
+        print(f"  Total tokens accepted: {num_accepted}")
+        print(f"  Acceptance rate: {accept_rate:.2%}")
+
+        # Verify acceptance rate is above minimum threshold (10%)
+        min_acceptance_rate = 0.10
+        assert accept_rate > min_acceptance_rate, (
+            f"Request {i + 1}: Acceptance rate {accept_rate:.2%} is below minimum threshold {min_acceptance_rate:.0%}"
+        )
 
     print("\n" + "=" * 80)
-    print(f"SUCCESS! Acceptance rate {accept_rate:.2%} > {min_acceptance_rate:.0%} threshold")
+    print("SUCCESS! All requests passed acceptance rate threshold")
     print("=" * 80)
diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml
index 40878c5fe90f..c4ed2f088641 100644
--- a/tests/integration/test_lists/test-db/l0_h100.yml
+++ b/tests/integration/test_lists/test-db/l0_h100.yml
@@ -116,8 +116,8 @@ l0_h100:
   - accuracy/test_llm_api_autodeploy.py::TestNemotronH::test_auto_dtype[True]
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_fp8
   - accuracy/test_llm_api_autodeploy.py::TestNemotronMOE::test_bf16
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[1]
-  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec[4]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec_output[draft_target]
+  - examples/test_ad_speculative_decoding.py::test_autodeploy_spec_dec_output[eagle3]
   - examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_acceptance_rate
 - condition:
     ranges: