From 4fff951957d96c4d903b6fc98397d89a4b45c2cc Mon Sep 17 00:00:00 2001 From: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com> Date: Thu, 6 Aug 2026 22:48:56 -0700 Subject: [PATCH 1/3] [None][fix] Release MoE workspaces before erasing the CUDA graph pool 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 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, that executor is torn down between the two warmups; the workspace is still held when clear() erases the pool, and the next clear_all_workspaces() frees it against a pool that no longer exists -- SIGSEGV inside the caching allocator: _Rb_tree_decrement -> free_block -> local_raw_delete -> ~TensorImpl -> moeOp.cpp:979 (WorkspaceInfo) -> moeOp.cpp:363 (clearWorkspaces) Releasing the workspaces at the top of clear() keeps them inside the pool's lifetime. Blast radius is exactly FusedMoeRunner: moeOp.cpp is the only thop translation unit that allocates under isCapturing and caches the result. Established on the Inkling NVFP4 tp=4/ep=4 + CUDA graph reproducer, five arms on one tree with the call site chosen by env: no call SIGSEGV before graph.reset() clean after graph.reset() clean no call, CUDA graph off clean after empty_cache() SIGSEGV The last arm is why this is a cause and not a coincidence: it makes the same call in the same process and still crashes, so the fix does not work by perturbing allocator layout. It also locates the deadline precisely -- both sides of graph.reset() are clean, and only crossing empty_cache() faults, so the boundary is where release_cached_blocks() erases the PrivatePool, not reset(), which merely drops its use_count. The accompanying test asserts the ordering only. The fault itself cannot be staged without the real allocator layout: in pure torch a held block keeps cudaMalloc_count above zero, so the pool is never erased and a stale handle and an erased pool cannot coexist. Ordering is checkable with no GPU at all, so the guard runs in any lane. Not Inkling-specific: tp=4/ep=4 is the most common layout in the accuracy suite (20 configs across 10 model classes). Those models simply do not land the workspace where erasure catches it; the hazard is in the shared path. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 16 +++++ .../executor/test_cuda_graph_runner_clear.py | 72 +++++++++++++++++++ 2 files changed, 88 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_cuda_graph_runner_clear.py diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 07618d87a1a9..7de0f58b7bea 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -845,6 +845,22 @@ def pad_batch(self, def clear(self): """Releases all captured graphs and the associated memory pool.""" + # Release the C++ MoE workspaces before the pool that backs them is + # erased. FusedMoeRunner allocates a workspace under isCapturing + # (moeOp.cpp getWorkspaceInfo, the only thop unit that does), so it + # comes from 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() would free the workspace against a pool that + # no longer exists, segfaulting in the caching allocator's free_block(). + # + # The deadline is the empty_cache() below, which is where + # release_cached_blocks() actually erases the PrivatePool; graph.reset() + # only drops its use_count. + from ..custom_ops.torch_custom_ops import MoERunner + MoERunner.clear_all_workspaces() for graph in self.graphs.values(): graph.reset() self.graphs.clear() 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}" + ) From 615a44d00173a69e579d8712b9e2aa96913a88d6 Mon Sep 17 00:00:00 2001 From: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com> Date: Fri, 7 Aug 2026 00:09:20 -0700 Subject: [PATCH 2/3] [None][fix] Move the MoE workspace release after graph teardown The release was at the top of clear(), which is the wrong side of a hazard the codebase already documents: PyExecutor.shutdown() warns that freeing a GPU workspace referenced by raw pointers inside captured CUDA graphs, ahead of the graph teardown, can trigger a device-wide cudaErrorIllegalAddress. The MoE workspace is exactly such a buffer -- getWorkspaceInfo() allocates it under isCapturing so the graph can replay against a stable address. Both bounds are load-bearing, and the reproducer showed the second one directly: placing the call after graph.reset() is clean, and only moving it past empty_cache() crashes again. So the window is between the two -- after no graph references the workspace, before release_cached_blocks() erases the pool. Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 34 +++++++++++-------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 7de0f58b7bea..e8147b2989ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -845,24 +845,28 @@ def pad_batch(self, def clear(self): """Releases all captured graphs and the associated memory pool.""" - # Release the C++ MoE workspaces before the pool that backs them is - # erased. FusedMoeRunner allocates a workspace under isCapturing - # (moeOp.cpp getWorkspaceInfo, the only thop unit that does), so it - # comes from 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() would free the workspace against a pool that - # no longer exists, segfaulting in the caching allocator's free_block(). + 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(). # - # The deadline is the empty_cache() below, which is where - # release_cached_blocks() actually erases the PrivatePool; graph.reset() - # only drops its use_count. + # 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() - for graph in self.graphs.values(): - graph.reset() self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() From e9625544399e86cca99588df04fdb56d42531a2e Mon Sep 17 00:00:00 2001 From: kleinc Date: Fri, 7 Aug 2026 12:50:36 -0700 Subject: [PATCH 3/3] [None][test] Add a model-free reproducer for the CUDA graph pool hazard Allocations made during CUDA graph capture and still held when the pool is torn down fault the caching allocator when they are finally freed. That is the hazard CUDAGraphRunner.clear() closes by releasing the MoE workspaces before empty_cache(), and until now it could only be observed on a full Inkling NVFP4 run at tp=4/ep=4. The reproduction needs two ingredients together, and the test pins both: expandable_segments pre-expansion result on yes SIGSEGV on no clean off yes clean off no clean expandable_segments:True being required matches the necessary condition measured on the real crash, which is the evidence that the two are the same hazard. Under CUDA_LAUNCH_BLOCKING=1 the backtrace holds no kernel frame at all: the fault is a host-side dereference of a neighbour Block* inside try_merge_blocks, reached from a tensor destructor, which is the same allocator path the Inkling stack enters from ~WorkspaceInfo. 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 under test. Signed-off-by: kleinc --- .../executor/test_cuda_graph_pool_hazard.py | 147 ++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 tests/unittest/_torch/executor/test_cuda_graph_pool_hazard.py 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