[https://nvbugs/6221483][fix] AutoDeploy: Fix Eagle metadata host syncs

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -295,7 +295,7 @@ def copy_to_device(self) -> None:
                     trunc_h_buf = self._trunc_host_bufs[name][:copy_bytes]
                     trunc_d_buf.copy_(trunc_h_buf, non_blocking=True)
 
-    def copy_to_host(self) -> None:
+    def copy_to_host(self, non_blocking: bool = False) -> None:
         """Copy from device buffer to host buffer.
 
         Mirrors ``copy_to_device``: uses the current length of the truncatable tensor
@@ -306,7 +306,7 @@ def copy_to_host(self) -> None:
             if self._total_bytes > 0:
                 h_buffer = self._host_buffer[: self._total_bytes]
                 d_buffer = self._device_buffer[: self._total_bytes]
-                h_buffer.copy_(d_buffer, non_blocking=True)
+                h_buffer.copy_(d_buffer, non_blocking=non_blocking)
 
             # Copy each truncatable tensor independently, truncated to current length
             for name in self._truncatable_names:
@@ -316,7 +316,7 @@ def copy_to_host(self) -> None:
                     copy_bytes = length * dtype.itemsize
                     trunc_d_buf = self._trunc_device_bufs[name][:copy_bytes]
                     trunc_h_buf = self._trunc_host_bufs[name][:copy_bytes]
-                    trunc_h_buf.copy_(trunc_d_buf, non_blocking=True)
+                    trunc_h_buf.copy_(trunc_d_buf, non_blocking=non_blocking)
 
     def resize(self, name: str, new_capacity: int) -> None:
         """Resize a truncatable tensor's capacity.
@@ -1175,6 +1175,27 @@ def _is_required(self, name: str, check_both: bool = True) -> bool:
         """
         return self._is_active(name, check_both) or self._is_active_host_prep(name, check_both)
 
+    def _active_host_update_args(
+        self, arg_names: Set[str], active_args_override: Optional[Set[str]] = None
+    ) -> List[str]:
+        """Return host args that need mirroring after an in-graph metadata update.
+
+        ``active_args_override`` lets a caller narrow host mirroring to the graph inputs the next
+        consumer actually reads. It is treated as a filter: only active host args whose names appear
+        in the override are mirrored. The override may contain names that are not active graph args
+        (e.g. a submodule's full placeholder set, which also includes inter-module tensors such as
+        ``inputs_embeds``/``hidden_states``); such entries are simply ignored. The caller is
+        responsible for including every host argument the next consumer may read.
+        """
+        needs_d2h_sync = [
+            k + self._host_suffix
+            for k in arg_names
+            if self._is_active(k + self._host_suffix, check_both=False)
+        ]
+        if active_args_override is None:
+            return needs_d2h_sync
+        return [arg_name for arg_name in needs_d2h_sync if arg_name in active_args_override]
+
     def _stage_arg(
         self,
         name: str,
@@ -1583,11 +1604,16 @@ def run_host_prepare_for_attention_forward(self) -> None:
             host_function(**{arg: self.get_arg(arg) for arg in args})
 
     @nvtx_range("ad_offset_pos_and_cache_")
-    def offset_pos_and_cache_(self, offset: torch.Tensor) -> None:
+    def offset_pos_and_cache_(
+        self, offset: torch.Tensor, active_args_override: Optional[Set[str]] = None
+    ) -> None:
         """Offset position and cache-related metadata for active arguments.
 
         Args:
             offset: 1D tensor [batch_size] with per-sequence position offsets.
+            active_args_override: Optional graph-input names for the next in-forward consumer. When
+                provided, host mirroring is limited to those active host args. The caller is
+                responsible for including every host argument the next consumer may read.
         """
         # check if we need a d2h sync
         _REQUIRES_UPDATE = {
@@ -1599,11 +1625,7 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None:
             "seq_len_with_cache",
             "use_initial_states",
         }
-        needs_d2h_sync = [
-            k + self._host_suffix
-            for k in _REQUIRES_UPDATE
-            if self._is_active(k + self._host_suffix, check_both=False)
-        ]
+        needs_d2h_sync = self._active_host_update_args(_REQUIRES_UPDATE, active_args_override)
         sync_to_host = any(needs_d2h_sync)
         if sync_to_host:
             ad_logger.debug(f"d2h sync required in offset_pos_and_cache_ for {needs_d2h_sync}")
@@ -1694,7 +1716,7 @@ def offset_pos_and_cache_(self, offset: torch.Tensor) -> None:
         # TODO: May need to dissect what fields are needed in the forward pass to reduce
         # data movement.
         if sync_to_host:
-            self._input_buffer.copy_to_host()
+            self._input_buffer.copy_to_host(non_blocking=False)
 
     @nvtx_range("ad_offset_with_new_lens_")
     def offset_with_new_lens_(self, new_lens_ungathered: torch.Tensor) -> None:
@@ -1718,13 +1740,18 @@ def offset_with_new_lens_(self, new_lens_ungathered: torch.Tensor) -> None:
         self.offset_pos_and_cache_(increment)
 
     @nvtx_range("ad_switch_to_generate_")
-    def switch_to_generate_(self) -> None:
+    def switch_to_generate_(self, active_args_override: Optional[Set[str]] = None) -> None:
         """Switch all sequences metadata to generate (decode) mode.
 
         Transitions the batch from any layout (prefill/extend/decode or mixed) to
         an all-decode layout where each sequence has exactly 1 token. We assume that we just take
         the last position of each sequence for the metadata.
 
+        Args:
+            active_args_override: Optional graph-input names for the next in-forward consumer. When
+                provided, host mirroring is limited to those active host args. The caller is
+                responsible for including every host argument the next consumer may read.
+
         NOTE: update device tensors first and mirror back to host only when an updated host-side
         argument is active.
 
@@ -1757,11 +1784,7 @@ def switch_to_generate_(self) -> None:
             "position_ids",
             "use_initial_states",
         }
