From 39f8aa6eedffa2ea7419deb912df1aa0f1b29706 Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 00:01:52 +0800 Subject: [PATCH] Neuron DSpark: 6-row steps run the MoE as one region graph (5 graphs a layer) A world all-reduce costs ~3.2 ms per graph that holds one, whatever it carries (microbenchmark at 64 ranks: one all-reduce a graph 3.5 ms, ten chained in one graph 3.6 ms; fp32 / bf16 and 1 / 6 rows alike), so the step pays per collective-bearing graph, not per byte. The MoE's gate, experts, all-reduce and shared expert now stay one graph at decode sizes (prompt chunks still split). bench62 (512 in, TP=64, nextn=5): engine step 805 -> 740 ms. Random-token prompts make greedy outputs chaotic, so a graph-boundary change that moves bf16 rounding changes the generated tokens and the acceptance; the step time is the comparable figure. Co-Authored-By: Claude Opus 5.5 --- deepinfer/neuron_ops/static_moe.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index d184ef6..85de881 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -251,7 +251,7 @@ def _nki_moe_forward( x = packed[:, : x.shape[1]].to(torch.bfloat16) idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32) w = packed[:, x.shape[1] + topk :] - regions.split(coarse=True) + regions.split() local_idx = idx - moe.neuron_rank_tensor() * num_local overflow = None if capacity and capacity < x.shape[0]: @@ -265,10 +265,10 @@ def _nki_moe_forward( gate_up, fmt="bf16", num_per_channels=None, topk_weights=w_e, clamp_value=moe.swiglu_clamp_value ) acc = acc + nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e).float() - regions.split(coarse=True) + regions.split() hier = getattr(moe, "_neuron_hier_ar", None) out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] - regions.split(coarse=True) + regions.split() if moe.has_shared: shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight) shared = swiglu_forward_and_per_token_cast(