MoE Expert FFN GEMM on the NPU — Milestone M0 viability (Qwen3-A3B) ✅ VIABLE

Date: 2026-06-26 Status: ✅ VIABLE — milestone M0 (Branch A, NPU-centric) of the MoE-on-NPU saturation build. The per-expert FFN GEMM on the NPU beats single-thread CPU NEON at every Qwen3-A3B expert shape and every per-expert M∈{8..128} — 3–10× at M=8, no crossover, bit-exact vs the CPU int8 reference. Tools: tools/moe_expert_probe.c (1840abb) · src/npu.c (ORK_DIRECT_I4 direct int4→int8 inflate 570cb68; cacheable 0x403 fill; chan_scales_f32 NULL-guard 980cea9). RK3588 10.3.0.236, governors pinned.

Question

The prior MoE Expert Offload work concluded the integrated mul_mat_id path was ~7× under CPU on prefill (Tier 5c), walled by the CPU-side NEON f32→int8 quant+tile repack — NOT the NPU GEMM. So the open question for the saturation build was narrower and more fundamental:

Isolating just the per-expert FFN GEMM (the work that actually lands on the NPU), and using the new materialization primitives — the direct int4→int8 tiled inflate (ORK_DIRECT_I4, 570cb68) + cacheable 0x403 fill — is computing a sparse-MoE expert on the NPU competitive with the CPU per-expert?

This is the matmul-level go/no-go that gates the whole NPU-centric branch. The earlier ~7× was an integration number (Q4_K→int8 repack churn dominated); M0 measures the primitive itself.

Method

tools/moe_expert_probe.c times one expert's three projections at the real Qwen3-A3B shapes, sweeping per-expert M (rows routed to that expert):

• gate / up: K=2048 × N=768
• down: K=768 × N=2048
• (+ the wider 4096 / 1408 variants for the larger A3B configs)
• per-expert M ∈ {8, 16, 32, 64, 128}

NPU path: weights inflated directly int4→int8 into the resident dma_buf (ORK_DIRECT_I4, no f32 round-trip, cacheable 0x403), then run_i8 / run_stream_i8. CPU reference: single-thread NEON int8 GEMM. Correctness checked against the CPU int8 reference (bit-exact).

Result — VIABLE at every shape and every M

• NPU beats CPU at every expert shape and every per-expert M∈{8..128} — 3–10× faster at M=8, no crossover anywhere in the sweep (the small-M end is where the submit floor would have bitten; it doesn't, because the GEMM is the only cost left).
• The direct int4→int8 inflate cut per-expert weight materialization 16–30× vs the old nibble→f32→tile_f32_i8 repack — the cost collapses to the GEMM itself. This is exactly the primitive the Tier 5c "new wall" (the NEON quant+tile being 77% of MoE-chain time) was waiting on: the int4-resident expert never needs a Q4_K→f32→int8 round-trip.
• Bit-exact vs the CPU int8 reference across the sweep.

Dispatch finding (which primitive per projection)

• run_i8 (single submit) is already multi-core — the auto-tuner spreads the N-tiles across the 3 cores. run_chain_i8 is single-core, so chaining experts that way is slower per expert.
• → use run_stream_i8 for the chainable gate/up (they satisfy the chain envelope: K%512==0, K≤4096), and run_i8 for down (K=768 is fine; it gets the multi-core N-tile spread for free).
• Fixed a chan_scales_f32 NULL-guard along the way (980cea9).

Caveat (honest framing of the CPU baseline)

The CPU reference here is single-thread NEON. A 4-thread CPU GEMM narrows the M=8 margin to ~1.5–2.5× — NPU still wins, and crucially it frees all 4 CPU cores (the whole point of the saturation goal: CPU does prefetch/expand/non-matmul ops while the NPU does the expert matmuls). So even where the raw GEMM margin is modest, the system-level win (4 cores reclaimed) is the payoff. This is the phase-aware saturation thesis at the matmul level.

Verdict

✅ VIABLE. The per-expert FFN GEMM on the NPU is competitive-to-dominant (3–10× at M=8 vs 1-thread CPU; ~1.5–2.5× vs 4-thread + frees the CPU), bit-exact, with the direct inflate collapsing materialization 16–30×. This flips the matmul-level go/no-go that the integrated ~7× (Tier 5c) could not — the wall there was the CPU-side Q4_K repack, not the NPU. M0 is done for Branch A.

Follow-on (the M1a gap)

The mixed-.orkpack convert (Layer-Streaming) does NOT yet persist the routed _exps / MUL_MAT_ID expert tensors — only the backbone (incl. the router + the dense shared-expert) currently persists. Persisting the routed experts is the M1a implementation target (see the saturation milestone thread on the roadmap).