-        needs_d2h_sync = [
-            k + self._host_suffix
-            for k in _REQUIRES_UPDATE
-            if self._is_active(k + self._host_suffix, check_both=False)
-        ]
+        needs_d2h_sync = self._active_host_update_args(_REQUIRES_UPDATE, active_args_override)
         sync_to_host = any(needs_d2h_sync)
 
         # --- input_ids (device) ---
@@ -1790,7 +1813,7 @@ def switch_to_generate_(self) -> None:
         # TODO: May need to dissect what fields are needed in the forward pass to reduce
         # data movement.
         if sync_to_host:
-            self._input_buffer.copy_to_host()
+            self._input_buffer.copy_to_host(non_blocking=False)
 
     def copy_(self, name: str, src: torch.Tensor, strict: bool = True) -> None:
         """Copy a tensor into the buffer. USE WITH CAUTION!

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -412,18 +412,16 @@ def cap_max_batch_size_to_max_num_tokens(self):
         return self
 
     @model_validator(mode="after")
-    def disable_cudagraph_for_speculative_flashinfer(self):
+    def reject_cudagraph_for_speculative_flashinfer(self):
         if (
             self.speculative_config is not None
             and self.attn_backend == "flashinfer"
             and self.is_cuda_graph_enabled()
         ):
-            ad_logger.warning(
+            raise ValueError(
                 "Speculative decoding with FlashInfer attention does not currently support CUDA "
-                "graph replay in AutoDeploy; falling back to compile_backend='torch-simple'."
+                "graph replay in AutoDeploy. Use compile_backend='torch-simple' instead."
             )
-            self.compile_backend = "torch-simple"
-            self.update_transforms_with_shortcuts()
         return self
 
     ### UTILITY METHODS ############################################################################

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_eagle.py

@@ -33,7 +33,7 @@
 
 from dataclasses import dataclass
 from types import SimpleNamespace
-from typing import Any, ClassVar, Dict, Optional, Union
+from typing import Any, ClassVar, Dict, Optional, Set, Union
 
 import torch
 import torch.nn as nn
@@ -912,10 +912,14 @@ def _forward_prefill_only(self, input_ids: torch.Tensor, position_ids: torch.Ten
     #  KV-cache forward (inference after graph transforms)               #
     # ================================================================== #
 
+    @staticmethod
+    def _submodule_placeholder_names(submodule: nn.Module) -> Set[str]:
+        return {node.name for node in submodule.graph.nodes if node.op == "placeholder"}
+
     @staticmethod
     def _filter_kwargs_for_submodule(kwargs: dict, submodule: nn.Module) -> dict:
         """Filter kwargs to only include those accepted by submodule's forward (GraphModule)."""
-        expected_names = {node.name for node in submodule.graph.nodes if node.op == "placeholder"}
+        expected_names = EagleWrapper._submodule_placeholder_names(submodule)
         return {k: v for k, v in kwargs.items() if k in expected_names}
 
     @staticmethod
@@ -1096,6 +1100,7 @@ def _forward_with_kv_cache(self, csi: CachedSequenceInterface):
         next_new_tokens[:, 0] = csi.info.maybe_gather_and_squeeze(csi.get_arg("input_ids"))
 
         # ---- Phase 5: Draft loop ----
+        draft_arg_names = self._submodule_placeholder_names(self.draft_model)
         for draft_idx in range(self.max_draft_len):
             # run forward pass on the draft model in shape [num_sequences, 1]
             draft_output = self.draft_model(
@@ -1123,9 +1128,9 @@ def _forward_with_kv_cache(self, csi: CachedSequenceInterface):
             # switch to generate (if not done already), store new tokens, and offset cache
             # can be skipped for last iteration since after we return metadata will be reset
             if draft_idx < self.max_draft_len - 1:
-                csi.info.switch_to_generate_()
+                csi.info.switch_to_generate_(active_args_override=draft_arg_names)
                 csi.info.copy_("input_ids", draft_tokens)
-                csi.info.offset_pos_and_cache_(c_offset)
+                csi.info.offset_pos_and_cache_(c_offset, active_args_override=draft_arg_names)
 
         # ---- Phase 6: Package output ----
         return EagleWrapperOutput(