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..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 @@ -290,6 +290,15 @@ 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