[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm SymmBuffer by qiaoxj07 · Pull Request #14213 · NVIDIA/TensorRT-LLM

qiaoxj07 · 2026-05-17T06:18:01Z

Summary

When attention DP is enabled, the DG MegaMoE SymmBuffer per-rank pool
is sized as dp_size * ep_size * max_num_tokens * num_topk instead of
dp_size * max_num_tokens * num_topk — the EP factor is double-counted.
On DEP=32 DSv4 (max_num_tokens=1024, num_topk=6, hidden=7168,
moe_intermediate=3072) this inflates the SymmBuffer to 95 GiB per rank
and OOMs the GEN ranks during model load on a 276 GiB GB300 GPU.

Root cause: ModelConfig.__post_init__ pre-multiplies
moe_max_num_tokens = max_num_tokens * dp_size. MegaMoEDeepGemm.__init__
then forwards this value to get_symm_buffer_for_mega_moe(...), whose
layout::Workspace / get_num_max_pool_tokens (DeepGEMM) further
scales the pool by num_ranks (= ep_size). Because the supported
full-ADP topology asserts ep_size == parallel_size == dp_size, the
dp_size factor is applied twice.

Fix: divide self.max_num_tokens by ep_size (ceil) when
use_dp is true and ep_size > 1. The DG kernel still sees a pool
sized for the worst case where every cluster token routes to one rank's
local experts — but the EP factor is no longer doubled.

Effect (theoretical, DSv4 routed expert, full ADP, max_num_tokens=1024):

Topology	Before	After
DEP=32	95 GiB	~3 GiB
DEP=16	24 GiB	~0.6 GiB
DEP=8	3 GiB	~0.15 GiB
non-ADP	unchanged	unchanged

Non-ADP paths (use_dp == False) and ADP+EP=1 paths are not touched.

Test plan

Rerun DSv4 disagg DEP=32 megamoe_deepgemm GEN that previously OOM'd
on rank 3 / rank 7 — verify SymmBuffer line reports ~3 GiB and GEN
workers reach the benchmark loop.
Rerun DSv4 disagg DEP=16 baseline to confirm no perf regression.
Sanity run on a non-ADP (enable_attention_dp: false) MegaMoE
config to confirm self.max_num_tokens is unchanged in that
path.
L0 / DSv4 unit tests under
tests/unittest/_torch/modules/test_fused_moe.py and
tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py —
none of them exercise the ADP+EP>1 path covered by the fix, but
keep them green.

🤖 Generated with Claude Code

…ymmBuffer 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>

qiaoxj07 · 2026-05-17T06:20:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-17T06:27:39Z

PR_Github #48746 [ run ] triggered by Bot. Commit: 20d832d Link to invocation

tensorrt-cicd · 2026-05-17T08:41:52Z

PR_Github #48746 [ run ] completed with state SUCCESS. Commit: 20d832d
/LLM/main/L0_MergeRequest_PR pipeline #38513 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 · 2026-05-19T02:39:14Z

/bot skip --comment "CI passed before, and only changed the code format in the latest commit."

tensorrt-cicd · 2026-05-19T02:45:10Z

PR_Github #49048 [ skip ] triggered by Bot. Commit: baf5aa7 Link to invocation

tensorrt-cicd · 2026-05-19T02:52:26Z

PR_Github #49048 [ skip ] completed with state SUCCESS. Commit: baf5aa7
Skipping testing for commit baf5aa7

Link to invocation

The MegaMoEDeepGemm wrapper previously read ``model_config.moe_max_num_tokens`` to size its DG SymmBuffer pool and then divided by ``ep_size`` (per NVIDIA#14213) to undo the auto-multiply that ``ModelConfig.__post_init__`` applies for allgather-based backends. Under ``FUSED_COMM`` scheduler (always used by mega_moe), ``FusedCommMoEScheduler._strip_adp_padding`` slices ``x`` to the local rank's tokens before ``run_moe``, so the per-call ``x.shape[0]`` is bounded by the engine's per-instance ``max_num_tokens`` and never by ``dp_size * max_num_tokens``. The auto-multiply is therefore semantically wrong for mega_moe and the divide is a workaround for the wrong source. Two failure modes of the multiply-then-divide dance: 1. If a deployed binary lacks the divide (e.g. binaries built before NVIDIA#14213 landed on this branch), the auto-multiplied value is used as- is, so the cubin's ``num_max_tokens_per_rank`` is baked at ``dp_size * engine.max_num_tokens`` instead of ``engine.max_num_tokens`` (e.g. 65664 vs 8448 on DSv4-Pro @ DEP=8). Measured cost on B300 DEP=8, balanced 8192 tok/rank: 4.03 ms -> 3.39 ms median per-call (-19%), plus the SymmBuffer pool drops from ~3.6 GB to ~470 MB (-7.6x). 2. If a user explicitly sets ``moe_config.max_num_tokens`` in yaml, ``__post_init__`` skips the auto-multiply, but the wrapper still divides, so the per-rank cap becomes ``yaml_value / ep_size`` (too small by ep_size). The line-667 runtime assert then fires for every request whose per-rank tokens exceed that. Read ``model_config.max_num_tokens`` directly instead. The other consumers of ``moe_max_num_tokens`` (fused_moe_cutlass, fused_moe_deepgemm, communication_factory) keep the auto-multiplied value, which is correct for those allgather paths. Test plan: - ``pre-commit run`` clean on the modified file (ruff, ruff-format, codespell, copyright, DCO). - TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[MEGAMOE_DEEPGEMM] on 4x B300 NVFP4 + static EPLB: PASSED. lm-eval GSM8K accuracy 95.262 vs reference 95.110 (threshold 91.907). Runtime ~15 min. - bench_moe at 8192 tok/rank DEP=8 (--analysis kernels --iters 12 --warmup 2) confirms cubin params after fix: sm100_fp8_fp4_mega_moe_impl<8448, 7168, 3072, 384, 6, 2, 192, ..., 414720, 6635520, ...> with mega_moe per-call median 2.62 ms. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> (cherry picked from commit 5a72fb3) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com> (cherry picked from commit 5a72fb3) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

qiaoxj07 requested a review from a team as a code owner May 17, 2026 06:18

qiaoxj07 requested review from QiJune and removed request for a team May 17, 2026 06:18

github-actions Bot assigned qiaoxj07 May 17, 2026

qiaoxj07 requested a review from lfr-0531 May 18, 2026 01:45

lfr-0531 requested a review from Barry-Delaney May 18, 2026 10:03

Barry-Delaney approved these changes May 18, 2026

View reviewed changes

lfr-0531 added the deepseek-v4 label May 19, 2026

[None][fix] Apply MegaMoEDeepGemm formatting

baf5aa7

Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 merged commit 5a72fb3 into NVIDIA:feat/deepseek_v4 May 19, 2026
6 checks passed

mingyangHao mentioned this pull request May 27, 2026

[Draft][None][fix] MegaMoEDeepGemm: use engine max_num_tokens directly #14605

Closed

1 task

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 10, 2026

[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm S…

92894fc

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 11, 2026

[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm S…

2a6d439

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request Jun 12, 2026

[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm S…

5cd160b

…ymmBuffer (NVIDIA#14213) Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm SymmBuffer - #14213

[None][fix] Avoid dp_size x ep_size double-count in MegaMoEDeepGemm SymmBuffer#14213
lfr-0531 merged 2 commits into
NVIDIA:feat/deepseek_v4from
qiaoxj07:fix/megamoe-symm-buffer-adp-double-count

qiaoxj07 commented May 17, 2026 •

edited

Loading

Uh oh!

qiaoxj07 commented May 17, 2026

Uh oh!

tensorrt-cicd commented May 17, 2026

Uh oh!

tensorrt-cicd commented May 17, 2026

Uh oh!

lfr-0531 commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

qiaoxj07 commented May 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

qiaoxj07 commented May 17, 2026

Uh oh!

tensorrt-cicd commented May 17, 2026

Uh oh!

tensorrt-cicd commented May 17, 2026

Uh oh!

lfr-0531 commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

tensorrt-cicd commented May 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

qiaoxj07 commented May 17, 2026 •

edited

Loading