default eagle3 num layers and insertion into transforms now happens in llm_args.py

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Author: govind-ramnarayan

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -1,14 +1,21 @@
 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
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 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.")
@@ -401,6 +414,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.

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -263,33 +263,6 @@ def construct_draft_llm_args(
     return draft_llm_args
 
 
-def _populate_hidden_state_capture_config(ad_config: LlmArgs) -> LlmArgs:
-    """Ensure hidden-state capture transforms follow Eagle config.
-
-    When users supply an EagleDecodingConfig, automatically wire the
-    DetectHiddenStatesForCapture transform so it captures the layers
-    requested by the Eagle config, and make sure the corresponding
-    cached residual add insertion is enabled.
-    """
-
-    spec_config = ad_config.speculative_config
-    if not isinstance(spec_config, EagleDecodingConfig):
-        return ad_config
-
-    spec_dec_mode = spec_config.spec_dec_mode
-    if not (spec_dec_mode.is_eagle3() or spec_dec_mode.is_eagle3_one_model()):
-        return ad_config
-
-    # DetectHiddenStatesForCapture configuration
-    capture_key = "detect_hidden_states_for_capture"
-    capture_cfg = ad_config.transforms.get(capture_key, {"stage": "pattern_matcher"})
-    if spec_config.eagle3_layers_to_capture is not None:
-        capture_cfg["eagle3_layers_to_capture"] = spec_config.eagle3_layers_to_capture
-    ad_config.transforms[capture_key] = capture_cfg
-
-    return ad_config
-
-
 def create_draft_kv_cache_manager_maybe(
     draft_model_engine: Optional[PyTorchModelEngine],
     ad_config: LlmArgs,
@@ -348,9 +321,6 @@ def build_from_config(cls, ad_config: LlmArgs):
 
         factory = ad_config.create_factory()
 
-        # Auto-configure hidden-state capture when using Eagle speculative decoding.
-        ad_config = _populate_hidden_state_capture_config(ad_config)
-
         # initialize seq info object
         seq_info = SequenceInfo(
             max_seq_len=max_seq_len,

tests/integration/defs/examples/test_ad_speculative_decoding.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import os
-from typing import Set
 
 import pytest
 from build_and_run_ad import ExperimentConfig, main
@@ -52,27 +51,17 @@ def get_model_paths():
     return base_model, draft_target_model, eagle_model
 
 
-def get_default_layers_to_capture(num_layers: int) -> Set[int]:
-    if num_layers <= 5:
-        raise ValueError("Not enough hidden layers for default EAGLE3 capture")
-    return {1, num_layers // 2 - 1, num_layers - 4}
-
-
 def make_spec_config(mode: str, spec_model_path: str):
     if mode == "draft_target":
         return DraftTargetDecodingConfig(
             max_draft_len=DRAFT_TARGET_MAX_DRAFT_LEN, speculative_model_dir=spec_model_path
         )
     if mode == "eagle":
-        num_layers = 32  # hardcoded num layers for Llama-3.1-8B-Instruct
-
-        # Note: Valid layer indices must be provided in EagleDecodingConfig when using AutoDeploy.
-        # TODO: Should we have some defaulting behavior in the Engine, similar to PyTorch backend?
         return EagleDecodingConfig(
             max_draft_len=EAGLE_MAX_DRAFT_LEN,
             speculative_model_dir=spec_model_path,
             eagle3_one_model=False,
-            eagle3_layers_to_capture=get_default_layers_to_capture(num_layers),
+            eagle3_layers_to_capture=None,
         )
     raise ValueError(f"Unknown speculative mode: {mode}")