From 3c38fd5829a6458829b9a49d8d40ab1d788e51ee Mon Sep 17 00:00:00 2001 From: whn09 Date: Mon, 26 Oct 2026 06:36:48 +0800 Subject: [PATCH] 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 --- deepinfer/neuron_ops/nki/blockquant.py | 40 ++++++++++++++++++++++++++ deepinfer/neuron_ops/nki_gemm.py | 10 +++++++ deepinfer/neuron_ops/static_moe.py | 19 ++++++++++-- 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/deepinfer/neuron_ops/nki/blockquant.py b/deepinfer/neuron_ops/nki/blockquant.py index 1b46620..e87a786 100644 --- a/deepinfer/neuron_ops/nki/blockquant.py +++ b/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. diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index be1af38..cc3716b 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/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" diff --git a/deepinfer/neuron_ops/static_moe.py b/deepinfer/neuron_ops/static_moe.py index 85de881..0ecf76a 100644 --- a/deepinfer/neuron_ops/static_moe.py +++ b/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]