diff --git a/deepinfer/engine/llm_engine.py b/deepinfer/engine/llm_engine.py index acc295c..0a92acb 100644 --- a/deepinfer/engine/llm_engine.py +++ b/deepinfer/engine/llm_engine.py @@ -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: diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 0ecf76a..4201b46 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/deepinfer/neuron_ops/static_moe.py @@ -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)