diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..a254eccaf9d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -100,7 +100,7 @@ class CUDAGraphRunner: and low-level execution (capturing, resource management, replaying) for multiple graphs, keyed by (batch size, draft_len, is_first_draft). """ - WARMUP_STEPS = 2 + WARMUP_STEPS = 1 def __init__(self, config: CUDAGraphRunnerConfig): self.config = config @@ -132,6 +132,7 @@ def __init__(self, config: CUDAGraphRunnerConfig): # tensor reallocation from invalidating addresses baked into existing # CUDA graphs. Use allow_capture() context manager during warmup. self._capture_allowed = False + self.is_warmup_only = False def _create_shared_static_tensors(self): """Allocates static tensors sized for the largest possible batch.""" @@ -321,7 +322,7 @@ def maybe_get_cuda_graph( key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata) - if key in self.graphs: + if key in self.graph_metadata: return self.graph_metadata[key][ "attn_metadata"], self.graph_metadata[key]["spec_metadata"], key @@ -378,8 +379,8 @@ def capture(self, forward_fn: Callable, initial_inputs: Dict[str, Any], enable_spec_decode: bool = False, - postprocess_fn: Optional[Callable] = None): - """Captures the forward pass for a given batch size.""" + postprocess_fn: Optional[Callable] = None) -> Any: + """Warm up and/or capture the forward pass for a graph key.""" batch_size = key[0] # [CUDA graph spec decode padding] # We pad input IDs/position IDs to the maximum draft length (token per request). @@ -422,18 +423,23 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs['attn_metadata'].use_spec_decoding = True return forward_fn(capture_inputs) - # We have to do warm up runs to initialize PyTorch's - # internal states according to the docs: - # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics - # This also lets us initialize states in the attn_metadata. - graph = torch.cuda.CUDAGraph() + output = None with with_multi_stream(True), piecewise_cuda_graph(False): + # We have to do a warmup run to initialize PyTorch's internal + # states according to the docs: + # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics + # This also lets us initialize states in the attn_metadata and + # resize the shared attention workspace before any graph is captured. for _ in range(self.WARMUP_STEPS): - _setup_spec_decoding_and_forward(key, forward_fn, - capture_inputs) + output = _setup_spec_decoding_and_forward( + key, forward_fn, capture_inputs) if postprocess_fn is not None: postprocess_fn(capture_inputs) + if self.is_warmup_only: + return output + + graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, pool=self.memory_pool): output = _setup_spec_decoding_and_forward( key, forward_fn, capture_inputs) @@ -441,8 +447,10 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, postprocess_fn(capture_inputs) self.graphs[key] = graph - self.graph_outputs[key] = make_weak_ref(output) + graph_output = make_weak_ref(output) + self.graph_outputs[key] = graph_output self.memory_pool = graph.pool() + return graph_output def replay(self, key: KeyType, current_inputs: Dict[str, Any]) -> Optional[torch.Tensor]: @@ -704,7 +712,7 @@ class EncoderCUDAGraphRunner: Restricted to `TrtllmAttentionMetadata` — FlashInfer's per-batch planner state is not compatible with CUDA graph capture/replay. """ - WARMUP_STEPS = 2 + WARMUP_STEPS = 1 def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.config = config @@ -730,6 +738,7 @@ def __init__(self, config: EncoderCUDAGraphRunnerConfig): self.cuda_graph_meta_buffers = get_memory_buffers() self._capture_allowed = False + self.is_warmup_only = False # CUDA graph H2D memcpy nodes require pinned host sources. In CC mode # prefer_pinned() is false: pageable host buffers are preferred, so the @@ -916,7 +925,7 @@ def maybe_get_cuda_graph( or not is_padding_successful: return None, None - if key in self.graphs: + if key in self.graph_metadata: return self.graph_metadata[key]["attn_metadata"], key # New key not yet captured. Only create metadata if capture is @@ -980,8 +989,8 @@ def capture( key: EncoderKeyType, forward_fn: Callable[[Dict[str, Any]], Any], inputs: Dict[str, Any], - ) -> None: - """Capture a CUDA graph for the given key.""" + ) -> Any: + """Warm up and/or capture the forward pass for a graph key.""" _, padded_num_tokens, _ = key sliced_static_tensors = { @@ -1003,17 +1012,21 @@ def capture( attn_md = capture_inputs["attn_metadata"] - self.graph_metadata[key] = { - "attn_metadata": attn_md, - } + self.graph_metadata[key] = {"attn_metadata": attn_md} - graph = torch.cuda.CUDAGraph() - # Warmup runs required by CUDA graph semantics. See - # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics + output = None with with_multi_stream(True), piecewise_cuda_graph(False): + # Warmup runs required by CUDA graph semantics. See + # https://pytorch.org/docs/stable/notes/cuda.html#cuda-graph-semantics + # Warmups initialize PyTorch and attention metadata state, and + # resize the shared attention workspace before any graph is captured. for _ in range(self.WARMUP_STEPS): - forward_fn(capture_inputs) + output = forward_fn(capture_inputs) + + if self.is_warmup_only: + return output + graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, pool=self.memory_pool): if self._capture_h2d_copy: # H2D copies for captured inside the graph: at replay @@ -1034,8 +1047,10 @@ def capture( "Encoder CUDA graph does not support nested tensor outputs. " "Disable encoder CUDA graphs for models with ragged outputs.") self.graphs[key] = graph - self.graph_outputs[key] = make_weak_ref(output) + graph_output = make_weak_ref(output) + self.graph_outputs[key] = graph_output self.memory_pool = graph.pool() + return graph_output def replay( self, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 62a442972830..f990e622bdb4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1138,7 +1138,18 @@ def warmup(self, resource_manager: ResourceManager) -> None: # NVSHMEM). gc.collect() torch.cuda.empty_cache() + # Warm up every graph shape before capturing any graph. Attention + # kernels can switch implementations at smaller batch sizes and require + # a larger workspace, so the first pass grows the workspace to its + # maximum size. The second pass runs the final per-shape warmup and + # captures without resizing the workspace. with self.cuda_graph_runner.allow_capture(): + self.cuda_graph_runner.is_warmup_only = True + try: + self._run_cuda_graph_warmup(resource_manager) + finally: + self.cuda_graph_runner.is_warmup_only = False + self.cuda_graph_runner.padding_dummy_requests = {} self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") if can_run_general_warmup: @@ -1466,23 +1477,27 @@ def _get_graphs_to_capture( for draft_len in draft_lengths] def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): - """Captures CUDA graphs for various batch sizes and draft lengths.""" + """Warm up or capture CUDA graphs for the configured graph shapes.""" if not (self.cuda_graph_runner.enabled or self._torch_compile_piecewise_cuda_graph): return self._capture_generation_cuda_graphs(resource_manager) - self._capture_piecewise_cuda_graphs(resource_manager) + # Piecewise graphs have separate capture machinery and do not use the + # whole-model attention workspace. Capture them only on the second pass. + if not self.cuda_graph_runner.is_warmup_only: + self._capture_piecewise_cuda_graphs(resource_manager) def _capture_generation_cuda_graphs(self, resource_manager: ResourceManager): - """Captures CUDA graphs for pure generation steps.""" + """Warm up or capture pure-generation CUDA graph shapes.""" if not self.cuda_graph_runner.enabled: return - logger.info( - f"Creating CUDA graph instances for {len(self._cuda_graph_batch_sizes)} batch sizes." - ) + operation = ("warmup" + if self.cuda_graph_runner.is_warmup_only else "capture") + logger.info(f"Running CUDA graph {operation} for " + f"{len(self._cuda_graph_batch_sizes)} batch sizes.") spec_resource_manager = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) @@ -1490,7 +1505,7 @@ def _capture_generation_cuda_graphs(self, cuda_graph_batch_sizes = sorted(self._cuda_graph_batch_sizes, reverse=True) - # Determine which graphs to capture + # Determine which graph shapes to process. graphs_to_capture = self._get_graphs_to_capture(cuda_graph_batch_sizes, spec_resource_manager) graphs_to_capture = sorted(graphs_to_capture, reverse=True) @@ -1617,7 +1632,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: f"not enough KV cache space.") continue logger.info( - f"Run generation-only CUDA graph warmup ({label}) " + f"Run generation-only CUDA graph {operation} ({label}) " f"for batch size={bs}, draft_len={draft_len}, " f"max_seq_len={max_seq_len}") self.enable_spec_decode = draft_len > 0 or self.is_draft_model or ( @@ -5197,7 +5212,17 @@ def warmup_encoder(self) -> None: torch.cuda.empty_cache() self._run_autotuner_warmup_encoder() + # Warm up every encoder graph shape before capturing any graph. Some + # attention kernels switch implementations at smaller shapes and need + # a larger workspace, so the first pass grows the workspace to its + # maximum size. The second pass runs the final per-shape warmup and + # captures without resizing the workspace. with self.encoder_cuda_graph_runner.allow_capture(): + self.encoder_cuda_graph_runner.is_warmup_only = True + try: + self._run_cuda_graph_warmup_encoder() + finally: + self.encoder_cuda_graph_runner.is_warmup_only = False self._run_cuda_graph_warmup_encoder() # Pre-populate the memory pool with max-shape allocations to reduce @@ -5247,14 +5272,14 @@ def _run_autotuner_warmup_encoder(self) -> None: AutoTuner.get().print_profiling_cache() def _run_cuda_graph_warmup_encoder(self) -> None: - """Captures whole-model CUDA graphs for the encode-only path.""" + """Warm up or capture whole-model encode-only CUDA graphs.""" if not self.encoder_cuda_graph_runner.enabled: return self._capture_encoder_cuda_graphs() def _capture_encoder_cuda_graphs(self) -> None: - """Capture whole-model encoder CUDA graphs for all feasible keys. + """Warm up or capture encoder CUDA graphs for all feasible keys. Feasibility filter (also used in source): nt >= prev_sl + bs (enough tokens for this sl bucket) @@ -5270,8 +5295,9 @@ def _capture_encoder_cuda_graphs(self) -> None: num_tokens_list = sorted(self._cuda_graph_num_tokens) seq_lens_list = sorted(self._cuda_graph_seq_lens) - num_captured = 0 - logger.info("Capturing encoder CUDA graphs ...") + operation = "warmup" if runner.is_warmup_only else "capture" + num_processed = 0 + logger.info(f"Running encoder CUDA graph {operation} ...") for bs in batch_sizes: if bs > self.batch_size: continue @@ -5290,13 +5316,14 @@ def _capture_encoder_cuda_graphs(self) -> None: if inputs is None: continue - logger.info(f"Encoder CUDA graph capture: " + logger.info(f"Encoder CUDA graph {operation}: " f"bs={bs}, nt={nt}, sl={sl}") self.encoder_forward(inputs) torch.cuda.synchronize() - num_captured += 1 + num_processed += 1 - logger.info(f"Captured {num_captured} encoder CUDA graph(s).") + logger.info(f"Completed encoder CUDA graph {operation} for " + f"{num_processed} graph shape(s).") @torch.inference_mode() @with_model_extra_attrs(lambda self: self.model.extra_attrs) @@ -5349,7 +5376,9 @@ def encoder_forward(self, inputs: Dict[str, Any], return self._forward_step(model_inputs, **forward_kwargs) - if self.encoder_cuda_graph_runner.needs_capture(key): + needs_capture = self.encoder_cuda_graph_runner.needs_capture( + key) + if needs_capture: def forward_fn( capture_inputs: Dict[str, Any]) -> Dict[str, Any]: @@ -5359,16 +5388,20 @@ def forward_fn( return self._forward_step(capture_inputs, **forward_kwargs) - self.encoder_cuda_graph_runner.capture( + capture_outputs = self.encoder_cuda_graph_runner.capture( key, forward_fn, { **model_inputs, "_forward_kwargs": forward_kwargs }) - with MoeLoadBalancerIterContext(moe_load_balancer): - graph_outputs = self.encoder_cuda_graph_runner.replay( - key, { - **model_inputs, "_forward_kwargs": forward_kwargs - }) + if self.encoder_cuda_graph_runner.is_warmup_only: + graph_outputs = capture_outputs + else: + with MoeLoadBalancerIterContext(moe_load_balancer): + graph_outputs = self.encoder_cuda_graph_runner.replay( + key, { + **model_inputs, "_forward_kwargs": + forward_kwargs + }) # Return a clone to avoid sharing data_ptr with the static buffers. outputs = {} @@ -5542,7 +5575,8 @@ def forward(self, gather_ids=gather_ids, gather_context_logits=gather_context_logits) else: - if self.cuda_graph_runner.needs_capture(key): + needs_capture = self.cuda_graph_runner.needs_capture(key) + if needs_capture: def capture_forward_fn(inputs: Dict[str, Any]): with MoeLoadBalancerIterContext(moe_load_balancer): @@ -5554,13 +5588,16 @@ def capture_forward_fn(inputs: Dict[str, Any]): def capture_postprocess_fn(inputs: Dict[str, Any]): self._postprocess_inputs(inputs) - self.cuda_graph_runner.capture( + capture_outputs = self.cuda_graph_runner.capture( key, capture_forward_fn, inputs, enable_spec_decode=self.enable_spec_decode, postprocess_fn=capture_postprocess_fn) + if self.cuda_graph_runner.is_warmup_only: + outputs = capture_outputs + elif needs_capture: # Pre-replay: set DSA slot mappings for current batch's draft cache (fixes 2nd warmup) saved_draft = prepare_attn_metadata_for_draft_replay( attn_metadata, draft_kv_cache_manager)