diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..65f05ff1912b 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2ec62d5d3cae..73b6152c503d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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(