MoE Expert Offload to NPU (GGML_OP_MUL_MAT_ID) — Qwen3-35B-A3B

Date: 2026-06-23 (updated 2026-06-24) Status: ⚖️ Not yet viable, but closing — handler built & routing solved; offload was a heavy net loss as first built (~80× prefill, ~60× decode), since narrowed to ~7.2× under CPU on prefill by chaining + a NEON repack + a fused Q4_K→int8 dequant (see the 2026-06-24 follow-up at the bottom). Experts stay on CPU (board default, ORK_MOE_NPU off). Not closed — the gap is still closing. Model: Qwen3.6-35B-A3B (Q4_K_XL GGUF, ~21GB) on RK3588. Baseline MoE decode ~6.3 t/s.

Problem

ggml-ork had no GGML_OP_MUL_MAT_ID handler — supports_op only handled MUL_MAT. So all MoE expert matmuls run on CPU (dense attn goes to NPU; experts don't). Trace: decode shows only mul_mat_i8 (attn), zero mul_mat_id.

Decode handler (built, in ggml-ork.cpp)

ggml_backend_ork_mul_mat_id_i8 (decode, src1->ne[2]==1): quantize the shared token activation once; gather the n_used selected experts (dequant Q4_K→int8, per-expert cached in wcache); batch them with ork_mm_run_chain_i8 (S=n_used, different weights, same activation — the validated chain primitive); scatter+dequant to dst[N,n_used]. Plus supports_op + offload_op MUL_MAT_ID cases. Compiles; supports_op accepts the decode expert ops.

Blocker (root-caused)

The experts never route to ork. Instrumenting ggml_backend_sched_backend_id_from_cur showed: ffn_*_exps.weight usage=WEIGHTS host=0 → from_buffer=CPU. The expert weights live in a non-host buffer — llama.cpp's CPU weight-repack buffer type (GGML_USE_CPU_REPACK) claims+repacks the Q4_K experts. ork's supports_buft only accepts host buffers (and the scheduler's offload fallback also requires is_host), so ork can't claim them. (Dense attn weights stay host/non-repacked → attn offloads fine.) The handler also needs the standard Q4_K layout for its to_float dequant.

Disabling repack globally (GGML_NO_REPACK=1) backfired: broke the working attn offload (680 ioctl failures, decode → 0.77 t/s), experts still didn't route. Repack is load-bearing; the fix must be surgical (exclude only the expert weights from repack).

Surgical fix — DONE; experts now run on NPU, but hit the IOMMU wall

The full routing chain was solved:

1. Repack exclusion (the fix): gate the MUL_MAT_ID branch of repack.cpp's buft supports_op behind !getenv("ORK_NO_EXPERT_REPACK"). With ORK_NO_EXPERT_REPACK=1 experts land in the standard host buffer → ork claims them. (Global GGML_NO_REPACK was wrong — broke attn.)
2. supports_op must accept MUL_MAT_ID for all n_tokens — the graph split is planned with a worst-case multi-token batch; the decode-only gate left experts on CPU. Handler loops over tokens.
3. Handler uses per-expert ork_mm_run_i8 (not run_chain_i8): pack_i8 splits K at 1024, so K=2048 gate/up experts are 2-slice (Sk=2), which the chain path rejects; run_i8 handles K-split.

Result: mul_mat_id_i8 fires (345×) — experts genuinely run on the NPU. But decode fails on NPU IOMMU exhaustion: the handler keeps every touched expert packed NPU-resident in wcache; at ~2067 resident packed weights bcreate returns CREATE: Bad address → decode res=-3 + soft-resets. ~18k expert weights, no eviction → all-resident doesn't scale.

✅ Pool fix works (IOMMU solved) — but ❌ decode offload is a net LOSS

Added ork_mm_repack_i8 (ork-driver main, commit ork_mm_repack_i8 …) — re-tiles int8 into an existing ork_w's DMA (no alloc/free). The handler now uses a fixed per-shape pool (ORK_MOE_CACHE slots/shape, default 384) reused round-robin via repack-in-place. Result: mul_mat_id_i8 fires 1936×, PACK FAIL 0, decode-fail 0 — the IOMMU exhaustion is GONE; MoE experts run on the NPU correctly.

BUT decode speed = tg4 = 0.11 t/s vs the 6.3 t/s CPU baseline — ~60× slower. Cause: per-expert run_i8 is ~1152 submits/token (8 experts × 3 projs × 48 layers; the submit floor dominates) + repack churn for cold experts. Conclusion: MoE expert offload to the NPU is a net LOSS for DECODE (which is bandwidth-bound regardless) — confirming the design intent that decode belongs on CPU. The infra (routing, repack, pool) is built and correct; it just shouldn't drive decode.

❌ Prefill group-by-expert measured — also a net LOSS (the submit floor is per-expert, not per-token)

Implemented the group-tokens-by-expert prefill path in ggml_backend_ork_mul_mat_id_i8: bucket every (token t, slot j) pair by expert id ids[j,t], then run ONE M=count matmul per active expert (M-padded to 32 like the dense path, zero-padded gather rows), dequant-scatter back to dst[N,n_used,n_tokens]. Pre-quantize each token's input once (broadcast n_b1==1). Handler now accepts any n_tokens; decode (M=1) degenerates to one matmul per selected expert.

Measured (RK3588, 10.3.0.236, -t 4, governors pinned, Qwen3.6-35B-A3B Q4_K_XL):

path	pp128
CPU baseline (experts on CPU, default)	22.59 t/s
NPU expert offload (ORK_NO_EXPERT_REPACK=1, group-by-expert)	~0.28 t/s (≈80× slower)

Offloading is correct (302 mul_mat_id_i8 calls, 0 PACK/run failures) — it's just ~80× slower.

Why grouping can't rescue it: the submit count is bounded by active experts, not tokens. A MoE forward touches ≈ all 128 experts × 3 projections × 48 layers ≈ 18k expert-matmul submits per batch, roughly batch-size-independent once every expert is hit (which happens by ~pp64). At the ~365µs submit floor that's ~6.7s of pure submit overhead per batch, before compute — plus Q4_K→int8 repack churn (cap=384/shape, but each layer routes 128 distinct experts so the pool recycles constantly). Grouping cut submits ~8× vs per-token (147k→18k for pp128) but 18k is still far too many at this point. The lever that helps is chaining experts within a layer into one submit — which was subsequently built (see the 2026-06-24 follow-up below); this initial group-by-expert measurement was ~80× off CPU.

VERDICT (2026-06-23, since revised — see follow-up): not yet viable as first built

As initially built (group-by-expert, per-expert submit), MoE-on-NPU was ~80× under CPU on prefill / ~60× on decode — the submit floor × (n_experts × n_proj × n_layers) dominated. Experts stay on CPU (board default, ORK_MOE_NPU off). The infra (routing, ork_mm_repack_i8 pool, group-by-expert handler) is built, correct, and preserved at ~/ggml-ork.moe-wip.cpp + scratchpad. The NPU's real value remains dense prefill matmuls (attn), which already offload and win. This was NOT the end of the line — the 2026-06-24 work below (chaining + NEON repack + fused dequant) closed the prefill gap to ~7.2×.

Original blocker (now resolved by the pool): NPU IOMMU memory management for experts

Keeping all touched experts resident exhausts the IOMMU (~2067 packed weights → CREATE: Bad address). LRU eviction (free + re-pack) was tried and made it WORSE — cap=1024 PACK-failed at wcache=1064 (vs 2067 with no eviction): ork_w_free+re-pack fragments the IOMMU (freed slots aren't reused), so it exhausts sooner; smaller caps churn more = worse.

Proper fix: a fixed, pre-allocated pool of reusable expert-weight DMA buffers — allocate a budgeted count once and repack in place round-robin (zero free/realloc → no fragmentation). Needs a new ork-driver API (e.g. ork_mm_repack_i8(ork_w*, K, N, B) that overwrites an existing packed weight's buffer instead of bcreate-ing a new one). Substantial ork-driver work — the next task.

Status: experts genuinely run on NPU (mul_mat_id_i8 339×); blocked on the IOMMU buffer pool. All changes env-gated/inert by default (board default = ~6.3 t/s baseline). WIP at ~/ggml-ork.moe-wip.cpp.

TODO — Prefill MoE offload — DONE & MEASURED (net loss, see verdict above)

The prefill group-by-expert path was implemented (n_tokens>1, M=count matmul per active expert) and benchmarked: ~0.28 t/s vs 22.59 CPU. It does NOT pay off — the submit floor scales with the number of active experts (≈ all of them per batch), not the batch size, so prefill hits the same ~18k-submit wall as decode. Closed; experts stay on CPU.

ROI note

Decode (dense or MoE) on RK3588 is DDR-bandwidth-bound — the NPU shares the same DDR, so offloading matmuls doesn't break the wall (~10–15% decode ceiling). The NPU's real value is prefill. See Exp 2026-06-23 Decode Profiling and Exp 2026-06-23 Chaining Resolution (the run_chain primitive).

---

FOLLOW-UP (2026-06-24): chaining + repack work closed prefill from ~80× to ~7.2× under CPU

The "~80× net loss, closing this line" verdict above was for the first implementation. Three subsequent levers (all gated ORK_MOE_NPU, experts still CPU-default) walked the prefill deficit down by >10×. All on Qwen3.6-35B-A3B Q4_K, RK3588, -t 4, pp128:

step	pp128 (NPU MoE)	vs CPU	what changed
group-by-expert (per-expert submit)	0.28 t/s	~80×	initial
+ expert chaining (run_chain_i8 Sk=2 via Bf, 1bfeed3/3f5ae01; ggml-ork chains a projection's experts in one submit)	0.97 t/s	~23×	submit floor 18k→144 ioctls/fwd
+ NEON-fused repack (ork_mm_pack_i8_f32/repack_i8_f32, 2b0ee20)	2.90 t/s	~7.9×	repack 95%→84%; the cache-hostile strided transpose store was the bulk
+ fused Q4_K→int8 dequant (ork_mm_pack_i8_dequant/repack_i8_dequant, b520ee7; ggml-ork b5bc5ce8b)	3.13 t/s	~7.2×	eliminates the full f32[N][K] staging round-trip (per-channel callback dequants into a reused K-scratch + NEON quant+tile in place)

(run_stream_i8 in the MoE path adds a further +5% → 3.28; kept in scratchpad, not committed.)

Key correction to the original verdict: 5b had concluded "even a free repack still loses" — that was based on the old strided repack and is now obsolete. A faster repack did close the gap substantially. New wall (profiled after the fused dequant): the Q4_K→f32 dequant is now 0 ms (folded into the pack), and pack/repack is still 77% of MoE-chain time but entirely the NEON f32→int8 quant+tile (~84 s). So the next lever is a cheaper int8 quant+tile (or keeping weights int8-resident to skip re-quant), NOT more dequant work. The earlier "repack→0 ⇒ CPU parity" extrapolation was too optimistic: the dequant was cheap; the quant+tile is the real cost. Verdict: still a net loss (experts stay on CPU), but not closed — the gap is ~7× and closing. The chain + repack primitives are general wins on main regardless.

Benchmark gotcha (cost two bench rounds): ORK_MOE_NPU=0 still enables NPU MoE — the gate is a getenv-presence check, not a value check. For a CPU baseline, unset the var (env -u ORK_MOE_NPU); use ORK_OFF=1 to force everything to CPU. Setting it to 0 silently ran NPU MoE in all configs (~3.1 t/s across the board) until the profile counters revealed MoE-chain was firing.