From 20d832d650f81e3cc0eed7f040ca9d5cfe438118 Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Sun, 17 May 2026 14:13:47 +0800 Subject: [PATCH 1/2] [None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm SymmBuffer Under attention DP, ``ModelConfig.__post_init__`` sets ``moe_max_num_tokens = max_num_tokens * dp_size``. The DG SymmBuffer allocator then scales the per-rank pool by ``num_ranks`` (= ep_size, which equals dp_size in the supported full-ADP topology asserted at construction). This double-counts the EP factor: the worst-case pool size becomes ``dp_size * ep_size * max_num_tokens * num_topk`` instead of ``dp_size * max_num_tokens * num_topk``. For a DEP=32 DSv4 run with max_num_tokens=1024, the over-sized SymmBuffer was 95 GiB per rank, OOMing the first two GEN nodes during weight load even on a 276 GiB GPU. Dividing the passed ``num_max_tokens_per_rank`` by ``ep_size`` when ADP is enabled cuts the buffer to ~3 GiB without changing correctness: the DG kernel still covers the worst case where every cluster token routes to one rank's local experts. Non-ADP paths (``use_dp == False``) and ADP+EP=1 are untouched. Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- .../modules/fused_moe/mega_moe/mega_moe_deepgemm.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py index d8f037e02d49..1a927337e2be 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -290,6 +290,17 @@ def __init__( or getattr(model_config, "max_num_tokens", 0) or 4096 ) + # Under attention DP, ``ModelConfig`` pre-multiplies + # ``moe_max_num_tokens`` by ``dp_size``. The DG SymmBuffer further + # scales the pool by ``num_ranks`` (= ep_size, which equals dp_size + # in the supported full-ADP topology asserted above). Without this + # divide the buffer is sized as ``dp_size * ep_size * + # max_num_tokens``, doubling the EP factor and exploding HBM (see + # ``layout::Workspace`` / ``get_num_max_pool_tokens`` in DG). + if self.use_dp and self.ep_size > 1: + self.max_num_tokens = max( + 1, (self.max_num_tokens + self.ep_size - 1) // self.ep_size + ) # Resolve the EP ProcessGroup at module construction — creating a # group at forward time would be collective on a non-synchronous From baf5aa7512d2a57ba0fe337b350479b234a3d403 Mon Sep 17 00:00:00 2001 From: Xianjie <5410381+qiaoxj07@users.noreply.github.com> Date: Tue, 19 May 2026 10:05:19 +0800 Subject: [PATCH 2/2] [None][fix] Apply MegaMoEDeepGemm formatting Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> --- .../_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py index 1a927337e2be..51328e6a4427 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -298,9 +298,7 @@ def __init__( # max_num_tokens``, doubling the EP factor and exploding HBM (see # ``layout::Workspace`` / ``get_num_max_pool_tokens`` in DG). if self.use_dp and self.ep_size > 1: - self.max_num_tokens = max( - 1, (self.max_num_tokens + self.ep_size - 1) // self.ep_size - ) + self.max_num_tokens = max(1, (self.max_num_tokens + self.ep_size - 1) // self.ep_size) # Resolve the EP ProcessGroup at module construction — creating a # group at forward time would be collective on a non-synchronous