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 <noreply@anthropic.com>

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)

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)

vendor/DeepEP

@@ -1 +1 @@
-Subproject commit f6aa38fe06cd4a0489e7730365d4e0114970a2f3
+Subproject commit 93eb6eb238127e96c6d7a4a625a6dad158348509