revert _compute_logits() args

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

Author: govind-ramnarayan

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -588,18 +588,10 @@ def _build_input_ids(request) -> Tuple[List[int], List[int]]:
         return last_logit_only
 
     @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]
 
-        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.
         logits = logits.float()
 
@@ -626,7 +618,16 @@ def forward(
         self.iter_counter += 1
 
         # compute all logits
-        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)
 
         # gather+cat logits
         logits_flat = torch.cat(