[https://nvbugs/6337051][fix] Refresh inputs on graph capture to avoid unintended modification

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -378,7 +378,8 @@ def capture(self,
                 forward_fn: Callable,
                 initial_inputs: Dict[str, Any],
                 enable_spec_decode: bool = False,
-                postprocess_fn: Optional[Callable] = None):
+                postprocess_fn: Optional[Callable] = None,
+                prepare_inputs_fn: Optional[Callable] = None):
         """Captures the forward pass for a given batch size."""
         batch_size = key[0]
         # [CUDA graph spec decode padding]
@@ -413,6 +414,15 @@ def capture(self,
             "spec_metadata": initial_inputs.get("spec_metadata", None),
         }
 
+        def _refresh_capture_inputs():
+            # A warmup forward can mutate attention/speculative metadata
+            # Rebuild it through the same path used before a normal replay
+            if prepare_inputs_fn is None:
+                return
+            refreshed_inputs, _ = prepare_inputs_fn()
+            capture_inputs.update(refreshed_inputs)
+            capture_inputs.update(sliced_static_tensors)
+
         def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable,
                                              capture_inputs: Dict[str, Any]):
             is_first_draft = key[2]
@@ -428,17 +438,19 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable,
         # This also lets us initialize states in the attn_metadata.
         graph = torch.cuda.CUDAGraph()
         with with_multi_stream(True), piecewise_cuda_graph(False):
-            for _ in range(self.WARMUP_STEPS):
+            for warmup_iter in range(self.WARMUP_STEPS):
+                if warmup_iter != 0:
+                    _refresh_capture_inputs()
                 _setup_spec_decoding_and_forward(key, forward_fn,
                                                  capture_inputs)
                 if postprocess_fn is not None:
                     postprocess_fn(capture_inputs)
 
+            _refresh_capture_inputs()
+            # Capture has no kernel execution, so no post-process or refresh.
             with torch.cuda.graph(graph, pool=self.memory_pool):
                 output = _setup_spec_decoding_and_forward(
                     key, forward_fn, capture_inputs)
-            if postprocess_fn is not None:
-                postprocess_fn(capture_inputs)
 
         self.graphs[key] = graph
         self.graph_outputs[key] = make_weak_ref(output)

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -5469,22 +5469,28 @@ def forward(self,
                 else:
                     spec_metadata = None
 
-            # Fill slot-ID buffer for scatter inside draft loop
-            if (self.enable_spec_decode and spec_tree_manager is not None
-                    and spec_tree_manager.use_dynamic_tree
-                    and not self.is_draft_model):
-                spec_tree_manager.slot_storage.fill_all_slot_ids(
-                    padded_requests.context_requests,
-                    padded_requests.generation_requests,
-                )
+            def prepare_forward_inputs_fn(
+            ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]:
+                # Fill slot-ID buffer for scatter inside draft loop
+                if (self.enable_spec_decode and spec_tree_manager is not None
+                        and spec_tree_manager.use_dynamic_tree
+                        and not self.is_draft_model):
+                    spec_tree_manager.slot_storage.fill_all_slot_ids(
+                        padded_requests.context_requests,
+                        padded_requests.generation_requests,
+                    )
 
-            inputs, gather_ids = self._prepare_inputs(
-                padded_requests, kv_cache_manager, attn_metadata, spec_metadata,
-                new_tensors_device, cache_indirection_buffer,
-                num_accepted_tokens_device, req_id_to_old_request,
-                resource_manager, can_run_graph)
-            self._prepare_inputs_event = torch.cuda.Event()
-            self._prepare_inputs_event.record()
+                inputs, gather_ids = self._prepare_inputs(
+                    padded_requests, kv_cache_manager, attn_metadata,
+                    spec_metadata, new_tensors_device, cache_indirection_buffer,
+                    num_accepted_tokens_device, req_id_to_old_request,
+                    resource_manager, can_run_graph)
+
+                self._prepare_inputs_event = torch.cuda.Event()
+                self._prepare_inputs_event.record()
+                return inputs, gather_ids
+
+            inputs, gather_ids = prepare_forward_inputs_fn()
 
             with with_shared_pool(self.cuda_graph_runner.get_graph_pool()):
                 if not can_run_graph:
@@ -5512,7 +5518,8 @@ def capture_postprocess_fn(inputs: Dict[str, Any]):
                             capture_forward_fn,
                             inputs,
                             enable_spec_decode=self.enable_spec_decode,
-                            postprocess_fn=capture_postprocess_fn)
+                            postprocess_fn=capture_postprocess_fn,
+                            prepare_inputs_fn=prepare_forward_inputs_fn)
 
                         # Pre-replay: set DSA slot mappings for current batch's draft cache (fixes 2nd warmup)
                         saved_draft = prepare_attn_metadata_for_draft_replay(