diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 07618d87a1a9..e8147b2989ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -847,6 +847,26 @@ def clear(self): """Releases all captured graphs and the associated memory pool.""" for graph in self.graphs.values(): graph.reset() + # Release the C++ MoE workspaces here: after the graphs are destroyed, + # before empty_cache() erases the pool that backs them. + # + # FusedMoeRunner allocates this workspace under isCapturing (moeOp.cpp + # getWorkspaceInfo, the only thop unit that does), so it is a block of + # this graph's private pool, and caches it in mStreamWorkspaces keyed by + # the capture stream. That map hangs off MoERunner.runner_dict, a class + # attribute, so it outlives the executor that captured the graph: with + # KV-cache-size estimation on the estimation executor is torn down, and + # the next warmup's clear_all_workspaces() then frees the workspace + # against an erased pool, segfaulting in the allocator's free_block(). + # + # Both bounds are load-bearing. After graph.reset(), because captured + # graphs hold raw pointers into this workspace and freeing it first + # risks the device-wide cudaErrorIllegalAddress that + # PyExecutor.shutdown() documents for exactly this ordering. Before + # empty_cache(), which is where release_cached_blocks() erases the + # PrivatePool; graph.reset() only drops its use_count. + from ..custom_ops.torch_custom_ops import MoERunner + MoERunner.clear_all_workspaces() self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() diff --git a/tests/unittest/_torch/executor/test_cuda_graph_pool_hazard.py b/tests/unittest/_torch/executor/test_cuda_graph_pool_hazard.py new file mode 100644 index 000000000000..902a7670d799 --- /dev/null +++ b/tests/unittest/_torch/executor/test_cuda_graph_pool_hazard.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Reproduces the hazard that CUDAGraphRunner.clear() must avoid. + +`clear()` releases the C++ MoE workspaces before `torch.cuda.empty_cache()` +erases the CUDA graph's private memory pool. This test shows why that ordering +matters, without a model, without MoE and without TRT-LLM in the picture: an +allocation made during graph capture and still held when the pool is torn down +faults the caching allocator when it is finally freed. + +The reproduction needs two things together, and neither alone is enough: + + expandable_segments pre-expansion result + ------------------- ------------- ------ + on yes SIGSEGV + on no clean + off yes clean + off no clean + +`expandable_segments:True` is also a measured necessary condition on the real +crash this guards against: the same Inkling NVFP4 run that segfaults reliably +runs clean with that allocator mode switched off and no other change. + +A SIGSEGV cannot be caught in-process, so each arm runs in a fresh subprocess +and the assertion is on its exit status. Fresh matters -- an already initialised +CUDA context changes the allocator layout that is under test. +""" + +import os +import subprocess +import sys +import textwrap + +import pytest +import torch + +# One large allocation that grows the segment once, then freed so everything +# below is carved out of that space without growing it again. +_PRE_EXPAND_MIB = 256 +_N_BLOCKS = 64 +_BLOCK_BYTES = 2 << 20 + +_SCRIPT = textwrap.dedent(""" + import gc + import sys + + import torch + + RELEASE_EARLY = sys.argv[1] == "release_early" + + x = torch.zeros(1024, device="cuda") + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + blocks = [] + with torch.cuda.graph(graph): + pre = torch.empty({pre_expand} << 20, dtype=torch.int8, device="cuda") + del pre + for _ in range({n_blocks}): + blocks.append(torch.empty({block_bytes}, dtype=torch.int8, device="cuda")) + x.add_(1.0) + torch.cuda.synchronize() + + # Hold a scattered subset, so the freed blocks between them cannot coalesce. + # This stands in for the C++ FusedMoeRunner holding its workspace tensor. + held = blocks[::2] + del blocks + gc.collect() + + if RELEASE_EARLY: + # What the fix does: drop the held allocations while the pool still + # exists, before empty_cache() tears it down. + del held + held = None + gc.collect() + + # Exactly CUDAGraphRunner.clear(): reset the graphs, drop the pool handle, + # then empty_cache(). + del x + graph.reset() + del graph + gc.collect() + torch.cuda.empty_cache() + + if held is not None: + del held + gc.collect() + torch.cuda.empty_cache() + + print("survived") +""").format(pre_expand=_PRE_EXPAND_MIB, + n_blocks=_N_BLOCKS, + block_bytes=_BLOCK_BYTES) + + +def _run(arm: str, expandable: bool) -> int: + env = {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"} if expandable else {} + return subprocess.run([sys.executable, "-c", _SCRIPT, arm], + env={**os.environ, **env}, + capture_output=True, + timeout=300).returncode + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_releasing_before_pool_teardown_avoids_the_fault(): + """The ordering clear() implements: release first, and the teardown is safe.""" + assert _run("release_early", expandable=True) == 0 + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_holding_across_pool_teardown_faults(): + """Holding a capture-pool allocation past the teardown crashes the allocator. + + This is the hazard clear() closes. If a future PyTorch release fixes the + allocator, this assertion is the thing that will tell us -- at which point + the early release becomes belt-and-braces rather than load-bearing, and this + test should be revisited rather than silently relaxed. + """ + rc = _run("hold", expandable=True) + assert rc != 0, ("expected the held allocation to fault the allocator; " + "if this now passes cleanly, the PyTorch allocator behaviour " + "has changed and CUDAGraphRunner.clear() should be re-examined") + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_fault_requires_expandable_segments(): + """Control: the identical arm is clean without expandable_segments. + + Pins down which allocator mode the hazard belongs to, and keeps the test + above from being read as "holding an allocation is always fatal". + """ + assert _run("hold", expandable=False) == 0 diff --git a/tests/unittest/_torch/executor/test_cuda_graph_runner_clear.py b/tests/unittest/_torch/executor/test_cuda_graph_runner_clear.py new file mode 100644 index 000000000000..1071a5c7974c --- /dev/null +++ b/tests/unittest/_torch/executor/test_cuda_graph_runner_clear.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""CUDAGraphRunner.clear() must release the C++ MoE workspaces first. + +Guards a SIGSEGV that needs no GPU to prevent but cannot be reproduced without +one: FusedMoeRunner allocates its workspace under isCapturing (moeOp.cpp +getWorkspaceInfo -- the only thop unit that does), so the tensor comes from the +graph's private memory pool, and caches it in mStreamWorkspaces. That map hangs +off MoERunner.runner_dict, a class attribute, so it outlives the executor that +captured the graph. With KV-cache-size estimation on, that executor is torn +down between the two warmups; if the workspace is still held when +empty_cache() erases the pool, the next clear_all_workspaces() frees it against +a dead pool and faults inside the caching allocator's free_block(). + +Reproducing the fault needs the real allocator layout -- in pure torch a held +block keeps cudaMalloc_count above zero, so the pool is never erased and the +crash cannot be staged. The ordering is checkable without any of that, which is +what this does. +""" + +from unittest import mock + +from tensorrt_llm._torch.pyexecutor import cuda_graph_runner as cgr + + +class _FakeGraph: + def __init__(self, order): + self._order = order + + def reset(self): + self._order.append("graph.reset") + + +def test_clear_releases_moe_workspaces_before_erasing_the_pool(): + order = [] + + runner = object.__new__(cgr.CUDAGraphRunner) + runner.graphs = {("k",): _FakeGraph(order)} + runner.graph_outputs = {} + runner.graph_metadata = {} + runner.padding_dummy_requests = {} + runner.memory_pool = object() + + with mock.patch( + "tensorrt_llm._torch.custom_ops.torch_custom_ops.MoERunner.clear_all_workspaces", + side_effect=lambda: order.append("clear_all_workspaces"), + ): + with mock.patch("torch.cuda.empty_cache", side_effect=lambda: order.append("empty_cache")): + runner.clear() + + assert "clear_all_workspaces" in order, ( + "clear() never released the C++ MoE workspaces; whoever frees them " + "next does it against a pool this method has already erased" + ) + # empty_cache() is the deadline: that is where release_cached_blocks() + # erases the PrivatePool. graph.reset() only drops its use_count, so + # landing between reset() and empty_cache() would also be correct. + assert order.index("clear_all_workspaces") < order.index("empty_cache"), ( + f"MoE workspaces must be released before the pool is erased; got {order}" + )