[https://nvbugs/6248648][fix] Wrap warmup+capture in try/except to graph.reset() the orphan then re-raise…

tensorrt_llm/_torch/compilation/piecewise_optimizer.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import dataclasses
 from typing import Callable, List, Optional, Sequence, Union
 from unittest.mock import patch
@@ -14,7 +17,7 @@
 from ..utils import (get_model_extra_attrs,
                      get_per_request_piecewise_cuda_graph_flag,
                      get_piecewise_cuda_graph_flag, make_weak_ref,
-                     set_piecewise_running)
+                     safe_reset_cuda_graph, set_piecewise_running)
 from .multi_stream.auto_multi_stream import multi_stream_schedule
 from .utils import (get_capture_piecewise_cuda_graph_flag,
                     get_optional_trtllm_op, is_call_function)
@@ -179,18 +182,20 @@ def __init__(
                 callable=default_callable,
             )
 
-    def clear_cuda_graphs(self):
+    def clear_cuda_graphs(self) -> None:
         """Release captures while retaining buckets for a later warmup."""
         for entry in self.entries.values():
             if entry.cuda_graph is not None:
-                entry.cuda_graph.reset()
+                safe_reset_cuda_graph(entry.cuda_graph,
+                                      "during piecewise cleanup")
             entry.cuda_graph = None
             entry.warmup_count = 0
             entry.input_addresses = None
             entry.output_addresses = None
             entry.output = None
 
     def __call__(self, *args):
