[https://nvbugs/6581063][fix] Keep the verified fix (geometry-only cache key + deterministic eviction from… - #17529
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan includes up to 12 reviews per rolling hour; 11 remain after this review. WalkthroughThe 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. ChangesMegaMoE buffer lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to 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: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
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 winUse a structural
SymmBufferprotocol for the cache.The repository does not expose a concrete DeepGEMM
SymmBuffertype. Define a private protocol forbuffer,group, anddestroy(), and use it for the cache and_free_symm_buffer(). ReplaceDict[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
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.pytensorrt_llm/executor/worker.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
| # 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() |
There was a problem hiding this comment.
🩺 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.pyRepository: 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)))
PYRepository: 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
left a comment
There was a problem hiding this comment.
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: |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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. |
There was a problem hiding this comment.
Two gaps with the sys.modules probe:
- 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
importinside a narrowtry: ... except ImportErrorwould at least fail loudly at the right time, or add a test that asserts the string resolves. - This is the MPI worker path only. As the comment right below notes, the Ray path owns its process group in
RayWorkerWrapperand 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.
BowenFu
left a comment
There was a problem hiding this comment.
Changes required: cache eviction is not yet safe or verified.
- The cache compares external
SymmBuffer.groupby 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 firstdestroy()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-1with 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>
5092ac6 to
965bc1f
Compare
|
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. |
|
Superseded by #18182, which keeps this PR's fix (geometry-only cache key + deterministic eviction from |
…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>
…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>
…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>
Summary
_MEGA_MOE_SYMM_BUFFER_CACHEwas keyed onid(self._ep_pg)and never evicted, so NVLink symmetric-memory workspaces (allocated viaempty_strided_p2p, outside the caching allocator, soempty_cache()cannot reclaim them) outlived their LLM in a reused MPI worker and stacked a second allocation per LLM until the driver OOM'd.worker.shutdown()beforedestroy_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_pgand freeing the stale buffer in place via a shared_free_symm_bufferhelper.pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm -vTest plan
Links
Dev Engineer Review
_ep_pgidentity. This enables reuse across LLM instances._ep_pg. Stale buffers are freed through_free_symm_buffer.worker.shutdown()releases cached buffers before it destroys process groups.release_symm_buffer_cache()API releases cached buffers and logs reclaimed memory.CODING_GUIDELINES.mdand validated with the related regression tests.QA Engineer Review
tests/integration/test_lists/waives.txt.TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm.