From 74884f5b3dea1949a20637070b4d3cef76aa51a8 Mon Sep 17 00:00:00 2001 From: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> Date: Wed, 3 Jun 2026 08:48:17 -0700 Subject: [PATCH 1/2] [nvbugs/6248648][fix] Harden CUDAGraph capture/teardown against destructor abort Signed-off-by: tensorrt-cicd <90828364+tensorrt-cicd@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 38 +++++++++++++++---- 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 4019c037f513..d515c4783b4c 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -675,13 +675,23 @@ def _setup_spec_decoding_and_forward(key: KeyType, forward_fn: Callable, 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) + try: + 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) + except Exception: + # Reset the partially-captured graph now, while the CUDA generator + # state is still valid. Otherwise this orphaned graph (it was never + # stored in self.graphs) is destroyed later during GC, when the + # generator state may already be gone, and ~CUDAGraph()'s + # unregister_graph can abort the process via terminate(), masking + # the real capture-time error. + self._safe_reset_graph(graph, "after a capture failure") + raise self.graphs[key] = graph graph_output = make_weak_ref(output) @@ -1013,10 +1023,22 @@ def pad_batch(self, scheduled_requests.generation_requests = scheduled_requests.generation_requests[: -padding_size] + @staticmethod + def _safe_reset_graph(graph: torch.cuda.CUDAGraph, context: str): + # graph.reset() can raise (e.g. a stale CUDA generator state inside + # ~CUDAGraph()); swallow it so one failing reset cannot abort the rest + # of the teardown or mask an earlier, more relevant error. + try: + graph.reset() + except Exception: + logger.warning("Failed to reset CUDA graph %s.", context) + def clear(self): """Releases all captured graphs and the associated memory pool.""" + # Reset each graph independently so a failure tearing down one graph + # does not abort cleanup of the remaining graphs. for graph in self.graphs.values(): - graph.reset() + self._safe_reset_graph(graph, "during cleanup") self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() From 2c49ca238e84b591b865593cf7d795fefb70185d Mon Sep 17 00:00:00 2001 From: William Zhang Date: Sat, 8 Aug 2026 20:26:09 -0700 Subject: [PATCH 2/2] [nvbugs/6248648][fix] Harden active CUDA graph teardown paths Signed-off-by: William Zhang --- .../_torch/compilation/piecewise_optimizer.py | 57 ++-- .../_torch/pyexecutor/cuda_graph_runner.py | 200 +++++++------ tensorrt_llm/_torch/utils.py | 9 + .../test_piecewise_cuda_graph_cleanup.py | 100 +++++++ .../executor/test_cuda_graph_cleanup.py | 274 ++++++++++++++++++ 5 files changed, 529 insertions(+), 111 deletions(-) create mode 100644 tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py create mode 100644 tests/unittest/_torch/executor/test_cuda_graph_cleanup.py diff --git a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py b/tensorrt_llm/_torch/compilation/piecewise_optimizer.py index 73164f885660..813538b624a3 100644 --- a/tensorrt_llm/_torch/compilation/piecewise_optimizer.py +++ b/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,11 +182,12 @@ 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 @@ -191,6 +195,7 @@ def clear_cuda_graphs(self): 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() diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index d515c4783b4c..8c8d5178e7f9 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/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,47 +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) - 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() + @contextlib.contextmanager + def _restore_state_after_forward() -> Iterator[None]: try: - 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) - except Exception: - # Reset the partially-captured graph now, while the CUDA generator - # state is still valid. Otherwise this orphaned graph (it was never - # stored in self.graphs) is destroyed later during GC, when the - # generator state may already be gone, and ~CUDAGraph()'s - # unregister_graph can abort the process via terminate(), masking - # the real capture-time error. - self._safe_reset_graph(graph, "after a capture failure") - raise + 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 + 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, @@ -1023,22 +1041,10 @@ def pad_batch(self, scheduled_requests.generation_requests = scheduled_requests.generation_requests[: -padding_size] - @staticmethod - def _safe_reset_graph(graph: torch.cuda.CUDAGraph, context: str): - # graph.reset() can raise (e.g. a stale CUDA generator state inside - # ~CUDAGraph()); swallow it so one failing reset cannot abort the rest - # of the teardown or mask an earlier, more relevant error. - try: - graph.reset() - except Exception: - logger.warning("Failed to reset CUDA graph %s.", context) - - def clear(self): + def clear(self) -> None: """Releases all captured graphs and the associated memory pool.""" - # Reset each graph independently so a failure tearing down one graph - # does not abort cleanup of the remaining graphs. for graph in self.graphs.values(): - self._safe_reset_graph(graph, "during cleanup") + safe_reset_cuda_graph(graph, "during decoder cleanup") self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() @@ -1758,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 @@ -1775,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: @@ -1846,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() diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index cb62ec99f76b..8a221a609777 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/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() diff --git a/tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py b/tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py new file mode 100644 index 000000000000..e60f1968f4d8 --- /dev/null +++ b/tests/unittest/_torch/compilation/test_piecewise_cuda_graph_cleanup.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext +from typing import Callable +from unittest.mock import patch + +import pytest +import torch +from torch.fx import symbolic_trace + +from tensorrt_llm._torch.compilation.piecewise_optimizer import PiecewiseRunner +from tensorrt_llm._torch.compilation.utils import capture_piecewise_cuda_graph +from tensorrt_llm._torch.utils import piecewise_cuda_graph + + +class _ResetFailureGraph: + def reset(self) -> None: + raise RuntimeError("reset failed") + + +class _RecordingGraph: + def __init__(self) -> None: + self.reset_count = 0 + + def reset(self) -> None: + self.reset_count += 1 + + +def _make_runner(default_callable: Callable[[], object] | None = None) -> PiecewiseRunner: + return PiecewiseRunner( + graph=symbolic_trace(torch.nn.Identity()), + name="test", + compile_time_num_tokens=1, + runtime_num_tokens_idx=None, + capture_num_tokens=[1, 2], + graph_pool_handle=None, + default_callable=default_callable or (lambda: None), + enable_inductor=False, + is_first_runner=False, + is_last_runner=False, + ) + + +def test_clear_continues_after_reset_failure() -> None: + runner = _make_runner() + remaining_graph = _RecordingGraph() + runner.entries[1].cuda_graph = _ResetFailureGraph() + runner.entries[2].cuda_graph = remaining_graph + + runner.clear_cuda_graphs() + + assert remaining_graph.reset_count == 1 + for entry in runner.entries.values(): + assert entry.cuda_graph is None + assert entry.warmup_count == 0 + assert entry.input_addresses is None + assert entry.output_addresses is None + assert entry.output is None + + +def test_capture_failure_resets_graph_before_entry_commit() -> None: + def fail_capture() -> None: + raise ValueError("capture failed") + + runner = _make_runner(fail_capture) + runner.entries[1].warmup_count = 3 + graph = _RecordingGraph() + capture_stream = object() + restored_stream = object() + extra_attrs = {} + + with ( + piecewise_cuda_graph(True), + capture_piecewise_cuda_graph(True), + patch("torch.cuda.CUDAGraph", return_value=graph), + patch("torch.cuda.graph", return_value=nullcontext()), + patch("torch.cuda.current_stream", side_effect=[capture_stream, restored_stream]), + patch( + "tensorrt_llm._torch.compilation.piecewise_optimizer.get_model_extra_attrs", + return_value=extra_attrs, + ), + pytest.raises(ValueError, match="capture failed"), + ): + runner() + + assert graph.reset_count == 1 + assert runner.entries[1].cuda_graph is None + assert extra_attrs["global_stream"] is restored_stream diff --git a/tests/unittest/_torch/executor/test_cuda_graph_cleanup.py b/tests/unittest/_torch/executor/test_cuda_graph_cleanup.py new file mode 100644 index 000000000000..45bffff68d90 --- /dev/null +++ b/tests/unittest/_torch/executor/test_cuda_graph_cleanup.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import nullcontext +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + EncoderCUDAGraphRunner, + EncoderCUDAGraphRunnerConfig, + KeyType, +) + + +class _ResetFailureGraph: + def reset(self) -> None: + raise RuntimeError("reset failed") + + +class _RecordingGraph: + def __init__(self) -> None: + self.reset_count = 0 + + def reset(self) -> None: + self.reset_count += 1 + + +class _SpecDecMode: + def needs_kv_cache_recompute(self) -> bool: + return False + + +class _SpecConfig: + spec_dec_mode = _SpecDecMode() + + +class _AttentionMetadata: + def __init__(self) -> None: + self.kv_lens_cuda = torch.tensor([3]) + self._seq_lens = torch.tensor([1]) + self._seq_lens_cuda = torch.tensor([1]) + self.num_seqs = 1 + self.update_count = 0 + + def on_update_kv_lens(self) -> None: + self.update_count += 1 + + +def _make_decoder_runner() -> CUDAGraphRunner: + runner = CUDAGraphRunner( + CUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[], + max_cuda_graph_batch_size=0, + max_beam_width=1, + max_num_tokens=1, + spec_config=None, + cuda_graph_mem_pool=None, + use_mrope=False, + original_max_draft_len=0, + original_max_total_draft_tokens=0, + is_draft_model=False, + enable_attention_dp=False, + is_encoder_decoder=False, + batch_size=0, + mapping=None, + dist=None, + kv_cache_manager_key=None, + ) + ) + runner.shared_static_tensors = { + "input_ids": torch.empty(1), + "position_ids": torch.empty((1, 1)), + } + return runner + + +def _make_encoder_runner() -> EncoderCUDAGraphRunner: + return EncoderCUDAGraphRunner( + EncoderCUDAGraphRunnerConfig( + use_cuda_graph=False, + cuda_graph_padding_enabled=False, + cuda_graph_batch_sizes=[], + cuda_graph_num_tokens=[], + cuda_graph_seq_lens=[], + max_cuda_graph_batch_size=0, + max_cuda_graph_num_tokens=0, + max_num_tokens=1, + max_seq_len=1, + cuda_graph_mem_pool=None, + ) + ) + + +def test_decoder_clear_continues_after_reset_failure() -> None: + runner = _make_decoder_runner() + remaining_graph = _RecordingGraph() + runner.graphs = {"broken": _ResetFailureGraph(), "remaining": remaining_graph} + runner.graph_outputs = {"output": object()} + runner.graph_metadata = {"metadata": object()} + runner.padding_dummy_requests = {"request": object()} + + with ( + patch("torch.cuda.empty_cache"), + patch("tensorrt_llm._torch.utils.logger.warning") as warning, + ): + runner.clear() + + assert remaining_graph.reset_count == 1 + warning.assert_called_once() + assert warning.call_args.args == ( + "Failed to reset CUDA graph during decoder cleanup: reset failed", + ) + assert runner.graphs == {} + assert runner.graph_outputs == {} + assert runner.graph_metadata == {} + assert runner.padding_dummy_requests == {} + assert runner.memory_pool is None + + +def test_decoder_warmup_failure_restores_state_without_publishing_metadata() -> None: + runner = _make_decoder_runner() + runner.config.spec_config = _SpecConfig() + metadata = _AttentionMetadata() + + def fail_after_mutating_state(inputs: dict[str, object]) -> None: + inputs["attn_metadata"].kv_lens_cuda.add_(2) + raise ValueError("warmup failed") + + with pytest.raises(ValueError, match="warmup failed"): + runner.capture( + (1, 0, False, False, False), + fail_after_mutating_state, + {"attn_metadata": metadata}, + enable_spec_decode=True, + ) + + torch.testing.assert_close(metadata.kv_lens_cuda, torch.tensor([3])) + assert metadata.update_count == 1 + assert runner.graphs == {} + assert runner.graph_outputs == {} + assert runner.graph_metadata == {} + + +def test_decoder_warmup_publishes_metadata_for_capture() -> None: + runner = _make_decoder_runner() + runner.is_warmup_only = True + metadata = _AttentionMetadata() + spec_metadata = object() + output = torch.tensor([1]) + key = KeyType(1, 0, False, False, False) + + result = runner.capture( + key, + lambda _inputs: output, + { + "attn_metadata": metadata, + "spec_metadata": spec_metadata, + }, + ) + + assert result is output + assert runner.graph_metadata[key] == { + "attn_metadata": metadata, + "spec_metadata": spec_metadata, + } + assert runner.graphs == {} + assert runner.graph_outputs == {} + + +def test_decoder_capture_failure_resets_graph_without_publishing_metadata() -> None: + runner = _make_decoder_runner() + runner.WARMUP_STEPS = 0 + graph = _RecordingGraph() + + def fail_capture(_inputs: dict[str, object]) -> None: + raise ValueError("capture failed") + + with ( + patch("torch.cuda.CUDAGraph", return_value=graph), + patch("torch.cuda.graph", return_value=nullcontext()), + pytest.raises(ValueError, match="capture failed"), + ): + runner.capture( + (1, 0, False, False, False), + fail_capture, + {"attn_metadata": object()}, + ) + + assert graph.reset_count == 1 + assert runner.graphs == {} + assert runner.graph_outputs == {} + assert runner.graph_metadata == {} + + +def test_encoder_rejects_nested_output_without_orphaning_graph() -> None: + runner = _make_encoder_runner() + runner.WARMUP_STEPS = 0 + runner.shared_static_tensors = { + "input_ids": torch.empty(1), + "position_ids": torch.empty((1, 1)), + } + runner.shared_static_tensors_cpu = runner.shared_static_tensors + runner._arange_max = torch.arange(1, dtype=torch.int32) + runner._capture_h2d_copy = False + graph = _RecordingGraph() + nested_output = torch.nested.nested_tensor([torch.tensor([1.0]), torch.tensor([1.0, 2.0])]) + + with ( + patch("torch.cuda.CUDAGraph", return_value=graph), + patch("torch.cuda.graph", return_value=nullcontext()), + pytest.raises(TypeError, match="nested tensor outputs"), + ): + runner.capture( + (1, 1, 1), + lambda _inputs: nested_output, + { + "attn_metadata": _AttentionMetadata(), + "input_ids": [1], + "seq_lens": [1], + }, + ) + + assert graph.reset_count == 1 + assert runner.graphs == {} + assert runner.graph_outputs == {} + assert runner.graph_metadata == {} + + +def test_encoder_warmup_publishes_metadata_for_capture() -> None: + runner = _make_encoder_runner() + runner.is_warmup_only = True + runner.shared_static_tensors = { + "input_ids": torch.empty(1), + "position_ids": torch.empty((1, 1)), + } + runner.shared_static_tensors_cpu = runner.shared_static_tensors + runner._arange_max = torch.arange(1, dtype=torch.int32) + runner._capture_h2d_copy = False + metadata = _AttentionMetadata() + output = torch.tensor([1]) + key = (1, 1, 1) + + with patch("torch.cuda.current_stream"): + result = runner.capture( + key, + lambda _inputs: output, + { + "attn_metadata": metadata, + "input_ids": [1], + "seq_lens": [1], + }, + ) + + assert result is output + assert runner.graph_metadata[key] == {"attn_metadata": metadata} + assert runner.graphs == {} + assert runner.graph_outputs == {}