+        """Run eagerly, capture a bucket, or replay its captured graph."""
         runtime_num_of_token = None
         if self.runtime_num_tokens_idx != None:
             runtime_num_of_token = int(
@@ -232,25 +237,37 @@ def __call__(self, *args):
 
             graph = torch.cuda.CUDAGraph()
 
-            # Torch's cuda graph will call gc.collect() internally. This will slow down the performance.
-            # We patch it to do nothing.
-            with patch("gc.collect", lambda: None):
-                # TODO: consider to use `make_graphed_callables()` when
-                # it's ready rather than capture it ourselves
-                # Graph Capture would override the stream. We need to setup the stream correctly.
-                extra_attrs = get_model_extra_attrs()
-                with torch.cuda.graph(graph, pool=self.graph_pool_handle):
-                    extra_attrs["global_stream"] = torch.cuda.current_stream()
-                    output = entry.callable(*args)
-                extra_attrs["global_stream"] = torch.cuda.current_stream()
+            try:
+                # Torch's cuda graph will call gc.collect() internally. This will slow down the performance.
+                # We patch it to do nothing.
+                with patch("gc.collect", lambda: None):
+                    # TODO: consider to use `make_graphed_callables()` when
+                    # it's ready rather than capture it ourselves
+                    # Graph Capture would override the stream. We need to setup the stream correctly.
+                    extra_attrs = get_model_extra_attrs()
+                    try:
+                        with torch.cuda.graph(graph,
+                                              pool=self.graph_pool_handle):
+                            extra_attrs[
+                                "global_stream"] = torch.cuda.current_stream()
+                            output = entry.callable(*args)
+                    finally:
+                        extra_attrs[
+                            "global_stream"] = torch.cuda.current_stream()
+
+                # Mark weak ref here. The intermediate activation tensor should be freed properly.
+                # Here we don't use python native weakref since we still need the object to be alive when the graph is replayed.
+                graph_output = make_weak_ref(output)
+                output_addresses = [
+                    i.data_ptr() for i in output if isinstance(i, torch.Tensor)
+                ]
+            except Exception:
+                safe_reset_cuda_graph(graph, "before piecewise graph commit")
+                raise
 
             entry.cuda_graph = graph
-            # Mark weak ref here. The intermediate activation tensor should be freed properly.
-            # Here we don't use python native weakref since we still need the object to be alive when the graph is replayed.
-            entry.output = make_weak_ref(output)
-            entry.output_addresses = [
-                i.data_ptr() for i in output if isinstance(i, torch.Tensor)
-            ]
+            entry.output = graph_output
+            entry.output_addresses = output_addresses
 
             entry.cuda_graph.replay()
 

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -1,5 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import bisect
 import contextlib
+import sys
 from dataclasses import dataclass
 from typing import (Any, Callable, Dict, Iterator, List, NamedTuple, Optional,
                     Tuple, TypeAlias)
@@ -22,7 +26,7 @@
 from ..speculative.interface import SpecMetadata
 from ..speculative.spec_sampler_base import SampleStateTensorsSpec
 from ..speculative.utils import get_draft_kv_cache_manager
-from ..utils import make_weak_ref, piecewise_cuda_graph
+from ..utils import make_weak_ref, piecewise_cuda_graph, safe_reset_cuda_graph
 from .llm_request import LlmRequest, get_draft_token_length
 from .resource_manager import (BaseResourceManager, ResourceManager,
                                ResourceManagerType)
@@ -642,7 +646,7 @@ def capture(self,
         saved_kv_lens_cuda = _save_spec_decode_capture_state(
             attn_metadata, enable_spec_decode)
 
-        self.graph_metadata[key] = {
+        graph_metadata = {
             "attn_metadata": attn_metadata,
             "spec_metadata": initial_inputs.get("spec_metadata", None),
         }
@@ -656,37 +660,61 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable,
                 capture_inputs['attn_metadata'].use_spec_decoding = True
             return forward_fn(capture_inputs)
 
+        @contextlib.contextmanager
+        def _restore_state_after_forward() -> Iterator[None]:
+            try:
+                yield
+            finally:
+                active_error = sys.exc_info()[1]
+                try:
+                    _restore_spec_decode_capture_state(attn_metadata,
+                                                       saved_kv_lens_cuda)
+                except RuntimeError as restore_error:
+                    if active_error is None:
+                        raise
+                    logger.warning(
+                        "Failed to restore speculative-decoding state after "
+                        f"CUDA graph forward failure: {restore_error}")
+
         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):
-                output = _setup_spec_decoding_and_forward(
-                    key, forward_fn, capture_inputs)
-                if postprocess_fn is not None:
-                    postprocess_fn(capture_inputs)
-                _restore_spec_decode_capture_state(attn_metadata,
-                                                   saved_kv_lens_cuda)
-
-            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)
-            if postprocess_fn is not None:
-                postprocess_fn(capture_inputs)
-            _restore_spec_decode_capture_state(attn_metadata,
-                                               saved_kv_lens_cuda)
+        graph = None
+        try:
+            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):
+                    with _restore_state_after_forward():
+                        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:
+                    self.graph_metadata[key] = graph_metadata
+                    return output
+
+                graph = torch.cuda.CUDAGraph()
+                with _restore_state_after_forward():
+                    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)
+
+            graph_output = make_weak_ref(output)
+            memory_pool = graph.pool()
+        except Exception:
+            if graph is not None:
+                safe_reset_cuda_graph(graph, "before graph commit")
+            raise
 
         self.graphs[key] = graph
-        graph_output = make_weak_ref(output)
         self.graph_outputs[key] = graph_output
