neuron: shared expert sliced over ranks, summed in the MoE all-reduce (DEEPINFER_NEURON_SHARED_SLICE=1)

18 of 64 ranks compute a 128-column slice of the 2304-wide shared expert (the rest add zeros, same graph);
the partial joins the routed experts' sum before the existing all-reduce. Decode step 591 -> 581 ms.
Summation order changes, so greedy outputs differ from before (coherent; tokens/step 2.70/3.10/3.34,
TPOT p50 184 ms). Tiny golden 9/9. Also DEEPINFER_NEURON_HOST_TIMING (host time per region call).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/model_loader.py

@@ -85,6 +85,9 @@ def get_model(
         from deepinfer.neuron_ops import nki_gemm
 
         logger.info(f"Neuron: {nki_gemm.prepare_model(model)} weights take the NKI block-quant GEMMs")
+        from deepinfer.neuron_ops import static_moe
+
+        logger.info(f"Neuron: {static_moe.prepare_shared_slices(model)} MoE layers compute a slice of the shared expert")
         from deepinfer.neuron_ops import nki_hc
 
         logger.info(f"Neuron: {nki_hc.prepare_model(model)} mHC fn matrices split into bf16 hi + lo")

deepinfer/neuron_ops/nki_gemm.py

@@ -120,6 +120,12 @@ def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> to
     return _linear_kernel(name, weight.shape[0])(x, weight, scale)[: x.shape[0]]
 
 
+def linear_qdq_raw(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """``linear_qdq`` for a bare fp8 weight (uint8 bytes ``[n, k]`` + kernel scale), e.g. a rank's shared-expert
+    slice (static_moe.prepare_shared_slices)."""
+    return _linear_kernel("fp8_linear", weight.shape[0])(x, weight, scale)[: x.shape[0]]
+
+
 def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.Tensor) -> torch.Tensor | None:
     """``linear_qdq`` for fp4 expert ``group`` that the kernel skips (zeros) unless ``flag`` (int32 ``[1, 1]``) is 1.
     None when the gated kernel does not apply (not fp4, or columns that do not split across the two cores)."""

deepinfer/neuron_ops/regions.py

@@ -79,6 +79,12 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
 
         def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
             with hide_cuda():
+                if _HOST_TIMING:
+                    started = time.perf_counter()
+                    out = compiled(mod, *args, **kwargs)
+                    _HOST[label][0] += 1
+                    _HOST[label][1] += time.perf_counter() - started
+                    return out
                 if not envs.DEEPINFER_NEURON_REGION_TIMING:
                     out = compiled(mod, *args, **kwargs)
                     if _SYNC_EVERY:
@@ -103,6 +109,16 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
 
 
 _MARK: list = [0.0, 0, ""]
+#: DEEPINFER_NEURON_HOST_TIMING: host seconds spent inside each region kind's compiled call (no synchronize: launch,
+#: guard evaluation, resume frames, and any wait for the execution queue), per step (take_host_timings).
+_HOST_TIMING = os.environ.get("DEEPINFER_NEURON_HOST_TIMING") == "1"
+_HOST: dict[str, list] = defaultdict(lambda: [0, 0.0])
+
+
+def take_host_timings() -> dict[str, list]:
+    out = {k: list(v) for k, v in _HOST.items()}
+    _HOST.clear()
+    return out
 
 
 @torch._dynamo.disable

deepinfer/neuron_ops/static_moe.py

@@ -39,6 +39,11 @@ _ALIGN = 16
 _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", "")
+#: DEEPINFER_NEURON_SHARED_SLICE: each rank computes a 128-column slice of the shared expert's intermediate dimension
+#: and adds it to the routed experts' partial sum before the MoE all-reduce, instead of every rank computing the whole
+#: shared expert after it (35 M fp8 weights a layer, dequantized on every rank).
+_SHARED_SLICE = os.environ.get("DEEPINFER_NEURON_SHARED_SLICE", "1") == "1"
+_SHARED_COLS = 128
 
 def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor:
     return t.contiguous().view(rows, -1).view(torch.uint8)
@@ -124,6 +129,52 @@ def _reduce_scatter(t: torch.Tensor, group: torch.distributed.ProcessGroup, size
     return funcol.wait_tensor(funcol.reduce_scatter_tensor(t, "sum", 0, group))
 
 
+def prepare_shared_slices(model: torch.nn.Module) -> int:
+    """Register this rank's shared-expert slice on every MoE (after nki_gemm.prepare_model). ``I // 128`` ranks own a
+    slice; the others get zero weights of the same shape, so every rank traces the same graph and their slice adds
+    exact zeros. Returns the number of MoE layers sliced."""
+    if not _SHARED_SLICE:
+        return 0
+    count = 0
+    for moe in model.modules():
+        gu = getattr(moe, "shared_gate_up_proj", None)
+        dn = getattr(moe, "shared_down_proj", None)
+        if not getattr(moe, "has_shared", False) or gu is None or dn is None:
+            continue
+        gw, dw = gu.weight, dn.weight
+        if getattr(gw, "_nki_scale", None) is None or getattr(dw, "_nki_scale", None) is None or gw.dtype != torch.float8_e4m3fn:
+            continue
+        size, rank = moe.ep_size, moe.ep_rank
+        inter = gw.shape[0] // 2
+        slices = inter // _SHARED_COLS
+        if size == 1 or inter % _SHARED_COLS or slices > size:
+            continue
+        c, b = _SHARED_COLS, _SHARED_COLS // 32
+        # On the host: strided slices of device tensors are not safe eagerly on a NeuronCore (one-time, at load).
+        device = gw.device
+        g_w, g_s, d_w, d_s = (t.cpu() for t in (gw._nki_weight, gw._nki_scale, dw._nki_weight, dw._nki_scale))
+        if rank < slices:
+            lo = rank * c
+            gu_w = torch.cat([g_w[lo : lo + c], g_w[inter + lo : inter + lo + c]])
+            gu_s = torch.cat([g_s[lo // 32 : lo // 32 + b], g_s[(inter + lo) // 32 : (inter + lo) // 32 + b]])
+            dn_w, dn_s = d_w[:, lo : lo + c], d_s[:, lo // 32 : lo // 32 + b]
+        else:
+            gu_w = torch.zeros((2 * c, g_w.shape[1]), dtype=g_w.dtype, device=g_w.device)
+            gu_s = torch.full((2 * b, g_s.shape[1]), 127, dtype=g_s.dtype, device=g_s.device)  # e8m0 127 = 1.0
+            dn_w = torch.zeros((d_w.shape[0], c), dtype=d_w.dtype, device=d_w.device)
+            dn_s = torch.full((d_s.shape[0], b), 127, dtype=d_s.dtype, device=d_s.device)
+        for name, t in (("gu_w", gu_w), ("gu_s", gu_s), ("dn_w", dn_w), ("dn_s", dn_s)):
+            moe.register_buffer(f"_neuron_shared_{name}", t.contiguous().to(device), persistent=False)
+        count += 1
+    return count
+
+
+def _shared_slice(moe: DeepseekMoE, x: torch.Tensor) -> torch.Tensor:
+    gate_up = nki_gemm.linear_qdq_raw(x, moe._neuron_shared_gu_w, moe._neuron_shared_gu_s)
+    act = swiglu_forward_and_per_token_cast(gate_up, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value)
+    return nki_gemm.linear_qdq_raw(act, moe._neuron_shared_dn_w, moe._neuron_shared_dn_s)
+
+
 def static_moe_forward(
     moe: DeepseekMoE,
     hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -283,11 +334,14 @@ def _nki_moe_forward(
             if down is None:
                 down = nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e)
             acc = acc + down.float()
+    sliced = moe.has_shared and getattr(moe, "_neuron_shared_gu_w", None) is not None
+    if sliced:
+        acc = acc + _shared_slice(moe, x).float()
     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()
-    if moe.has_shared:
+    if moe.has_shared and not sliced:
         shared = nki_gemm.linear_qdq(x_local, moe.shared_gate_up_proj.weight)
         shared = swiglu_forward_and_per_token_cast(
             shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value

deepinfer/worker/neuron_graph_runner.py

@@ -199,6 +199,21 @@ class NeuronGraphRunner(CudaGraphRunner):
             reqs = -(-batch_size // (1 + envs.DEEPINFER_USE_NEXTN_DECODE))
             nki_attention.set_table_bounds(reqs, getattr(input_metadata_device, "neuron_ctx_limit", None))
         if not envs.DEEPINFER_NEURON_REGION_TIMING:
+            if regions._HOST_TIMING:
+                started = time.perf_counter()
+                with hide_cuda():
+                    out = self.compiled[batch_size](input_metadata_device)
+                self._host_steps = getattr(self, "_host_steps", 0) + 1
+                if self._host_steps % 20 == 0 and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0):
+                    host = regions.take_host_timings()
+                    inside = sum(v[1] for v in host.values())
+                    logger.info(
+                        f"Neuron host bs={batch_size}: forward {1e3 * (time.perf_counter() - started):.1f} ms, in regions "
+                        f"{1e3 * inside:.1f} ms; " + ", ".join(f"{k} {v[0]}x {1e3 * v[1]:.1f} ms" for k, v in sorted(host.items(), key=lambda kv: -kv[1][1]))
+                    )
+                else:
+                    regions.take_host_timings()
+                return out
             with hide_cuda():
                 return self.compiled[batch_size](input_metadata_device)
         torch.neuron.synchronize()