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

[KleinBlueC]

What breaks

CUDAGraphRunner.clear() erases the CUDA graph's private memory pool while the
C++ FusedMoeRunner still holds a workspace tensor allocated from it. The next
MoERunner.clear_all_workspaces() frees that tensor against a pool that no
longer exists, and the PyTorch caching allocator segfaults:

Segmentation fault, exit 139, on every TP rank but rank 0, during warmup
after model, weights and KV cache are up.

  std::_Rb_tree_decrement(std::_Rb_tree_node_base*)
  CUDACachingAllocator::Native::DeviceCachingAllocator::free_block
  c10::TensorImpl::~TensorImpl()
  cpp/tensorrt_llm/thop/moeOp.cpp:979   ~WorkspaceInfo
  cpp/tensorrt_llm/thop/moeOp.cpp:363   clearWorkspaces()

This is memory corruption, not exhaustion — running out of memory raises
torch.OutOfMemoryError; it does not walk the allocator's own free-block tree
off a cliff.

Why

1. moeOp.cpp getWorkspaceInfo() allocates the workspace unconditionally when
   isCapturing(stream), so the graph can replay against a stable buffer.
   PyTorch routes allocations on the capturing stream into the graph-private
   mempool, so this tensor is a private-pool block. moeOp.cpp is the only unit
   under cpp/tensorrt_llm/thop/ that allocates under isCapturing.

2. It is cached in mStreamWorkspaces, keyed by the capture stream, on a
   runner held in MoERunner.runner_dict — a class attribute, so
   process-lifetime. The runner outlives the executor that captured the graph.

3. With KV-cache-size estimation on (the default), an executor is built, warmed
   up (capturing graphs), then torn down before the real one is built. Shutdown
   reaches CUDAGraphRunner.clear(), which resets the graphs, drops the pool
   handle and calls empty_cache() — where release_cached_blocks() erases the
   PrivatePool.

4. The next warmup calls MoERunner.clear_all_workspaces(), which frees the
   step-1 tensor against the pool erased in step 3.

The deadline is empty_cache(), not graph.reset(): reset() only drops the
pool's use_count. This patch releases the workspaces between the two — after
the graphs are gone, before the pool is erased. See Side effects for why both
bounds matter.

Why Inkling hits it and other MoE models do not

The hazard is in the shared path, not in Inkling. tp=4 / ep=4 is the most
common layout in the accuracy suite — 20 configs across 10 model classes
(DeepSeek-V3Lite / V4Flash / R1, GPT-OSS, NemotronV3, Qwen3, Step3 …), all
capturing CUDA graphs successfully. So neither the config nor the CUDA driver is
the differentiator.

What varies is whether step 3 actually erases the pool, which happens only when
its segments are free at that instant — an allocator-layout question. Both knobs
that appear to control this bug are layout knobs: expert-parallel size moves the
MoE blocks around, and BF16 vs NVFP4 changes the weight footprint from ~40 GiB
to ~127 GiB per GPU. Inkling's layouts reach the erasing state; the other ten
models' evidently do not.

To be clear about the limit of that claim: this is characterised, not fully
explained. The hazard is demonstrably in the shared path and demonstrably
layout-sensitive, but I cannot give a segment-level account of why the other
models' blocks stay put. The fix does not depend on that answer — it removes the
window regardless of which layouts would have landed in it.

Validation

Reproduced on Inkling NVFP4 at tp=4 / ep=4 with CUDA graphs on, which
segfaults reliably, and clean with the patch. Moving the same call to after
empty_cache() crashes again, which is what rules out the fix working by merely
perturbing allocator layout rather than by closing the window.

Related: this is the crash that
PR #17062 currently guards
against at load time in InklingForCausalLM._assert_inkling_moe_parallel
("moe_tp_size=1 … segfaults during CUDA-graph capture"). That guard was added
because the cause was unknown; it can be removed once this lands.

Side effects