-        self.memory_pool = graph.pool()
+        self.graph_metadata[key] = graph_metadata
+        self.memory_pool = memory_pool
         return graph_output
 
     def replay(self, key: KeyType,
@@ -1013,10 +1041,10 @@ def pad_batch(self,
                 scheduled_requests.generation_requests = scheduled_requests.generation_requests[:
                                                                                                 -padding_size]
 
-    def clear(self):
+    def clear(self) -> None:
         """Releases all captured graphs and the associated memory pool."""
         for graph in self.graphs.values():
-            graph.reset()
+            safe_reset_cuda_graph(graph, "during decoder cleanup")
         self.graphs.clear()
         self.graph_outputs.clear()
         self.graph_metadata.clear()
@@ -1736,7 +1764,7 @@ def capture(
 
         attn_md = capture_inputs["attn_metadata"]
 
-        self.graph_metadata[key] = {"attn_metadata": attn_md}
+        graph_metadata = {"attn_metadata": attn_md}
 
         # Warmup must see the same runtime data as capture. In particular,
         # graph metadata initializes _seq_lens_cuda to ones, while
@@ -1753,43 +1781,54 @@ def capture(
         torch.cuda.current_stream().synchronize()
 
         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):
-                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,
-                                  capture_error_mode="thread_local"):
-                if self._capture_h2d_copy:
-                    # H2D copies for captured inside the graph: at replay
-                    # time it re-issues from the pinned static buffer without
-                    # an eager driver call.
-                    capture_inputs["input_ids"].copy_(
-                        sliced_static_tensors_cpu["input_ids"],
-                        non_blocking=True)
-                    capture_inputs["position_ids"].copy_(
-                        sliced_static_tensors_cpu["position_ids"],
-                        non_blocking=True)
-                    attn_md._seq_lens_cuda.copy_(attn_md._seq_lens,
-                                                 non_blocking=True)
-                output = forward_fn(capture_inputs)
-
-        if self._contains_nested_tensor(output):
-            raise TypeError(
-                "Encoder CUDA graph does not support nested tensor outputs. "
-                "Disable encoder CUDA graphs for models with ragged outputs.")
+        graph = None
+        try:
+            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):
+                    output = forward_fn(capture_inputs)
+
+                if self.is_warmup_only:
+                    self.graph_metadata[key] = graph_metadata
+                    return output
+
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(graph,
+                                      pool=self.memory_pool,
+                                      capture_error_mode="thread_local"):
+                    if self._capture_h2d_copy:
+                        # H2D copies for captured inside the graph: at replay
+                        # time it re-issues from the pinned static buffer without
+                        # an eager driver call.
+                        capture_inputs["input_ids"].copy_(
+                            sliced_static_tensors_cpu["input_ids"],
+                            non_blocking=True)
+                        capture_inputs["position_ids"].copy_(
+                            sliced_static_tensors_cpu["position_ids"],
+                            non_blocking=True)
+                        attn_md._seq_lens_cuda.copy_(attn_md._seq_lens,
+                                                     non_blocking=True)
+                    output = forward_fn(capture_inputs)
+
+            if self._contains_nested_tensor(output):
+                raise TypeError(
+                    "Encoder CUDA graph does not support nested tensor outputs. "
+                    "Disable encoder CUDA graphs for models with ragged outputs."
+                )
+            graph_output = make_weak_ref(output)
+            memory_pool = graph.pool()
+        except Exception:
+            if graph is not None:
+                safe_reset_cuda_graph(graph, "before encoder graph commit")
+            raise
+
         self.graphs[key] = graph
-        graph_output = make_weak_ref(output)
         self.graph_outputs[key] = graph_output
-        self.memory_pool = graph.pool()
+        self.graph_metadata[key] = graph_metadata
+        self.memory_pool = memory_pool
         return graph_output
 
     def retire_staging(self) -> None:
@@ -1824,9 +1863,10 @@ def replay(
     def get_graph_pool(self):
         return self.memory_pool
 
-    def clear(self):
+    def clear(self) -> None:
+        """Release all captured encoder graphs and their shared memory pool."""
         for graph in self.graphs.values():
-            graph.reset()
+            safe_reset_cuda_graph(graph, "during encoder cleanup")
         self.graphs.clear()
         self.graph_outputs.clear()
         self.graph_metadata.clear()

tensorrt_llm/_torch/utils.py

@@ -14,6 +14,7 @@
 
 from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor,
                                  get_sm_version, torch_dtype_to_str)
+from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.math_utils import ceil_div, pad_up
 from tensorrt_llm.quantization.utils import fp4_utils
@@ -387,6 +388,14 @@ def get_piecewise_cuda_graph_flag() -> bool:
     return _enable_piecewise_cuda_graph
 
 
+def safe_reset_cuda_graph(graph: torch.cuda.CUDAGraph, context: str) -> None:
+    """Reset a CUDA graph without letting teardown failure mask an error."""
+    try:
+        graph.reset()
+    except RuntimeError as error:
+        logger.warning(f"Failed to reset CUDA graph {context}: {error}")
+
+
 @contextlib.contextmanager
 def piecewise_cuda_graph(enable: bool):
     prev_enable = get_piecewise_cuda_graph_flag()