neuron: norms feeding only NKI kernels skip their dead fp8 copy (TTFT 1387 -> 1229 ms, step 457 -> 450 ms, tokens identical); capacity MoE wired into the layer region graph with its overflow flag riding the MoE all-reduce (still off: per-layer host sync makes prefill slower, 1823 ms)

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

deepinfer/model_executor/model_loader.py

@@ -85,6 +85,8 @@ 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")
+        if os.environ.get("DEEPINFER_NEURON_BF16_NORMS", "1") == "1":
+            logger.info(f"Neuron: {nki_gemm.bf16_norms(model)} norms feed only NKI kernels and skip their fp8 copy")
         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")

deepinfer/model_executor/models/decoder_layer.py

@@ -176,6 +176,14 @@ class DeepseekDecoderLayer(HCLayerMixin):
             # where it would be a graph output (and its encode would run) for nothing.
             hidden_states = hidden_states_bf16
         neuron_regions.split(coarse=True)
+        capacity = envs.DEEPINFER_NEURON_MOE_CAPACITY
+        if capacity and hidden_states_bf16.shape[0] == envs.DEEPINFER_NEURON_PREFILL_BUCKET:
+            # A prefill chunk: the experts over their routed rows only; the overflow check is a graph break, and an
+            # overflowed chunk reruns the dense path.
+            out, overflow = self.mlp.forward_capacity(hidden_states, hidden_states_bf16, input_metadata)
+            if bool(overflow):
+                out = self.mlp(hidden_states, hidden_states_bf16, input_metadata)
+            return out, residual
         return self.mlp(hidden_states, hidden_states_bf16, input_metadata), residual
 
     def _attn_half(

deepinfer/neuron_ops/nki_gemm.py

@@ -146,3 +146,25 @@ def linear_qdq_gated(x: torch.Tensor, w: torch.Tensor, group: int, flag: torch.T
 
 def has_kernel(w: torch.Tensor | None) -> bool:
     return ENABLED and w is not None and getattr(w, "_nki_scale", None) is not None and w.device.type == "neuron"
+
+
+def bf16_norms(model: torch.nn.Module) -> int:
+    """Stop the attention / ffn norms casting to fp8 where every consumer is an NKI kernel (prepare_model first).
+
+    The kernels quantize the bf16 hidden themselves, so the norm's fp8 copy and its scales were dead -- but they
+    crossed the graph boundary after the mHC pre block as outputs, and their encode (floor / select / power over the
+    whole hidden, mostly on Vector and GpSimd) was about half of that block (3.6 ms of a 512-row prefill chunk,
+    twice a layer). Returns the count of norms switched to bf16.
+    """
+    count = 0
+    for layer in model.modules():
+        attn, mlp = getattr(layer, "self_attn", None), getattr(layer, "mlp", None)
+        pairs = (
+            (getattr(layer, "attn_layernorm", None), has_kernel(getattr(getattr(attn, "q_a_kv_swa_proj", None), "weight", None))),
+            (getattr(layer, "ffn_layernorm", None), has_kernel(getattr(mlp, "gate_up_proj_weight", None)) and not getattr(mlp, "use_mega_moe", False)),
+        )
+        for norm, nki_only in pairs:
+            if nki_only and getattr(norm, "use_fp8", False):
+                norm.use_fp8 = False
+                count += 1
+    return count

deepinfer/neuron_ops/static_moe.py

@@ -339,6 +339,15 @@ def _nki_moe_forward(
         acc = acc + _shared_slice(moe, x).float()
     regions.split()
     hier = getattr(moe, "_neuron_hier_ar", None)
+    flag_reduced = None
+    if capacity and replicated and size > 1:
+        # The ranks' overflow flags ride the output's all-reduce as one extra row (a separate 64-rank all-reduce
+        # cost ~3.4 ms a layer): every rank then sees the sum, so all of them take the same path.
+        flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=acc.device)).to(torch.float32)
+        acc = torch.cat([acc, flag.reshape(1, 1).expand(1, acc.shape[1])], dim=0)
+        out = _all_reduce(acc, group, hier)
+        out, flag_reduced = out[:-1], out[-1, 0]
+    else:
         out = _all_reduce(acc, group, hier) if replicated else _reduce_scatter(acc, group, size, rank)[:num_tokens]
     regions.split()
     if moe.has_shared and not sliced:
@@ -347,6 +356,8 @@ def _nki_moe_forward(
             shared, fmt="bf16", num_per_channels=None, clamp_value=moe.swiglu_clamp_value
         )
         out = out + nki_gemm.linear_qdq(shared, moe.shared_down_proj.weight).float()
+    if capacity and flag_reduced is not None:
+        return out.to(torch.bfloat16), flag_reduced > 0
     if capacity:
         flag = (overflow if overflow is not None else torch.zeros((), dtype=torch.bool, device=out.device)).to(torch.float32)
         if size > 1: