made ADHiddenStatesManager class which extracts hidden size and dtype

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

Author: govind-ramnarayan

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -30,7 +30,7 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import get_draft_token_length
 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_resource_manager, get_spec_drafter
+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 (
@@ -111,6 +111,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:
@@ -360,48 +444,6 @@ def __init__(
         # start fresh with fixed seed
         torch.manual_seed(42)
 
-    def _prepare_hidden_state_capture(
-        self, ordered_requests: RequestList, resource_manager: ResourceManager
-    ) -> None:
-        spec_resource_manager = resource_manager.get_resource_manager(
-            ResourceManagerType.SPEC_RESOURCE_MANAGER
-        )
-        if spec_resource_manager is None or not isinstance(
-            spec_resource_manager, Eagle3ResourceManager
-        ):
-            return
-
-        caches = []
-        for name, tensor in self.cache_seq_interface.named_args.items():
-            if "hidden_states_cache" in name:
-                caches.append((name, tensor))
-
-        seq_lens = self.cache_seq_interface.info.seq_len
-        num_tokens = sum(seq_lens)
-        max_total_draft_tokens = getattr(spec_resource_manager, "max_total_draft_tokens", 0)
-
-        start_idx = 0
-        hidden_state_write_indices = []
-        for request, seq_len in zip(ordered_requests, seq_lens):
-            request_id = request.request_id
-            slot_id = spec_resource_manager.slot_manager.get_slot(request_id)
-            spec_resource_manager.start_indices[slot_id] = start_idx
-            hidden_state_write_indices.extend(range(start_idx, start_idx + seq_len))
-            start_idx += max(seq_len, max_total_draft_tokens + 1)
-            assert start_idx < spec_resource_manager.hidden_states.shape[0], (
-                f"start_idx {start_idx} exceeds hidden_states capacity {spec_resource_manager.hidden_states.shape[0]}"
-            )
-
-        assert len(hidden_state_write_indices) == num_tokens
-
-        self.hidden_state_write_indices_host = torch.tensor(
-            hidden_state_write_indices, dtype=torch.long
-        )
-
-        self.hidden_state_write_indices_gpu[:num_tokens].copy_(
-            self.hidden_state_write_indices_host, non_blocking=True
-        )
-
     @nvtx_range("ad_prepare_inputs")
     def _prepare_inputs(
         self,
@@ -414,7 +456,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
         extend_requests = [
@@ -425,7 +470,6 @@ def _prepare_inputs(
         ]
         gen_requests = extend_requests + generation_requests
         ordered_requests = context_requests + gen_requests
-
         # info to be extracted
         input_ids: List[List[int]] = []
         input_pos: List[int] = []
@@ -566,58 +610,38 @@ def _build_input_ids(request) -> Tuple[List[int], List[int]]:
                 scatter_ref=dummy_token,
             )
 
+        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
 
-        self._prepare_hidden_state_capture(ordered_requests, resource_manager)
-
         return last_logit_only
 
     @nvtx_range("ad_compute_logits")
     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]
-        self._capture_hidden_states_cache(resource_manager)
-
-        # TRTLLMSampler expects float32 logits. PyTorchModelEngine always casts to float32 regardless.
-        logits = logits.float()
-
-        # return a list of tensors
-        return self.cache_seq_interface.info.unnest_sequences(logits)
 
-    def _capture_hidden_states_cache(self, resource_manager: ResourceManager) -> None:
-        """Capture and print hidden_states_cache tensor passed to the model."""
         spec_resource_manager = resource_manager.get_resource_manager(
             ResourceManagerType.SPEC_RESOURCE_MANAGER
         )
-        if spec_resource_manager is None or not isinstance(
-            spec_resource_manager, Eagle3ResourceManager
+        if spec_resource_manager is not None and isinstance(
+            spec_resource_manager, ADHiddenStateManager
         ):
-            return
+            spec_resource_manager.capture_hidden_states(self.cache_seq_interface)
 
-        caches = []
-        for name, tensor in self.cache_seq_interface.named_args.items():
-            if "hidden_states_cache" in name:
-                caches.append((name, tensor))
-
-        if not caches:
-            return
-
-        seq_lens = self.cache_seq_interface.info.seq_len
-        num_tokens = sum(seq_lens)
-
-        used_caches = [cache[:num_tokens] for _, cache in caches]
-
-        eagle3_hidden_states = spec_resource_manager.hidden_states
-        hidden_states_cache_value = torch.cat(used_caches, dim=1) if used_caches else None
-        hidden_states_cache_value = hidden_states_cache_value.to(dtype=eagle3_hidden_states.dtype)
+        # TRTLLMSampler expects float32 logits. PyTorchModelEngine always casts to float32 regardless.
+        logits = logits.float()
 
-        token_idx = self.hidden_state_write_indices_gpu[:num_tokens]
-        eagle3_hidden_states[:, : hidden_states_cache_value.shape[1]].index_copy_(
-            0, token_idx, hidden_states_cache_value
-        )
+        # return a list of tensors
+        return self.cache_seq_interface.info.unnest_sequences(logits)
 
     def get_max_num_sequences(self) -> int:
         """Maximum number of sequences supported by the engine."""
@@ -837,15 +861,16 @@ def create_autodeploy_executor(ad_config: LlmArgs, tokenizer: Optional[Tokenizer
         ad_config=ad_config, target_engine=engine, dist_mapping=dist_mapping, mpi_dist=mpi_dist
     )
 
-    target_model_dtype = torch.bfloat16  # TODO: Get this from the model engine.
-    target_hidden_size = 4096  # TODO: Get this from the model engine.
-
-    spec_resource_manager = _get_spec_resource_manager(
-        target_model_engine=engine,
-        max_seq_len=engine.llm_args.max_seq_len,
-        model_dtype=target_model_dtype,
-        hidden_size=target_hidden_size,
-        draft_model_engine=draft_model_engine,
+    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

tensorrt_llm/_torch/speculative/__init__.py

@@ -5,10 +5,10 @@
 from .ngram import NGramDrafter, NGramPoolManager
 from .save_hidden_state import SaveHiddenStatesDrafter
 from .spec_tree_manager import SpecTreeManager
-from .utils import (_get_spec_resource_manager, get_num_extra_kv_tokens,
-                    get_num_spec_layers, get_spec_decoder, get_spec_drafter,
-                    get_spec_metadata, get_spec_resource_manager,
-                    get_spec_worker, update_spec_config_from_model_config)
+from .utils import (get_num_extra_kv_tokens, get_num_spec_layers,
+                    get_spec_decoder, get_spec_drafter, get_spec_metadata,
+                    get_spec_resource_manager, get_spec_worker,
+                    update_spec_config_from_model_config)
 
 __all__ = [
     "Eagle3SpecMetadata",
@@ -25,7 +25,6 @@
     "get_spec_drafter",
     "get_spec_metadata",
     "get_spec_resource_manager",
-    "_get_spec_resource_manager",
     "get_spec_worker",
     "update_spec_config_from_model_config",
     "suggest_spec_config",

tensorrt_llm/_torch/speculative/utils.py

@@ -108,50 +108,47 @@ def get_spec_metadata(spec_config,
     return None
 
 
-def _get_spec_resource_manager(
-        target_model_engine: "ModelEngine",
-        max_seq_len: int,
-        model_dtype: torch.dtype,
-        hidden_size: int,
-        draft_model_engine: Optional["PyTorchModelEngine"] = None):
-    spec_config = target_model_engine.spec_config
+def get_spec_resource_manager(model_engine, draft_model_engine=None):
+    spec_config = model_engine.spec_config
     if spec_config is None:
         return None
-    max_num_requests = target_model_engine.batch_size
-    max_num_tokens = target_model_engine.llm_args.max_num_tokens
+    model_config = model_engine.model.config
+    max_num_requests = model_engine.batch_size
+    max_seq_len = model_engine.max_seq_len
+    max_num_tokens = model_engine.max_num_tokens
     spec_dec_mode = spec_config.spec_dec_mode
     if spec_dec_mode.is_mtp_eagle_one_model():
         if spec_config.use_relaxed_acceptance_for_thinking:
             return MTPHiddenStatesManager(
                 spec_config,
-                model_dtype,
-                hidden_size,
+                model_config.torch_dtype,
+                model_config.hidden_size,
                 max_num_requests,
             )
         else:
             return None
     if spec_dec_mode.is_mtp_one_model():
         return MTPHiddenStatesManager(
             spec_config,
-            model_dtype,
-            hidden_size,
+            model_config.torch_dtype,
+            model_config.hidden_size,
             max_num_requests,
         )
     if spec_dec_mode.is_eagle3() or spec_dec_mode.is_mtp_eagle():
         assert draft_model_engine is not None, "Draft model engine is required for Eagle3 and MTP Eagle two model flow."
         return Eagle3ResourceManager(
             spec_config,
             draft_model_engine.model.config.torch_dtype,
-            hidden_size,
+            model_config.hidden_size,
             max_num_requests,
             max_seq_len,
             max_num_tokens,
         )
     if spec_dec_mode.is_save_hidden_states():
         return Eagle3ResourceManager(
             spec_config,
-            model_dtype,
-            hidden_size,
+            model_engine.model.config.torch_dtype,
+            model_config.hidden_size,
             max_num_requests,
             max_seq_len,
             max_num_tokens,
@@ -163,17 +160,6 @@ def _get_spec_resource_manager(
     return None
 
 
-def get_spec_resource_manager(
-        model_engine: "PyTorchModelEngine",
-        draft_model_engine: Optional["PyTorchModelEngine"] = None):
-    return _get_spec_resource_manager(
-        target_model_engine=model_engine,
-        max_seq_len=model_engine.max_seq_len,
-        model_dtype=model_engine.model.config.torch_dtype,
-        hidden_size=model_engine.model.config.hidden_size,
-        draft_model_engine=draft_model_engine)
-
-
 def get_spec_decoder(sampler_args: TorchSampler.Args,
                      spec_config: "DecodingBaseConfig"):
     if spec_config.spec_dec_mode.is_mtp_one_model():