coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 14:57:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Neuron DSpark: 6-row steps run the MoE as one region graph (5 graphs a layer)

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 00:01:52 +08:00
parent 5e36f90015
commit 39f8aa6eed
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
deepinfer/neuron_ops/static_moe.py
View File

@ -251,7 +251,7 @@ def _nki_moe_forward(
x = packed[:, : x.shape[1]].to(torch.bfloat16) x = packed[:, : x.shape[1]].to(torch.bfloat16)
idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32) idx = packed[:, x.shape[1] : x.shape[1] + topk].to(torch.int32)
w = packed[:, x.shape[1] + topk :] w = packed[:, x.shape[1] + topk :]
regions.split(coarse=True) regions.split()
local_idx = idx - moe.neuron_rank_tensor() * num_local local_idx = idx - moe.neuron_rank_tensor() * num_local
overflow = None overflow = None
if capacity and capacity < x.shape[0]: 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 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() 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) hier = getattr(moe, "_neuron_hier_ar", None)
out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens] 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: if moe.has_shared:
shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight) shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
shared = swiglu_forward_and_per_token_cast( shared = swiglu_forward_and_per_token_cast(