From ad23d2d870868a35e92207282975635d04151775 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 17:06:47 +0800 Subject: [PATCH] Neuron MoE capacity path: all-reduce the overflow flag; bench prompts differ per rank (#26) The overflow decision was rank-local, and the dense rerun has collectives of its own: one rank rerunning left the others waiting in its all-gather (NRT 'Collective operation pending'). The flag is now max-all-reduced over the EP group. The bench gave every rank the same prompt, which routes 64 copies of each token to the same experts (overflow on every layer); real requests differ, so the prompts now take the rank in their seed. Co-Authored-By: Claude Opus 5.5 --- deepinfer/neuron_ops/static_moe.py | 7 ++++++- examples/neuron/bench_neuron.py | 4 ++-- vendor/DeepEP | 2 +- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 9277970..ba86d74 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -230,5 +230,10 @@ def _nki_moe_forward( ) out = out + nki_gemm.linear_qdq(shared, moe.shared_down_proj.weight).float() if capacity: - return out.to(torch.bfloat16), overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device) + flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)).to(torch.float32) + if size > 1: + # Every rank must take the same path: the dense rerun has collectives of its own (a rank-local decision + # left the others waiting in its all-gather: "Collective operation pending"). + flag = funcol.wait_tensor(funcol.all_reduce(flag.reshape(1), "max", group)).reshape(()) + return out.to(torch.bfloat16), flag > 0 return out.to(torch.bfloat16) diff --git a/examples/neuron/bench_neuron.py b/examples/neuron/bench_neuron.py index 7f91c9b..5bcc6ec 100644 --- a/examples/neuron/bench_neuron.py +++ b/examples/neuron/bench_neuron.py @@ -119,12 +119,12 @@ def main(dp_size: int, dp_rank: int) -> None: lens = [int(x) for x in args.input_lens.split(",")] for w in range(args.warmup): - _run_one(llm, _random_prompt(lens[0], 10_000 + w, vocab) if busy else None, params) + _run_one(llm, _random_prompt(lens[0], 10_000 + w + 7919 * dp_rank, vocab) if busy else None, params) results = [] for length in lens: cell = [] for r in range(args.repeats): - got = _run_one(llm, _random_prompt(length, 1000 * length + r, vocab) if busy else None, params) + got = _run_one(llm, _random_prompt(length, 1000 * length + r + 7919 * dp_rank, vocab) if busy else None, params) if got is not None: ttft, latency, n = got tpot = (latency - ttft) / max(n - 1, 1) diff --git a/vendor/DeepEP b/vendor/DeepEP index f6aa38f..93eb6eb 160000 --- a/vendor/DeepEP +++ b/vendor/DeepEP @@ -1 +1 @@ -Subproject commit f6aa38fe06cd4a0489e7730365d4e0114970a2f3 +Subproject commit 93eb6eb238127e96c6d7a4a625a6dad158348509