From f96b924313563bb21ebad3eff4a5f5a4f29f9311 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 13 Jul 2026 09:33:02 +0000 Subject: [PATCH 1/5] [None][fix] stabilize attention workspace for CUDA graphs Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 8 +++++-- .../_torch/pyexecutor/model_engine.py | 24 +++++++++++++++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..7e81288eda15 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -663,13 +663,17 @@ def pad_batch(self, scheduled_requests.generation_requests = scheduled_requests.generation_requests[: -padding_size] - def clear(self): - """Releases all captured graphs and the associated memory pool.""" + def clear_captured_graphs(self) -> None: + """Releases captured graphs while retaining the associated memory pool.""" for graph in self.graphs.values(): graph.reset() self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() + + def clear(self): + """Releases all captured graphs and the associated memory pool.""" + self.clear_captured_graphs() self.padding_dummy_requests = {} del self.memory_pool self.memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 26f8562cc474..5c4d50300c8b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -8,7 +8,7 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union import torch import torch._dynamo.config @@ -701,6 +701,7 @@ def __init__( sparse_attention_config=self.sparse_attention_config, ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + self._cuda_graph_replay_enabled = True # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( @@ -877,6 +878,15 @@ def no_cuda_graph(self): finally: self.cuda_graph_runner.enabled = _run_cuda_graphs + @contextmanager + def no_cuda_graph_replay(self) -> Iterator[None]: + replay_enabled = self._cuda_graph_replay_enabled + self._cuda_graph_replay_enabled = False + try: + yield + finally: + self._cuda_graph_replay_enabled = replay_enabled + def _pad_batch_seed_mrope_delta_cache( self, padded_requests: ScheduledRequests) -> None: if not self.use_mrope or padded_requests.num_generation_requests == 0: @@ -1082,6 +1092,12 @@ def warmup(self, resource_manager: ResourceManager) -> None: gc.collect() torch.cuda.empty_cache() with self.cuda_graph_runner.allow_capture(): + # Discover the maximum attention workspace across every graph + # shape before capturing the graphs used at runtime. Some kernels + # select a larger-workspace implementation for a smaller batch size. + with self.no_cuda_graph_replay(): + self._run_cuda_graph_warmup(resource_manager) + self.cuda_graph_runner.clear_captured_graphs() self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") if can_run_general_warmup: @@ -5461,7 +5477,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): @@ -5480,6 +5497,9 @@ def capture_postprocess_fn(inputs: Dict[str, Any]): enable_spec_decode=self.enable_spec_decode, postprocess_fn=capture_postprocess_fn) + if not self._cuda_graph_replay_enabled: + outputs = self.cuda_graph_runner.graph_outputs[key] + 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) From 7735763a651c391870107941f88b47f53321f64f Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:04:34 +0000 Subject: [PATCH 2/5] [None][fix] isolate CUDA graph warmup memory pools Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py | 8 ++------ tensorrt_llm/_torch/pyexecutor/model_engine.py | 7 +++++-- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 7e81288eda15..e0bd91ac52a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -663,17 +663,13 @@ def pad_batch(self, scheduled_requests.generation_requests = scheduled_requests.generation_requests[: -padding_size] - def clear_captured_graphs(self) -> None: - """Releases captured graphs while retaining the associated memory pool.""" + def clear(self): + """Releases all captured graphs and the associated memory pool.""" for graph in self.graphs.values(): graph.reset() self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() - - def clear(self): - """Releases all captured graphs and the associated memory pool.""" - self.clear_captured_graphs() self.padding_dummy_requests = {} del self.memory_pool self.memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5c4d50300c8b..69a630a8b67b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1096,9 +1096,12 @@ def warmup(self, resource_manager: ResourceManager) -> None: # shape before capturing the graphs used at runtime. Some kernels # select a larger-workspace implementation for a smaller batch size. with self.no_cuda_graph_replay(): + memory_pool = self.cuda_graph_runner.memory_pool + self.cuda_graph_runner.memory_pool = None + self._run_cuda_graph_warmup(resource_manager) + self.cuda_graph_runner.clear() + self.cuda_graph_runner.memory_pool = memory_pool self._run_cuda_graph_warmup(resource_manager) - self.cuda_graph_runner.clear_captured_graphs() - self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce From 68417a71b195e057b6b619b11e45f359456f43b6 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:19:02 +0000 Subject: [PATCH 3/5] [None][fix] stabilize encoder CUDA graph workspace Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 39 ++++++++++++++----- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 69a630a8b67b..7461d04bdabb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1091,10 +1091,12 @@ def warmup(self, resource_manager: ResourceManager) -> None: # NVSHMEM). gc.collect() torch.cuda.empty_cache() + # Capture twice without replay. The first pass uses a disposable pool + # to discover the maximum attention workspace across every graph shape. + # Some kernels need more workspace for smaller shapes. Discard those + # graphs and pool, then capture the runtime graphs after the workspace + # is fully sized. with self.cuda_graph_runner.allow_capture(): - # Discover the maximum attention workspace across every graph - # shape before capturing the graphs used at runtime. Some kernels - # select a larger-workspace implementation for a smaller batch size. with self.no_cuda_graph_replay(): memory_pool = self.cuda_graph_runner.memory_pool self.cuda_graph_runner.memory_pool = None @@ -5135,8 +5137,18 @@ def warmup_encoder(self) -> None: torch.cuda.empty_cache() self._run_autotuner_warmup_encoder() + # Capture once with a disposable pool to discover the maximum encoder + # attention workspace across every graph shape. Some kernels need more + # workspace for smaller shapes. Discard those graphs and pool, then + # capture the runtime graphs after the workspace is fully sized. with self.encoder_cuda_graph_runner.allow_capture(): - self._run_cuda_graph_warmup_encoder() + with self.no_cuda_graph_replay(): + memory_pool = self.encoder_cuda_graph_runner.memory_pool + self.encoder_cuda_graph_runner.memory_pool = None + self._run_cuda_graph_warmup_encoder() + self.encoder_cuda_graph_runner.clear() + self.encoder_cuda_graph_runner.memory_pool = memory_pool + self._run_cuda_graph_warmup_encoder() # Pre-populate the memory pool with max-shape allocations to reduce # fragmentation at runtime. @@ -5287,7 +5299,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]: @@ -5302,11 +5316,16 @@ def 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 not self._cuda_graph_replay_enabled: + graph_outputs = self.encoder_cuda_graph_runner.graph_outputs[ + key] + 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 = {} From b538b1384c46b5d0bbf3823258bd43f79b87b3f4 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 14 Jul 2026 04:47:56 +0000 Subject: [PATCH 4/5] [None][fix] split CUDA graph warmup and capture Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 152 +++++++++++++----- .../_torch/pyexecutor/model_engine.py | 93 +++++------ 2 files changed, 157 insertions(+), 88 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e0bd91ac52a8..f2b55c284bc5 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,6 +1,7 @@ import bisect import contextlib from dataclasses import dataclass +from enum import Enum, auto from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeAlias) @@ -39,6 +40,48 @@ KeyType: TypeAlias = Tuple[int, int, bool, bool, bool] +class _CUDAGraphCaptureMode(Enum): + WARMUP_AND_CAPTURE = auto() + WARMUP_ONLY = auto() + CAPTURE_ONLY = auto() + + +class _CUDAGraphCaptureModeMixin: + _capture_mode = _CUDAGraphCaptureMode.WARMUP_AND_CAPTURE + + @contextlib.contextmanager + def warmup_only(self) -> Iterator[None]: + """Run eager CUDA graph warmups without capturing a graph.""" + previous_mode = self._capture_mode + self._capture_mode = _CUDAGraphCaptureMode.WARMUP_ONLY + try: + yield + finally: + self._capture_mode = previous_mode + + @contextlib.contextmanager + def capture_only(self) -> Iterator[None]: + """Capture graphs previously prepared by ``warmup_only``.""" + previous_mode = self._capture_mode + self._capture_mode = _CUDAGraphCaptureMode.CAPTURE_ONLY + try: + yield + finally: + self._capture_mode = previous_mode + + @property + def is_warmup_only(self) -> bool: + return self._capture_mode == _CUDAGraphCaptureMode.WARMUP_ONLY + + @property + def is_capture_only(self) -> bool: + return self._capture_mode == _CUDAGraphCaptureMode.CAPTURE_ONLY + + @property + def _should_warmup(self) -> bool: + return not self.is_capture_only + + @dataclass class CUDAGraphRunnerConfig: """Configuration for the CUDAGraphRunner, passed from the ModelEngine.""" @@ -92,7 +135,7 @@ class CUDAGraphRunnerConfig: sparse_attention_config: Optional[BaseSparseAttentionConfig] = None -class CUDAGraphRunner: +class CUDAGraphRunner(_CUDAGraphCaptureModeMixin): """ Manages the lifecycle and execution of CUDA graphs for the model engine. @@ -321,9 +364,13 @@ 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 + elif self.is_capture_only: + raise RuntimeError( + f"CUDA graph key {key} was not prepared by the warmup-only pass" + ) # Graph doesn't exist yet. If on-the-fly capture is not allowed, # fall back to eager so the caller doesn't need a separate check. @@ -378,8 +425,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). @@ -408,11 +455,6 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) - self.graph_metadata[key] = { - "attn_metadata": initial_inputs["attn_metadata"], - "spec_metadata": initial_inputs.get("spec_metadata", None), - } - def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs: Dict[str, Any]): is_first_draft = key[2] @@ -422,18 +464,36 @@ 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): - for _ in range(self.WARMUP_STEPS): - _setup_spec_decoding_and_forward(key, forward_fn, - capture_inputs) - if postprocess_fn is not None: - postprocess_fn(capture_inputs) - + if self._should_warmup: + # 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 + # and resize the shared attention workspace before any graph is captured. + for _ in range(self.WARMUP_STEPS): + output = _setup_spec_decoding_and_forward( + key, forward_fn, capture_inputs) + if postprocess_fn is not None: + postprocess_fn(capture_inputs) + + if self.is_capture_only: + metadata = self.graph_metadata.get(key) + if metadata is None: + raise RuntimeError( + f"CUDA graph key {key} was not prepared by the warmup-only pass" + ) + + self.graph_metadata[key] = { + "attn_metadata": initial_inputs["attn_metadata"], + "spec_metadata": initial_inputs.get("spec_metadata", None), + } + + 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 +501,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]: @@ -694,7 +756,7 @@ class EncoderCUDAGraphRunnerConfig: cuda_graph_mem_pool: Any -class EncoderCUDAGraphRunner: +class EncoderCUDAGraphRunner(_CUDAGraphCaptureModeMixin): """CUDA graph runner for no-cache encoder forward passes. Designed for encoder inputs with `input_ids` (flat [total_tokens]) and @@ -916,8 +978,12 @@ 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 + elif self.is_capture_only: + raise RuntimeError( + f"Encoder CUDA graph key {key} was not prepared by the warmup-only pass" + ) # New key not yet captured. Only create metadata if capture is # allowed (warmup time); otherwise fall back to eager. @@ -980,8 +1046,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 +1069,29 @@ def capture( attn_md = capture_inputs["attn_metadata"] - 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): - for _ in range(self.WARMUP_STEPS): - forward_fn(capture_inputs) - + if self._should_warmup: + # 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): + output = forward_fn(capture_inputs) + + if self.is_capture_only: + metadata = self.graph_metadata.get(key) + if metadata is None: + raise RuntimeError( + f"Encoder CUDA graph key {key} was not prepared by the warmup-only pass" + ) + + self.graph_metadata[key] = {"attn_metadata": attn_md} + + 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 +1112,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 7461d04bdabb..b3adc7fed3b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -8,7 +8,7 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch._dynamo.config @@ -701,7 +701,6 @@ def __init__( sparse_attention_config=self.sparse_attention_config, ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) - self._cuda_graph_replay_enabled = True # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( @@ -878,15 +877,6 @@ def no_cuda_graph(self): finally: self.cuda_graph_runner.enabled = _run_cuda_graphs - @contextmanager - def no_cuda_graph_replay(self) -> Iterator[None]: - replay_enabled = self._cuda_graph_replay_enabled - self._cuda_graph_replay_enabled = False - try: - yield - finally: - self._cuda_graph_replay_enabled = replay_enabled - def _pad_batch_seed_mrope_delta_cache( self, padded_requests: ScheduledRequests) -> None: if not self.use_mrope or padded_requests.num_generation_requests == 0: @@ -1091,18 +1081,15 @@ def warmup(self, resource_manager: ResourceManager) -> None: # NVSHMEM). gc.collect() torch.cuda.empty_cache() - # Capture twice without replay. The first pass uses a disposable pool - # to discover the maximum attention workspace across every graph shape. - # Some kernels need more workspace for smaller shapes. Discard those - # graphs and pool, then capture the runtime graphs after the workspace - # is fully sized. + # 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 captures without resizing it. with self.cuda_graph_runner.allow_capture(): - with self.no_cuda_graph_replay(): - memory_pool = self.cuda_graph_runner.memory_pool - self.cuda_graph_runner.memory_pool = None + with self.cuda_graph_runner.warmup_only(): self._run_cuda_graph_warmup(resource_manager) - self.cuda_graph_runner.clear() - self.cuda_graph_runner.memory_pool = memory_pool + self.cuda_graph_runner.padding_dummy_requests = {} + with self.cuda_graph_runner.capture_only(): self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") if can_run_general_warmup: @@ -1430,23 +1417,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) @@ -1454,7 +1445,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) @@ -1581,7 +1572,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 ( @@ -5137,17 +5128,14 @@ def warmup_encoder(self) -> None: torch.cuda.empty_cache() self._run_autotuner_warmup_encoder() - # Capture once with a disposable pool to discover the maximum encoder - # attention workspace across every graph shape. Some kernels need more - # workspace for smaller shapes. Discard those graphs and pool, then - # capture the runtime graphs after the workspace is fully sized. + # 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 capture pass can only start after the + # workspace has reached its maximum size. with self.encoder_cuda_graph_runner.allow_capture(): - with self.no_cuda_graph_replay(): - memory_pool = self.encoder_cuda_graph_runner.memory_pool - self.encoder_cuda_graph_runner.memory_pool = None + with self.encoder_cuda_graph_runner.warmup_only(): self._run_cuda_graph_warmup_encoder() - self.encoder_cuda_graph_runner.clear() - self.encoder_cuda_graph_runner.memory_pool = memory_pool + with self.encoder_cuda_graph_runner.capture_only(): self._run_cuda_graph_warmup_encoder() # Pre-populate the memory pool with max-shape allocations to reduce @@ -5197,14 +5185,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) @@ -5220,8 +5208,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 @@ -5240,13 +5229,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) @@ -5311,14 +5301,13 @@ 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 }) - if not self._cuda_graph_replay_enabled: - graph_outputs = self.encoder_cuda_graph_runner.graph_outputs[ - key] + 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( @@ -5512,15 +5501,15 @@ 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 not self._cuda_graph_replay_enabled: - outputs = self.cuda_graph_runner.graph_outputs[key] + 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( From 3c24394667f2c589ef9f6427e5d26bd6d5e36511 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 15 Jul 2026 08:36:33 +0000 Subject: [PATCH 5/5] [None][fix] simplify CUDA graph warmup state Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 123 +++++------------- .../_torch/pyexecutor/model_engine.py | 24 ++-- 2 files changed, 44 insertions(+), 103 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index f2b55c284bc5..a254eccaf9d2 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,7 +1,6 @@ import bisect import contextlib from dataclasses import dataclass -from enum import Enum, auto from typing import (Any, Callable, Dict, Iterator, List, Optional, Tuple, TypeAlias) @@ -40,48 +39,6 @@ KeyType: TypeAlias = Tuple[int, int, bool, bool, bool] -class _CUDAGraphCaptureMode(Enum): - WARMUP_AND_CAPTURE = auto() - WARMUP_ONLY = auto() - CAPTURE_ONLY = auto() - - -class _CUDAGraphCaptureModeMixin: - _capture_mode = _CUDAGraphCaptureMode.WARMUP_AND_CAPTURE - - @contextlib.contextmanager - def warmup_only(self) -> Iterator[None]: - """Run eager CUDA graph warmups without capturing a graph.""" - previous_mode = self._capture_mode - self._capture_mode = _CUDAGraphCaptureMode.WARMUP_ONLY - try: - yield - finally: - self._capture_mode = previous_mode - - @contextlib.contextmanager - def capture_only(self) -> Iterator[None]: - """Capture graphs previously prepared by ``warmup_only``.""" - previous_mode = self._capture_mode - self._capture_mode = _CUDAGraphCaptureMode.CAPTURE_ONLY - try: - yield - finally: - self._capture_mode = previous_mode - - @property - def is_warmup_only(self) -> bool: - return self._capture_mode == _CUDAGraphCaptureMode.WARMUP_ONLY - - @property - def is_capture_only(self) -> bool: - return self._capture_mode == _CUDAGraphCaptureMode.CAPTURE_ONLY - - @property - def _should_warmup(self) -> bool: - return not self.is_capture_only - - @dataclass class CUDAGraphRunnerConfig: """Configuration for the CUDAGraphRunner, passed from the ModelEngine.""" @@ -135,7 +92,7 @@ class CUDAGraphRunnerConfig: sparse_attention_config: Optional[BaseSparseAttentionConfig] = None -class CUDAGraphRunner(_CUDAGraphCaptureModeMixin): +class CUDAGraphRunner: """ Manages the lifecycle and execution of CUDA graphs for the model engine. @@ -143,7 +100,7 @@ class CUDAGraphRunner(_CUDAGraphCaptureModeMixin): 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 @@ -175,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.""" @@ -367,10 +325,6 @@ def maybe_get_cuda_graph( if key in self.graph_metadata: return self.graph_metadata[key][ "attn_metadata"], self.graph_metadata[key]["spec_metadata"], key - elif self.is_capture_only: - raise RuntimeError( - f"CUDA graph key {key} was not prepared by the warmup-only pass" - ) # Graph doesn't exist yet. If on-the-fly capture is not allowed, # fall back to eager so the caller doesn't need a separate check. @@ -455,6 +409,11 @@ def capture(self, capture_inputs = initial_inputs.copy() capture_inputs.update(sliced_static_tensors) + self.graph_metadata[key] = { + "attn_metadata": initial_inputs["attn_metadata"], + "spec_metadata": initial_inputs.get("spec_metadata", None), + } + def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, capture_inputs: Dict[str, Any]): is_first_draft = key[2] @@ -466,29 +425,16 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, output = None with with_multi_stream(True), piecewise_cuda_graph(False): - if self._should_warmup: - # 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 - # and resize the shared attention workspace before any graph is captured. - for _ in range(self.WARMUP_STEPS): - output = _setup_spec_decoding_and_forward( - key, forward_fn, capture_inputs) - if postprocess_fn is not None: - postprocess_fn(capture_inputs) - - if self.is_capture_only: - metadata = self.graph_metadata.get(key) - if metadata is None: - raise RuntimeError( - f"CUDA graph key {key} was not prepared by the warmup-only pass" - ) - - self.graph_metadata[key] = { - "attn_metadata": initial_inputs["attn_metadata"], - "spec_metadata": initial_inputs.get("spec_metadata", None), - } + # 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): + 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 @@ -756,7 +702,7 @@ class EncoderCUDAGraphRunnerConfig: cuda_graph_mem_pool: Any -class EncoderCUDAGraphRunner(_CUDAGraphCaptureModeMixin): +class EncoderCUDAGraphRunner: """CUDA graph runner for no-cache encoder forward passes. Designed for encoder inputs with `input_ids` (flat [total_tokens]) and @@ -766,7 +712,7 @@ class EncoderCUDAGraphRunner(_CUDAGraphCaptureModeMixin): 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 @@ -792,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 @@ -980,10 +927,6 @@ def maybe_get_cuda_graph( if key in self.graph_metadata: return self.graph_metadata[key]["attn_metadata"], key - elif self.is_capture_only: - raise RuntimeError( - f"Encoder CUDA graph key {key} was not prepared by the warmup-only pass" - ) # New key not yet captured. Only create metadata if capture is # allowed (warmup time); otherwise fall back to eager. @@ -1069,24 +1012,16 @@ def capture( attn_md = capture_inputs["attn_metadata"] + self.graph_metadata[key] = {"attn_metadata": attn_md} + output = None with with_multi_stream(True), piecewise_cuda_graph(False): - if self._should_warmup: - # 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): - output = forward_fn(capture_inputs) - - if self.is_capture_only: - metadata = self.graph_metadata.get(key) - if metadata is None: - raise RuntimeError( - f"Encoder CUDA graph key {key} was not prepared by the warmup-only pass" - ) - - self.graph_metadata[key] = {"attn_metadata": attn_md} + # 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): + output = forward_fn(capture_inputs) if self.is_warmup_only: return output diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b3adc7fed3b3..1ac2c86357f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1084,13 +1084,16 @@ def warmup(self, resource_manager: ResourceManager) -> None: # 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 captures without resizing it. + # maximum size. The second pass runs the final per-shape warmup and + # captures without resizing the workspace. with self.cuda_graph_runner.allow_capture(): - with self.cuda_graph_runner.warmup_only(): + 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 = {} - with self.cuda_graph_runner.capture_only(): - self._run_cuda_graph_warmup(resource_manager) + self._run_cuda_graph_warmup(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") if can_run_general_warmup: # Pre-populate the memory pool with max-shape allocations to reduce @@ -5130,13 +5133,16 @@ def warmup_encoder(self) -> None: 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 capture pass can only start after the - # workspace has reached its maximum size. + # 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(): - with self.encoder_cuda_graph_runner.warmup_only(): - self._run_cuda_graph_warmup_encoder() - with self.encoder_cuda_graph_runner.capture_only(): + 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 # fragmentation at runtime.