Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
trtllm-agent wants to merge 1 commit into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 66 additions & 2 deletions tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,9 +55,56 @@
# on the current TRT-LLM execution contract that MegaMoE layers run
# serially within a forward pass; concurrent MegaMoE forwards sharing
# a key would race on the same scratch buffers.
#
# Keyed on buffer geometry only, never on the EP group's identity: a
# worker process outlives the LLM built in it, so keying on the group
# would mint a fresh key per LLM and stack a second allocation on top
# of the first (and ``id()`` is recycled, so it can also alias a dead
# group). Since the key no longer distinguishes groups, a hit is
# re-validated against the caller's live group in ``_alloc_symm_buffer``.
# ``release_symm_buffer_cache`` bounds the lifetime.
_MEGA_MOE_SYMM_BUFFER_CACHE: Dict[tuple, object] = {}


def _free_symm_buffer(buffered: object) -> int:
"""Release one DG SymmBuffer's symmetric memory. Returns its size in bytes.

``SymmBuffer.destroy`` nulls only ``handle``/``buffer``/``group``/``x``/
``x_sf``, leaving 10 of the 12 tensor views sliced out of the allocation
in place -- and any surviving view pins the whole thing. So sweep every
remaining tensor attribute too, otherwise this frees nothing whenever the
object is held by a reference cycle rather than dropped outright.
"""
nbytes = buffered.buffer.nbytes
buffered.destroy()
for name, value in list(vars(buffered).items()):
if isinstance(value, torch.Tensor):
setattr(buffered, name, None)
return nbytes


def release_symm_buffer_cache() -> None:
"""Free every cached DG SymmBuffer. Call once per executor teardown.

The EP group these buffers were rendezvoused over is destroyed on
executor shutdown, so nothing may reuse them afterwards. They must be
dropped explicitly: the backing symmetric memory sits outside PyTorch's
caching allocator, so neither a GC nor ``torch.cuda.empty_cache()``
reclaims it, and a reused worker process would otherwise carry a full
activation workspace into the next LLM's memory budget.
"""
if not _MEGA_MOE_SYMM_BUFFER_CACHE:
return
# Detach before freeing so a raising free cannot leave a half-destroyed
# 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)

@brnguyen2 brnguyen2 Aug 12, 2026

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

logger.info(
f"[MegaMoE] released {len(buffers)} DG SymmBuffer(s): {total_bytes / 2**30:.2f} GiB"
)


# ---- Fused MXFP8 per-token quant backends --------------------------------
# We want: BF16 (m, H) → FP8 E4M3 (m, H) + packed-UE8M0 SF (m, H/32/4) int32.
# Three candidates, in preference order:
Expand Down Expand Up @@ -580,7 +627,8 @@ def _alloc_symm_buffer(self) -> None:
capture would fail on the host-side IPC handle exchange.

Buffers are shared across layers via ``_MEGA_MOE_SYMM_BUFFER_CACHE``
keyed on the (EP-PG, slot/expert/topk/shape/activation) tuple.
keyed on the (EP-size, slot/expert/topk/shape/activation) tuple; see
that cache's definition for why the key carries no group identity.
Sharing is safe only while MegaMoE layer forwards are issued
serially within a forward pass; concurrent MegaMoE forwards
sharing a key would race on the same scratch buffers.
Expand All @@ -600,7 +648,7 @@ def _alloc_symm_buffer(self) -> None:
if self._symm_buffer is not None:
return
key = (
id(self._ep_pg),
self.ep_size,
self.num_experts,
self.num_slots,
self.max_num_tokens,
Expand All @@ -610,6 +658,22 @@ def _alloc_symm_buffer(self) -> None:
self.activation,
)
cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key)
if cached is not None and cached.group is not self._ep_pg:

@brnguyen2 brnguyen2 Aug 12, 2026

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.

# Geometry-equal but rendezvoused over a different (already
# destroyed) EP group: an LLM torn down without
# release_symm_buffer_cache() leaves such an entry behind, and its
# buffer's peer mappings are dead. Free it here so the driver can
# reuse those pages for the allocation below -- symmetric memory is
# outside the caching allocator, so this is the only way to get it
# back within the process. Evict before freeing so a raising free
# cannot leave a half-destroyed buffer behind for the next lookup.
del _MEGA_MOE_SYMM_BUFFER_CACHE[key]
freed = _free_symm_buffer(cached)
cached = None
logger.info(
f"[MegaMoE] layer={self.layer_idx} released stale DG "
f"SymmBuffer: {freed / 2**30:.2f} GiB"
)
if cached is None:
cached = self._dg.get_symm_buffer_for_mega_moe(
self._ep_pg,
Expand Down
11 changes: 11 additions & 0 deletions tensorrt_llm/executor/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import gc
import os
import sys
import threading
import time
import traceback
Expand Down Expand Up @@ -134,6 +135,16 @@ def shutdown(self):
self._executor_config.checkpoint_loader.cleanup()
self._executor_config.checkpoint_loader = None

# 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.

@brnguyen2 brnguyen2 Aug 12, 2026

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.

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()
Comment on lines +138 to +146

@coderabbitai coderabbitai Bot Aug 12, 2026

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.


# Destroy torch distributed process groups so that NCCL communicators
# are torn down cleanly before MPI session shutdown and process exit.
# This is done here (not in PyExecutor.shutdown()) because the MPI
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype SKIP (https://nvbugs/6561677)
accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm SKIP (https://nvbugs/6581063)
accuracy/test_llm_api_pytorch.py::TestGLM52::test_nvfp4[tp_size=8-ep_size=8] SKIP (https://nvbugs/6507108)
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_guided_decoding_4gpus[one_model] SKIP (https://nvbugs/5596343)
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_1gpu[v1_kv_cache-True-True-triton-auto] SKIP (https://nvbugs/6026676)
Expand Down
Loading