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..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 @@ -63,7 +63,122 @@ # 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] + # 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" + ) + 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 +679,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 +700,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 +709,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 +722,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)