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

neuron: diagnostics for MoE collective skew (expert flag force, per-follower runtime profiling)

Browse Source 
DEEPINFER_NEURON_MOE_FLAG_FORCE=0|1 forces the gated experts' flags (timing only, wrong results);
DEEPINFER_NEURON_PROFILE_RANKS keeps NEURON_RT_INSPECT on for the listed followers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
whn09 2026-10-26 15:14:19 +08:00
parent d7d2f36ec1
commit b6d8df7484
2 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
deepinfer/engine/llm_engine.py
View File

@ -395,7 +395,10 @@ class LLMEngine:
# Neuron: the follower claims core cp_rank of CUDA_VISIBLE_DEVICES (deepinfer/platform.py); the spawned
# process copies this environment at start.
# Runtime profiling (NEURON_RT_INSPECT_*) stays on the driver: 64 ranks would overwrite one output dir.
overrides = {"DEEPINFER_NEURON_LOCAL_RANK": str(cp_rank), "NEURON_RT_INSPECT_ENABLE": "0"}
# DEEPINFER_NEURON_PROFILE_RANKS (e.g. "1,63") keeps it on for those followers (each writes a pid dir).
overrides = {"DEEPINFER_NEURON_LOCAL_RANK": str(cp_rank)}
if str(cp_rank) not in os.environ.get("DEEPINFER_NEURON_PROFILE_RANKS", "").split(","):
overrides["NEURON_RT_INSPECT_ENABLE"] = "0"
saved = {k: os.environ.get(k) for k in overrides}
os.environ.update(overrides)
try:

3
deepinfer/neuron_ops/static_moe.py
View File

@ -38,6 +38,7 @@ _ALIGN = 16
#: per-expert flag, blockquant.fp4_linear_2core_gated).
_GATED_EXPERTS = os.environ.get("DEEPINFER_NEURON_GATED_EXPERTS", "1") == "1"
_GATED_MAX_ROWS = 16
_FLAG_FORCE = os.environ.get("DEEPINFER_NEURON_MOE_FLAG_FORCE", "")
def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor:
return t.contiguous().view(rows, -1).view(torch.uint8)
@ -270,6 +271,8 @@ def _nki_moe_forward(
for e in range(num_local):
w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
flag = (local_idx == e).any().to(torch.int32).reshape(1, 1) if gated else None
if gated and _FLAG_FORCE in ("0", "1"): # timing diagnostic only: wrong results
flag = torch.full((1, 1), int(_FLAG_FORCE), dtype=torch.int32, device=x.device)
gate_up = nki_gemm.linear_qdq_gated(x, moe.gate_up_proj_weight, e, flag) if gated else None
if gate_up is None:
gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e)