[None][fix] Release MoE workspaces before erasing the CUDA graph pool

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()

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

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}"
+    )