Placement is constrained on both sides. The release sits between
graph.reset() and empty_cache(), and neither bound is arbitrary. Before
graph.reset() is unsafe: captured graphs hold raw pointers into this
workspace, and PyExecutor.shutdown() already documents that freeing such a
buffer ahead of graph teardown risks a device-wide cudaErrorIllegalAddress.
After empty_cache() is too late, because that is where the pool is erased.

The call is global; the method is not. MoERunner.clear_all_workspaces()
walks MoERunner.runner_dict, a class attribute, so it clears every
FusedMoeRunner in the process — not only the ones whose workspaces came from
the pool being erased.

That has a concrete consequence in this codebase. PyExecutor.shutdown() tears
engines down in a loop:

for engine in (self.model_engine, self.draft_model_engine):
    engine._release_cuda_graphs()

Each engine owns its own CUDAGraphRunner, graphs and pool. When the target
engine's clear() runs, it frees the draft engine's MoE workspace too — while
the draft engine's captured graphs are still alive and still hold raw pointers
into it. That is the same hazard the placement above is careful to avoid within
one runner, reintroduced across runners.

I have not observed this fire; it needs speculative decoding with a draft model
that captures graphs and runs MoE. But it is the same mechanism as a documented
one, so it should be weighed rather than assumed benign. The clean form of this
fix releases only the workspaces belonging to the pool being erased, which needs
a per-stream or per-pool entry point on the C++ side (moeOp.cpp currently
exposes only clear_workspaces(), all-or-nothing). This PR trades a narrower,
unobserved window for a reproducible crash; whether that trade is acceptable, or
whether the C++ API should be extended first, is a reviewer call.

Cost. One workspace re-allocation per stream on next MoE use, on a teardown
path that already calls empty_cache() — nothing on a steady-state path.
clear() is only reached from PyExecutor.shutdown() and the model engine's
cleanup, both verified teardown-only.

Numerics. None. The workspace is scratch; getWorkspaceInfo() already
reallocates it whenever capture is active or the required size grows.

No MoE. runner_dict is empty and the call is a no-op.

Not touched. EncoderCUDAGraphRunner.clear() has the same shape, but
whether an encoder tower ever allocates a MoE workspace is unconfirmed, so this
does not add a call there on the strength of the pattern alone.

Not measured. No performance or accuracy regression run on a configuration
that was already healthy. The reasoning above says there should be none.

Alternatives considered

This patch is the smallest change that closes the window. Two other shapes were
considered and are worth stating, because the side effect above is a property of
this shape rather than of the fix.

Split clear() into destroy_graphs() and release_pool(). Python only,
no C++ change. PyExecutor.shutdown() would then destroy every engine's graphs,
release the workspaces once, and only then erase the pools:

for engine in (...):  engine.destroy_graphs()
MoERunner.clear_all_workspaces()          # no graph anywhere holds a pointer,
for engine in (...):  engine.release_pool()   # and no pool is erased yet

That places the global call at the one instant when it is safe for every
engine, so the cross-engine window disappears. Cost is a wider API — one method
becomes three — and edits to PyExecutor.shutdown() and the model engine's
cleanup path. Arguably it is also the more honest decomposition: clear()
currently couples two things with different lifetimes, which is what created
the window in the first place.

Per-pool release on the C++ side. Add something like
clear_workspaces_for_stream() to moeOp.cpp, so the release granularity
matches the pool granularity and clear() frees only what it is erasing.
Semantically the cleanest, and it removes the global call entirely. Cost is a
C++ change plus a rebuild, and CUDAGraphRunner would have to know its capture
streams.

I went with the minimal version here on the grounds that it fixes a
reproducible crash with one call and no API change. If reviewers would rather
not take the cross-engine window, the split above is the cheaper of the two
alternatives and needs no C++ work.

Test

tests/unittest/_torch/executor/test_cuda_graph_runner_clear.py asserts that
clear() releases the MoE workspaces before empty_cache(), by mocking both
and checking their order. No GPU, no model, no CUDA, so it runs in any lane.
Triggering the fault needs the real allocator layout — in pure torch a held
block keeps the pool's cudaMalloc_count above zero, so the pool is never
erased and a stale handle and an erased pool cannot coexist. The ordering is
what regresses, and the ordering is checkable anywhere.

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.