From 7f17fb587db46edf0fff3618ab18807d1561a64e Mon Sep 17 00:00:00 2001 From: qgai Date: Mon, 24 Aug 2026 23:24:04 -0700 Subject: [PATCH 1/2] [https://nvbugs/6581063][fix] Release MegaMoE symm buffers on executor 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 #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 --- .../fused_moe/mega_moe/mega_moe_deepgemm.py | 114 +++++++++++++++++- tensorrt_llm/executor/worker.py | 20 +++ tests/integration/test_lists/waives.txt | 1 - 3 files changed, 129 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index 0258879b2651..2b0137d4d2be 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -63,7 +63,110 @@ # 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. -_MEGA_MOE_SYMM_BUFFER_CACHE: Dict[tuple, object] = {} +# +# 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, each entry +# records the ProcessGroup it was allocated over and a hit is +# re-validated against the caller's live group in +# ``_take_cached_symm_buffer``. The group is recorded on this side of +# the DeepGEMM boundary on purpose: relying on ``SymmBuffer.group`` +# would make correctness depend on DeepGEMM storing the passed group +# verbatim, and a version bump that normalizes it would turn every hit +# into a false stale -- freeing a buffer another layer already holds. +# ``release_symm_buffer_cache`` bounds the lifetime. +_MEGA_MOE_SYMM_BUFFER_CACHE: Dict[tuple, Tuple[object, 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 the other 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. + + Tolerates an already-destroyed buffer (``buffer`` is ``None``) so a + double free degrades to a no-op instead of taking down the teardown path. + """ + buffer = getattr(buffered, "buffer", None) + nbytes = buffer.nbytes if buffer is not None else 0 + buffered.destroy() + for name, value in list(vars(buffered).items()): + if isinstance(value, torch.Tensor): + setattr(buffered, name, None) + return nbytes + + +def _take_cached_symm_buffer(key: tuple, ep_pg: object) -> Optional[object]: + """Return the cached SymmBuffer for ``key`` if it was allocated over + ``ep_pg``; free and drop a geometry-equal entry from another group. + + A geometry-equal entry recorded against a different (already destroyed) + EP group means an LLM was torn down without ``release_symm_buffer_cache``, + and its buffer's peer mappings are dead. Free it here so the driver can + reuse those pages for the caller's allocation -- 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. + """ + entry = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key) + if entry is None: + return None + cached, cached_pg = entry + if cached_pg is ep_pg: + return cached + del _MEGA_MOE_SYMM_BUFFER_CACHE[key] + freed = _free_symm_buffer(cached) + logger.info( + f"[MegaMoE] released stale DG SymmBuffer from a previous EP group: {freed / 2**30:.2f} GiB" + ) + return None + + +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. + + Only the MPI worker path (``GenerationExecutorWorker.shutdown``) calls + this today. The Ray and RPC worker paths do not: their worker process + exits with the LLM, so the driver reclaims the memory at process exit. + Any path that starts reusing a worker process across LLMs must call this + on its teardown too. + """ + 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. + entries = list(_MEGA_MOE_SYMM_BUFFER_CACHE.values()) + _MEGA_MOE_SYMM_BUFFER_CACHE.clear() + # Free one at a time and keep going on failure: the entries are already + # out of the cache, so any buffer skipped by a raising predecessor would + # be unreclaimable for the rest of the process lifetime. + released = 0 + total_bytes = 0 + for buffered, _ in entries: + try: + total_bytes += _free_symm_buffer(buffered) + released += 1 + except Exception as e: + logger.error( + f"[MegaMoE] failed to release a DG SymmBuffer during executor " + f"teardown, continuing with the remaining buffers: {e}" + ) + logger.info( + f"[MegaMoE] released {released}/{len(entries)} DG SymmBuffer(s): " + f"{total_bytes / 2**30:.2f} GiB" + ) # ---- Fused MXFP8 per-token quant backends -------------------------------- @@ -564,7 +667,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. @@ -584,7 +688,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, @@ -593,7 +697,7 @@ def _alloc_symm_buffer(self) -> None: self.intermediate_size, self.dg_activation, ) - cached = _MEGA_MOE_SYMM_BUFFER_CACHE.get(key) + cached = _take_cached_symm_buffer(key, self._ep_pg) if cached is None: cached = self._dg.get_symm_buffer_for_mega_moe( self._ep_pg, @@ -606,7 +710,7 @@ def _alloc_symm_buffer(self) -> None: mma_type="fp8xfp4", activation=self.dg_activation, ) - _MEGA_MOE_SYMM_BUFFER_CACHE[key] = cached + _MEGA_MOE_SYMM_BUFFER_CACHE[key] = (cached, self._ep_pg) # Log only on the first layer; deeper layers reuse the cache # and would otherwise spam N copies of an identical line. log_fn = logger.info if self.layer_idx == 0 else logger.debug diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 34779c032536..61ca7de6e5e8 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -38,6 +38,12 @@ # touches the PyTorch model zoo should not pay for it. _MODELING_UTILS_MODULE = "tensorrt_llm._torch.models.modeling_utils" +# Probed (not imported) in shutdown() so a non-MegaMoE run does not pull in +# the MoE stack just to release an empty cache. Must stay in sync with the +# real module path -- a mismatch silently degrades the release to a no-op. +_MEGA_MOE_DEEPGEMM_MODULE = ( + "tensorrt_llm._torch.moe.fused_moe.mega_moe.mega_moe_deepgemm") + class GenerationExecutorWorker(RpcWorkerMixin, BaseWorker): @@ -139,6 +145,20 @@ 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. Never let a failure here escape: + # doing_shutdown is already set, so a retry would no-op, and skipping + # the teardown below would leave NCCL communicators alive -- turning a + # MegaMoE-only failure into a worker teardown hang. + mega_moe = sys.modules.get(_MEGA_MOE_DEEPGEMM_MODULE) + if mega_moe is not None: + try: + mega_moe.release_symm_buffer_cache() + except Exception as e: + logger.error( + f"Failed to release MegaMoE symm buffers on shutdown: {e}") + # 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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b7da811e1181..fe792c43e5df 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -31,7 +31,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_sched accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_python_scheduler[mtp_nextn=2-attention_dp=False-cuda_graph=False-overlap_scheduler=False-enable_chunked_prefill=True] SKIP (https://nvbugs/6507095) accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6655987) 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::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-no_overlap_scheduler] SKIP (https://nvbugs/6644462) 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) From 06f76eb597a6123fb85f0e1290f7f60b96f286a3 Mon Sep 17 00:00:00 2001 From: qgai Date: Tue, 25 Aug 2026 01:28:40 -0700 Subject: [PATCH 2/2] [https://nvbugs/6581063][fix] Do not fail LLM init when freeing a stale symm buffer raises The stale-group hit in _take_cached_symm_buffer runs on the next LLM's init path, and the buffer being freed was rendezvoused over an already-destroyed EP group, so its destroy() faces dead peer mappings. An escaping exception there would fail the new LLM's construction -- the exact reused-worker scenario this fix targets -- whereas the pre-fix code merely leaked and allocated fresh, so a raising free was survivable. Catch it, log, and fall through to the fresh allocation, mirroring the per-buffer error handling release_symm_buffer_cache already has. Addresses review feedback from Bowen Fu on PR #18182. Signed-off-by: qgai --- .../moe/fused_moe/mega_moe/mega_moe_deepgemm.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py index 2b0137d4d2be..a1a1b411abae 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -120,7 +120,19 @@ def _take_cached_symm_buffer(key: tuple, ep_pg: object) -> Optional[object]: if cached_pg is ep_pg: return cached del _MEGA_MOE_SYMM_BUFFER_CACHE[key] - freed = _free_symm_buffer(cached) + # This runs on the next LLM's init path and the stale buffer was + # rendezvoused over an already-destroyed group, so a raising free must + # not escape and fail the new LLM -- fall through to a fresh allocation + # instead (no worse than the pre-fix behavior, which leaked and + # allocated). + try: + freed = _free_symm_buffer(cached) + except Exception as e: + logger.error( + f"[MegaMoE] failed to release a stale DG SymmBuffer from a " + f"previous EP group, allocating fresh: {e}" + ) + return None logger.info( f"[MegaMoE] released stale DG SymmBuffer from a previous EP group: {freed / 2**30:.2f} GiB" )