Skip to content

[https://nvbugs/6581063][fix] Keep the verified fix (geometry-only cache key + deterministic eviction from… - #17529

Closed
trtllm-agent wants to merge 1 commit into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6581063
Closed

[https://nvbugs/6581063][fix] Keep the verified fix (geometry-only cache key + deterministic eviction from…#17529
trtllm-agent wants to merge 1 commit into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6581063

Conversation

@trtllm-agent

@trtllm-agent trtllm-agent commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Root cause: _MEGA_MOE_SYMM_BUFFER_CACHE was keyed on id(self._ep_pg) and never evicted, so NVLink symmetric-memory workspaces (allocated via empty_strided_p2p, outside the caching allocator, so empty_cache() cannot reclaim them) outlived their LLM in a reused MPI worker and stacked a second allocation per LLM until the driver OOM'd.
  • Fix: Keep the verified fix (geometry-only cache key + deterministic eviction from worker.shutdown() before destroy_process_group()), and close the stale-group-hit window geometry-only keying opened by re-validating each cache hit against the caller's live _ep_pg and freeing the stale buffer in place via a shared _free_symm_buffer helper.
  • Original test: pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm -v
  • Automated fix generated by repair-bot

Test plan

  • Verify fix on the same GPU type as the original failure
  • Check for regressions in related tests

Links

Dev Engineer Review

  • The cache now uses buffer geometry and EP size instead of _ep_pg identity. This enables reuse across LLM instances.
  • Cache hits validate the caller’s live _ep_pg. Stale buffers are freed through _free_symm_buffer.
  • worker.shutdown() releases cached buffers before it destroys process groups.
  • The new release_symm_buffer_cache() API releases cached buffers and logs reclaimed memory.
  • The changes address NVLink symmetric-memory leaks in reused MPI workers.
  • The implementation should be checked against CODING_GUIDELINES.md and validated with the related regression tests.

QA Engineer Review

  • Modified test-list file: tests/integration/test_lists/waives.txt.
  • Removed the waiver for TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm.
  • No test functions were added, modified, or removed.
  • The test is now eligible to run instead of being skipped.
  • Verdict: needs follow-up because CBTS coverage data is unavailable.

@coderabbitai

coderabbitai Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 974c39c8-b444-411e-bb88-687963226b67

📥 Commits

Reviewing files that changed from the base of the PR and between 9997d3f and 965bc1f.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
  • tensorrt_llm/executor/worker.py
  • tests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt
🚧 Files skipped from review as they are similar to previous changes (2)
  • tensorrt_llm/executor/worker.py
  • tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py

Included review availability: Your plan includes up to 12 reviews per rolling hour; 11 remain after this review.


Walkthrough

The change adds SymmBuffer cache cleanup helpers, keys cached buffers by geometry and EP size, evicts buffers linked to different live EP groups, and releases the cache during worker shutdown. The related DeepGEMM integration test waiver is removed.

Changes

MegaMoE buffer lifecycle

Layer / File(s) Summary
SymmBuffer cache cleanup and validation
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
SymmBuffer helpers clear tensor views, free symmetric memory, detach cache entries, and log released buffers. Cache keys use geometry and EP size. Entries from a different live EP group are evicted and freed.
Worker shutdown integration
tensorrt_llm/executor/worker.py, tests/integration/test_lists/waives.txt
Worker shutdown releases the loaded DeepGEMM symmetric-memory cache before distributed process-group destruction. The DeepGEMM integration test waiver is removed.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 965bc

The PR adds targeted validation and deterministic cleanup for the workspace cache; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: schetlur-nv, asfiyab-nvidia, dpitman-nvda, junyixu-nv, niukuo

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the NVBugs fix and its main changes: geometry-only cache keys and deterministic eviction.
Description check ✅ Passed The description explains the root cause, fix, testing, and bug link, but it omits the template's PR Checklist section.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py (1)

62-79: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use a structural SymmBuffer protocol for the cache.

The repository does not expose a concrete DeepGEMM SymmBuffer type. Define a private protocol for buffer, group, and destroy(), and use it for the cache and _free_symm_buffer(). Replace Dict[tuple, object] with built-in generic types.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py` around
lines 62 - 79, Define a private structural SymmBuffer protocol exposing buffer,
group, and destroy(), then use that protocol for _MEGA_MOE_SYMM_BUFFER_CACHE and
the buffered parameter of _free_symm_buffer(). Replace Dict[tuple, object] with
the appropriate built-in generic annotation while preserving the existing
cleanup behavior.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/executor/worker.py`:
- Around line 127-135: Update shutdown() so the
mega_moe.release_symm_buffer_cache() call is wrapped in a try/finally structure,
ensuring destroy_process_group() and subsequent CUDA cleanup always execute even
when cache release raises. Preserve the existing sys.modules lookup and shutdown
guard behavior.

---

Nitpick comments:
In `@tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py`:
- Around line 62-79: Define a private structural SymmBuffer protocol exposing
buffer, group, and destroy(), then use that protocol for
_MEGA_MOE_SYMM_BUFFER_CACHE and the buffered parameter of _free_symm_buffer().
Replace Dict[tuple, object] with the appropriate built-in generic annotation
while preserving the existing cleanup behavior.
🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 15db361a-175c-48e4-bf0e-c00d33688215

📥 Commits

Reviewing files that changed from the base of the PR and between 43c2386 and 5092ac6.

📒 Files selected for processing (3)
  • tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
  • tensorrt_llm/executor/worker.py
  • tests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
  • tests/integration/test_lists/waives.txt

Comment on lines +127 to +135
# MegaMoE's NVLink symmetric-memory activation workspaces are
# rendezvoused over the EP group, so they must go before the
# destroy_process_group() below. Probed via sys.modules so a
# non-MegaMoE run does not import the MoE stack to release an
# empty cache.
mega_moe = sys.modules.get(
"tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm")
if mega_moe is not None:
mega_moe.release_symm_buffer_cache()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect every cache-release call and the shutdown control flow.
rg -n -C 12 --glob '*.py' \
  'release_symm_buffer_cache|_free_symm_buffer|destroy_process_group|doing_shutdown' \
  tensorrt_llm

# Locate DeepGEMM SymmBuffer destruction implementations and exception behavior.
rg -n -C 8 --glob '*.py' --glob '*.pyi' \
  'class .*SymmBuffer|def destroy\(' .

Repository: NVIDIA/TensorRT-LLM

Length of output: 50375


🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- worker shutdown ---'
sed -n '1,180p' tensorrt_llm/executor/worker.py

printf '%s\n' '--- cache helper ---'
sed -n '1,125p' tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
rg -n -C 12 'def release_symm_buffer_cache|def _free_symm_buffer|SymmBuffer|destroy\(' \
  tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 35783


🏁 Script executed:

#!/bin/bash
set -euo pipefail

python3 - <<'PY'
import ast
from pathlib import Path

worker = ast.parse(Path("tensorrt_llm/executor/worker.py").read_text())
shutdown = next(
    node for node in ast.walk(worker)
    if isinstance(node, ast.FunctionDef) and node.name == "shutdown"
)

calls = []
for node in ast.walk(shutdown):
    if isinstance(node, ast.Call):
        func = node.func
        if isinstance(func, ast.Attribute):
            name = func.attr
        elif isinstance(func, ast.Name):
            name = func.id
        else:
            name = ast.dump(func)
        calls.append((node.lineno, name))

for line, name in sorted(calls):
    if name in {
        "release_symm_buffer_cache",
        "destroy_process_group",
        "empty_cache",
        "collect",
    }:
        print(f"{line}: {name}")

print("shutdown_try_finally_blocks:",
      sum(isinstance(node, ast.Try) and node.finalbody
          for node in ast.walk(shutdown)))

helper = ast.parse(
    Path("tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py")
    .read_text()
)
release = next(
    node for node in ast.walk(helper)
    if isinstance(node, ast.FunctionDef)
    and node.name == "release_symm_buffer_cache"
)
print("release_try_finally_blocks:",
      sum(isinstance(node, ast.Try) and node.finalbody
          for node in ast.walk(release)))
print("release_has_per_buffer_exception_handler:",
      any(isinstance(node, ast.Try) and node.handlers
          for node in ast.walk(release)))
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 354


Preserve worker teardown when cache release fails.

If release_symm_buffer_cache() raises, shutdown() skips destroy_process_group() and CUDA cleanup. Since doing_shutdown is already True, later calls do not retry. Run the remaining teardown in a finally block.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/executor/worker.py` around lines 127 - 135, Update shutdown() so
the mega_moe.release_symm_buffer_cache() call is wrapped in a try/finally
structure, ensuring destroy_process_group() and subsequent CUDA cleanup always
execute even when cache release raises. Preserve the existing sys.modules lookup
and shutdown guard behavior.

@brnguyen2 brnguyen2 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two things to settle before this merges.

Waiver removal. The reported failure is order-dependent — it needs the reused worker pool carrying state from the tests that ran before it. A standalone pytest ...::test_gsm8k_dep8_megamoe_deepgemm pass doesn't exercise that, so it isn't evidence the waiver can come out. Please run the post-merge stage that contains this test with the waiver removed (/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-1") and paste the result; if the cross-test residue turns out to have other contributors, this fix can land on its own and the waiver removal split off.

Tests. _free_symm_buffer and release_symm_buffer_cache are pure Python — a unit test with a stub object (tensor attrs + a destroy() that nulls only a couple of them) would pin both the "sweep every remaining tensor view" behavior and the cache-emptied-on-release contract without a GPU. That's cheap insurance for logic whose failure mode is an invisible leak.

self.activation,
)
cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key)
if cached is not None and cached.group is not self._ep_pg:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes correctness depend on DeepGEMM storing the process group you passed verbatim as SymmBuffer.group. If it ever normalizes it (wraps it, resolves dist.group.WORLD to the default PG object, stores a group name), the identity check is permanently true and layer 1 frees the buffer layer 0 already assigned to self._symm_buffer — a use-after-free on symmetric memory, much worse than the leak being fixed.

Don't depend on someone else's attribute: record the group alongside the buffer in the cache and compare against your own record.

cached_pg, cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key, (None, None))
if cached is not None and cached_pg is not self._ep_pg:
    ...
_MEGA_MOE_SYMM_BUFFER_CACHE[key] = (self._ep_pg, cached)

Separately, the comment asserts the mismatching group is "already destroyed" — nothing checks that. Two LLMs alive at once in one process with the same geometry and different EP groups (ep_size equal) hit this path and the second one frees the first one's live workspace. If that's considered out of contract, say so explicitly here; if not, gate the free on the old group actually being torn down.

# buffer in the cache for a later lookup to trip over.
buffers = list(_MEGA_MOE_SYMM_BUFFER_CACHE.values())
_MEGA_MOE_SYMM_BUFFER_CACHE.clear()
total_bytes = sum(_free_symm_buffer(buffered) for buffered in buffers)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sum(...) over a generator stops at the first exception, so one buffer whose destroy() raises leaves every later buffer in the list unfreed — and they've already been dropped from the cache, so nothing can ever reclaim them. Given the whole point is deterministic release, loop and keep going:

total_bytes = 0
for buffered in buffers:
    try:
        total_bytes += _free_symm_buffer(buffered)
    except Exception:
        logger.exception("[MegaMoE] failed to release a DG SymmBuffer")

# rendezvoused over the EP group, so they must go before the
# destroy_process_group() below. Probed via sys.modules so a
# non-MegaMoE run does not import the MoE stack to release an
# empty cache.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two gaps with the sys.modules probe:

  1. The module path is a bare string with nothing tying it to the real module. If the file moves, this silently degrades to a no-op and the leak comes back with no signal. A lazy import inside a narrow try: ... except ImportError would at least fail loudly at the right time, or add a test that asserts the string resolves.
  2. This is the MPI worker path only. As the comment right below notes, the Ray path owns its process group in RayWorkerWrapper and doesn't run through here, so a Ray-launched MegaMoE run still leaks the workspace. Worth either handling it there too or noting the gap explicitly.

@leslie-fang25
leslie-fang25 removed their request for review August 12, 2026 02:14

@BowenFu BowenFu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes required: cache eviction is not yet safe or verified.

  • The cache compares external SymmBuffer.group by identity and frees on any mismatch. Two live same-geometry buffers with different process groups can therefore free memory still used by the first model. Track ownership explicitly and free only after the old group is proven retired, or enforce the exclusive-lifetime invariant.
  • sum(...) stops on the first destroy() exception after all entries were removed from the cache, leaking every later buffer permanently.
  • Worker shutdown must still destroy the process group and clean CUDA if cache release raises.
  • Add focused CPU regressions for ownership and exception paths, then run DGX_B200-8_GPUs-PyTorch-1 with the waiver removed.

These are required before the OOM waiver can be removed.

The DG SymmBuffer cache was keyed on id(self._ep_pg) and never evicted, so
NVLink symmetric-memory activation workspaces outlived the LLM that allocated
them. A worker process is reused across LLMs (a new executor is built per LLM
while the CUDA context persists) and the EP group is destroyed on executor
shutdown, so the next LLM computed a different key and allocated a second
buffer beside the first instead of reusing it. These come from
empty_strided_p2p, outside PyTorch's caching allocator, so the
torch.cuda.empty_cache() already on the worker shutdown path cannot reclaim
them; ~40 GiB of 178 GiB was gone before the failing test allocated anything
and _alloc_symm_buffer died with "CUDA driver error: out of memory". id() is
also recycled once a group is freed, so a new group could land on a dead
group's id and hit a buffer rendezvoused over a destroyed group.

Key the cache on buffer geometry so a later LLM in a reused worker reuses the
buffer, and release the cache from the worker's shutdown path, just before the
EP group these buffers were rendezvoused over is destroyed. Eviction is
deterministic rather than reachability-based on purpose: SymmBuffer.__init__
keeps a strong reference to its own group, so a weakly-held owner would stay
reachable through the cached buffer and never be collected. Releasing also
clears every remaining tensor attribute, because SymmBuffer.destroy() nulls
only a few of the views sliced out of the allocation and any survivor pins the
whole buffer.

Unwaive the test now that it passes.

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
@Barry-Delaney
Barry-Delaney force-pushed the repair-bot-bug6581063 branch from 5092ac6 to 965bc1f Compare August 17, 2026 06:26
@coderabbitai

coderabbitai Bot commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

@sunnyqgg

Copy link
Copy Markdown
Collaborator

Superseded by #18182, which keeps this PR's fix (geometry-only cache key + deterministic eviction from worker.shutdown()) and additionally addresses the review feedback: exception-safe teardown in shutdown() (a failing release can no longer skip destroy_process_group()), per-buffer release instead of sum() over a generator, stale-hit validation against the ProcessGroup recorded on the TRT-LLM side instead of SymmBuffer.group, a named probe-string constant asserted by a regression test, and CPU regression tests wired into l0_a10.yml. Closing in favor of #18182.

@sunnyqgg sunnyqgg closed this Aug 25, 2026
sunnyqgg added a commit to sunnyqgg/TensorRT-LLM that referenced this pull request Aug 25, 2026
…r teardown

The DG SymmBuffer cache was keyed on id(self._ep_pg) and never evicted, so
NVLink symmetric-memory activation workspaces outlived the LLM that allocated
them. A worker process is reused across LLMs (a new executor is built per LLM
while the CUDA context persists) and the EP group is destroyed on executor
shutdown, so the next LLM computed a different key and allocated a second
buffer beside the first instead of reusing it. These come from
empty_strided_p2p, outside PyTorch's caching allocator, so the
torch.cuda.empty_cache() already on the worker shutdown path cannot reclaim
them; ~40 GiB of 178 GiB was gone before the failing test allocated anything
and _alloc_symm_buffer died with "CUDA driver error: out of memory". id() is
also recycled once a group is freed, so a new group could land on a dead
group's id and hit a buffer rendezvoused over a destroyed group.

Key the cache on buffer geometry so a later LLM in a reused worker reuses the
buffer, and release the cache from the worker's shutdown path, just before the
EP group these buffers were rendezvoused over is destroyed. Eviction is
deterministic rather than reachability-based on purpose: SymmBuffer.__init__
keeps a strong reference to its own group, so a weakly-held owner would stay
reachable through the cached buffer and never be collected. Releasing also
clears every remaining tensor attribute, because SymmBuffer.destroy() nulls
only a few of the views sliced out of the allocation and any survivor pins the
whole buffer.

Supersedes PR NVIDIA#17529, additionally addressing its review feedback:
- worker.shutdown() no longer lets a failing release skip
  destroy_process_group() and the CUDA cleanup below it: doing_shutdown is
  already set on entry, so an escaping exception would leave NCCL
  communicators alive with no retry possible.
- release_symm_buffer_cache() frees buffer-by-buffer and keeps going on
  failure, instead of sum() over a generator that abandons the remaining
  buffers (already evicted, hence unreclaimable) on the first exception.
- Cache entries record the ProcessGroup TRT-LLM passed at allocation, and
  stale-hit validation checks that record instead of SymmBuffer.group, so a
  DeepGEMM bump that normalizes the stored group cannot turn every hit into
  a false stale that frees a live buffer another layer already holds.
- The sys.modules probe string lives in a named constant, and a regression
  test asserts it resolves, so a module move fails loudly instead of
  silently reverting to the leak.
- _free_symm_buffer() tolerates an already-destroyed buffer.
- Document that the Ray/RPC worker paths rely on per-LLM process exit for
  reclamation and must call the release if they ever reuse processes.
- Add CPU regression tests (TestMegaMoeSymmBufferRelease in the existing
  tests/unittest/executor/test_base_worker.py, already wired in l0_cpu and
  l0_a100) covering the cache release/sweep, continue-on-raising-destroy,
  stale-group eviction with recorded-group validation, probe-string
  resolution, and shutdown surviving a release failure.

Unwaive the test now that it passes.

Signed-off-by: qgai <qgai@nvidia.com>
sunnyqgg added a commit to sunnyqgg/TensorRT-LLM that referenced this pull request Aug 25, 2026
…r teardown

The DG SymmBuffer cache was keyed on id(self._ep_pg) and never evicted, so
NVLink symmetric-memory activation workspaces outlived the LLM that allocated
them. A worker process is reused across LLMs (a new executor is built per LLM
while the CUDA context persists) and the EP group is destroyed on executor
shutdown, so the next LLM computed a different key and allocated a second
buffer beside the first instead of reusing it. These come from
empty_strided_p2p, outside PyTorch's caching allocator, so the
torch.cuda.empty_cache() already on the worker shutdown path cannot reclaim
them; ~40 GiB of 178 GiB was gone before the failing test allocated anything
and _alloc_symm_buffer died with "CUDA driver error: out of memory". id() is
also recycled once a group is freed, so a new group could land on a dead
group's id and hit a buffer rendezvoused over a destroyed group.

Key the cache on buffer geometry so a later LLM in a reused worker reuses the
buffer, and release the cache from the worker's shutdown path, just before the
EP group these buffers were rendezvoused over is destroyed. Eviction is
deterministic rather than reachability-based on purpose: SymmBuffer.__init__
keeps a strong reference to its own group, so a weakly-held owner would stay
reachable through the cached buffer and never be collected. Releasing also
clears every remaining tensor attribute, because SymmBuffer.destroy() nulls
only a few of the views sliced out of the allocation and any survivor pins the
whole buffer.

Supersedes PR NVIDIA#17529, additionally addressing its review feedback:
- worker.shutdown() no longer lets a failing release skip
  destroy_process_group() and the CUDA cleanup below it: doing_shutdown is
  already set on entry, so an escaping exception would leave NCCL
  communicators alive with no retry possible.
- release_symm_buffer_cache() frees buffer-by-buffer and keeps going on
  failure, instead of sum() over a generator that abandons the remaining
  buffers (already evicted, hence unreclaimable) on the first exception.
- Cache entries record the ProcessGroup TRT-LLM passed at allocation, and
  stale-hit validation checks that record instead of SymmBuffer.group, so a
  DeepGEMM bump that normalizes the stored group cannot turn every hit into
  a false stale that frees a live buffer another layer already holds.
- The sys.modules probe string lives in a named constant next to the import
  block instead of being buried inline in shutdown().
- _free_symm_buffer() tolerates an already-destroyed buffer.
- Document that the Ray/RPC worker paths rely on per-LLM process exit for
  reclamation and must call the release if they ever reuse processes.

Unwaive the test now that it passes; the DGX_B200 run of
TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm scheduled after the
heavy predecessor sequence is the validation gate.

Signed-off-by: qgai <qgai@nvidia.com>
sunnyqgg added a commit to sunnyqgg/TensorRT-LLM that referenced this pull request Aug 27, 2026
…r teardown

The DG SymmBuffer cache was keyed on id(self._ep_pg) and never evicted, so
NVLink symmetric-memory activation workspaces outlived the LLM that allocated
them. A worker process is reused across LLMs (a new executor is built per LLM
while the CUDA context persists) and the EP group is destroyed on executor
shutdown, so the next LLM computed a different key and allocated a second
buffer beside the first instead of reusing it. These come from
empty_strided_p2p, outside PyTorch's caching allocator, so the
torch.cuda.empty_cache() already on the worker shutdown path cannot reclaim
them; ~40 GiB of 178 GiB was gone before the failing test allocated anything
and _alloc_symm_buffer died with "CUDA driver error: out of memory". id() is
also recycled once a group is freed, so a new group could land on a dead
group's id and hit a buffer rendezvoused over a destroyed group.

Key the cache on buffer geometry so a later LLM in a reused worker reuses the
buffer, and release the cache from the worker's shutdown path, just before the
EP group these buffers were rendezvoused over is destroyed. Eviction is
deterministic rather than reachability-based on purpose: SymmBuffer.__init__
keeps a strong reference to its own group, so a weakly-held owner would stay
reachable through the cached buffer and never be collected. Releasing also
clears every remaining tensor attribute, because SymmBuffer.destroy() nulls
only a few of the views sliced out of the allocation and any survivor pins the
whole buffer.

Supersedes PR NVIDIA#17529, additionally addressing its review feedback:
- worker.shutdown() no longer lets a failing release skip
  destroy_process_group() and the CUDA cleanup below it: doing_shutdown is
  already set on entry, so an escaping exception would leave NCCL
  communicators alive with no retry possible.
- release_symm_buffer_cache() frees buffer-by-buffer and keeps going on
  failure, instead of sum() over a generator that abandons the remaining
  buffers (already evicted, hence unreclaimable) on the first exception.
- Cache entries record the ProcessGroup TRT-LLM passed at allocation, and
  stale-hit validation checks that record instead of SymmBuffer.group, so a
  DeepGEMM bump that normalizes the stored group cannot turn every hit into
  a false stale that frees a live buffer another layer already holds.
- The sys.modules probe string lives in a named constant next to the import
  block instead of being buried inline in shutdown().
- _free_symm_buffer() tolerates an already-destroyed buffer.
- Document that the Ray/RPC worker paths rely on per-LLM process exit for
  reclamation and must call the release if they ever reuse processes.

Unwaive the test now that it passes; the DGX_B200 run of
TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm scheduled after the
heavy predecessor sequence is the validation gate.

Signed-off-by: qgai <qgai@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants