Neuron MoE: skip unrouted experts inside the fp4 kernels at decode sizes (DEEPINFER_NEURON_GATED_EXPERTS, default on)

The fp4 expert GEMMs dequantize the whole weight to bf16 in HBM on every call
(device profile: the MoE graph reads 756 MB a layer), and a rank computes all
6 of its experts while a 6-row step routes to about one. fp4_linear_2core_gated
takes a per-expert flag (any row routed to it) and runs the dequant and GEMM in
an nl.fori_loop bounded by it, so an unrouted expert costs a zero fill. Decode
sizes only (<= 16 rows); prompt chunks route to every expert anyway.

bench70 (vllm-sonnet 512, TP=64, nextn=5): TPOT 258.9 -> 248.4 ms, step
738 -> 711 ms, tokens per step identical. Tiny cp=2 nextn=5 on device: 9/9
golden tokens, drafts unchanged.

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

deepinfer/neuron_ops/nki/blockquant.py

@@ -753,6 +753,46 @@ def fp4_linear_2core(x, w, scale):
 
 
 #: The low nibble of each uint16 lane of a uint32, i.e. two nibbles a lane-op.
+@nki.jit
+def fp4_linear_2core_gated(x, w, scale, flag):
+    """`fp4_linear_2core` that runs only when `flag` (int32 `[1, 1]`, 0 or 1) is set, else returns zeros.
+
+    For the routed experts of a decode step (DeepInfer, Neuron DSpark): a rank owns 6 experts of 384 and a 6-row step
+    routes to about one of them, so most expert GEMMs would dequantize a weight nobody uses. The flag bounds a dynamic
+    loop (`nl.fori_loop`), so the skipped expert costs a zero fill instead of the dequant and the GEMM.
+    """
+    m, k = x.shape
+    n = w.shape[0]
+    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
+    programs = nl.num_programs(0)
+    core = nl.program_id(0)
+    assert programs == 2, f"launch grid is {programs}; launch fp4_linear_2core_gated at grid 2"
+    assert n % (programs * PT) == 0, f"n={n} does not divide into {programs} shares of a multiple of {PT}"
+    span = n // programs
+    lo = core * span
+    _check(k, span)
+    mpad = -(-m // PT) * PT
+
+    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    zero = nl.ndarray((PT, span), dtype=nl.bfloat16, buffer=nl.sbuf)
+    nisa.memset(zero, 0.0)
+    for mb in nl.affine_range(mpad // PT):
+        nl.store(out[mb * PT : (mb + 1) * PT, lo : lo + span], value=zero)
+    flag_sb = nl.load(flag[0:1, 0:1], dtype=nl.int32)
+    count = nisa.register_alloc()
+    nisa.register_load(count, flag_sb)
+    xq = nl.ndarray((mpad, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
+    wd = nl.ndarray((span, k), dtype=nl.bfloat16, buffer=nl.private_hbm)
+
+    def body(_i):  # noqa: ANN001, ANN202
+        _stage_rows(x, m, k, xq, True)
+        _dequant_fp4(w[lo : lo + span, 0 : k // 2], scale[lo : lo + span], wd, span, k)
+        _gemm(xq, wd, out[0:mpad, lo : lo + span], mpad, k, span, out_scale=FP4_POST)
+
+    nl.fori_loop(0, count, body)  # the tracer frontend takes fori_loop, not dynamic_range
+    return out
+
+
 FP4_NIB_MASK2 = 0x000F000F
 #: `FP4_FIELD` doubled and signed: NKI takes a 32-bit bitwise operand as int32,
 #: the same way `E4M3_SGN2` is written.

deepinfer/neuron_ops/nki_gemm.py

@@ -120,5 +120,15 @@ 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_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)."""
+    scale = w._nki_scale[group]
+    weight = w._nki_weight[group]
+    if w.dtype != torch.int8 or not TWO_CORE or weight.shape[0] % 256 != 0:
+        return None
+    return _kernel("fp4_linear_2core_gated", 2)(x, weight, scale, flag)[: x.shape[0]]
+
+
 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"

deepinfer/neuron_ops/static_moe.py

@@ -33,6 +33,12 @@ _HOST = os.environ.get("DEEPINFER_NEURON_DEVICE", "neuron") == "cpu"
 _ALIGN = 16
 
 
+
+#: DEEPINFER_NEURON_GATED_EXPERTS: skip unrouted experts inside the fp4 kernels at decode sizes (dynamic loop on a
+#: per-expert flag, blockquant.fp4_linear_2core_gated).
+_GATED_EXPERTS = os.environ.get("DEEPINFER_NEURON_GATED_EXPERTS", "1") == "1"
+_GATED_MAX_ROWS = 16
+
 def _as_rows(t: torch.Tensor, rows: int) -> torch.Tensor:
     return t.contiguous().view(rows, -1).view(torch.uint8)
 
@@ -258,13 +264,22 @@ def _nki_moe_forward(
         acc, overflow = _grouped_experts(moe, x, local_idx, w, capacity)
     else:
         acc = torch.zeros((x.shape[0], moe.hidden_size), dtype=torch.float32, device=x.device)
+        # Decode steps: an expert no row routes to is skipped inside its kernels (a 6-row step routes to about one of
+        # a rank's 6 experts); prompt chunks route to all of them anyway.
+        gated = _GATED_EXPERTS and x.shape[0] <= _GATED_MAX_ROWS
         for e in range(num_local):
             w_e = torch.where(local_idx == e, w, torch.zeros_like(w)).sum(-1)
-            gate_up = nki_gemm.linear_qdq(x, moe.gate_up_proj_weight, group=e)
+            flag = (local_idx == e).any().to(torch.int32).reshape(1, 1) if gated else None
+            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)
             act = swiglu_forward_and_per_token_cast(
                 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()
+            down = nki_gemm.linear_qdq_gated(act, moe.down_proj_weight, e, flag) if gated else None
+            if down is None:
+                down = nki_gemm.linear_qdq(act, moe.down_proj_weight, group=e)
+            acc = acc + down.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]