Batch QMoE CPU QNBit expert GEMMs into grouped MLAS dispatches

[Tianlei Wu (tianleiwu)]

Description

Follow-up to #32644. That PR put block-wise int4/int8 QMoE CPU experts on the MLAS QNBit (MatMulNBits) kernels, but still issues one MlasQNBitGemmBatch call per active expert, so a layer costs 2 * num_active_experts thread-pool barriers.

MLAS already flattens a batch into a single parallel domain over (expert × N-block), so this batches the active experts instead: they are bucketed by token count and each bucket runs as one batched dispatch. Decode (one token per active expert) collapses to a single bucket, turning 2 * top_k dispatches per layer into 2.

This only applies where the expert loop is already serialized (fewer active experts than threads). With at least as many active experts as threads, one expert per thread remains faster and that path is unchanged.

Key changes

• Bucket the active experts by token count and run one batched MlasQNBitGemmBatch per bucket for FC1 and FC2, with SwiGLU applied across all gathered rows in between.
• fp32 bias is passed straight from the initializer; only fp16 bias needs a converted copy.
• Tests gain an intra-op thread-count knob so the same shapes can be driven through either dispatch path.

Motivation

CPU decode of MoE models is dominated by the per-expert GEMMs. The per-expert dispatch pattern is barrier-bound rather than bandwidth-bound: it sits near 30 GB/s regardless of expert count and gets worse above 32 threads, while the batched form scales to 68-126 GB/s.

Measured on a 96-core x64 host.

End-to-end on LFM2.5-8B-A1B int4 (22 QMoE layers, hidden 2048, moe_intermediate 1792, 32 experts, top-k 4) through onnxruntime-genai model_benchmark, 630 tokens per mode:

Metric	before	after
Token generation	36.5 ms/token (27.4 tok/s)	30.3 ms/token (33.0 tok/s)
Prompt processing (16 tok)	249 ms	142 ms

A single QMoE layer at those dimensions: 0.749 -> 0.438 ms at 32 threads.

The gain scales with threads / active_experts, since it is barrier cost that is being removed. At top-k 4 decode: 1.20x at 8 threads, 1.38x at 16, 1.78x at 32, 2.11x at 64. With few threads relative to active experts the expert-parallel path is selected instead and nothing changes.

Testing

• MoETest.* passes (24 tests).
• New *_PerExpertDispatch and *_SingleThread cases cap the intra-op pool so the same shapes and reference exercise the per-expert path and a single-threaded pool, checking the two dispatch paths against each other.
• Verified with a temporary probe that the grouped cases really take the new path and the PerExpertDispatch cases really do not.
• Greedy generation on LFM2.5-8B-A1B produced identical token IDs and identical text before and after.

Note on reproducibility

Bucketing sorts experts by token count, which changes the cross-expert accumulation order. Decode output is bit-identical; with multiple tokens per expert results can differ at float rounding level (observed in the 8th significant digit at 16 tokens).

[Copilot]

🟡 Changes recommended

Grouped execution can allocate hundreds of MiB of additional scratch memory for long prompts, risking OOM.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Batches active CPU QMoE QNBit expert GEMMs to reduce thread-pool barrier overhead.

Changes:

• Groups experts by token count for batched FC1/FC2 MLAS dispatches.
• Adds per-session thread controls and dispatch-path tests.
• Reuses fp32 bias data while converting fp16 bias once.

File summaries

File	Description
moe_quantization_cpu.cc	Implements grouped QNBit expert execution.
moe_test.cc	Adds thread-count controls and path coverage.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

Single-thread behavior contradicts the stated dispatch policy, and multi-bucket coverage is hardware-dependent.

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc:1665

• This predicate also selects grouped execution when the session has one intra-op thread and multiple active experts, because num_expert_threads is then 1 even though num_active_experts >= max_expert_threads. That contradicts the PR's stated invariant that configurations with at least as many active experts as threads are unchanged, and it adds route-wide staging and changes accumulation order for single-thread sessions. Either gate this on num_active_experts < max_expert_threads or document the broader behavior as intentional.
  onnxruntime/contrib_ops/cpu/moe/moe_quantization_cpu.cc:1737
• The new multi-bucket path is not exercised deterministically. The existing prefill case uses the default pool, so a host with 2–4 intra-op threads takes the per-expert path, while the newly pinned prefill test explicitly forces that same path with two threads. Add a prefill case with an intra-op count greater than num_experts so bucket construction and reuse across distinct token counts are always covered.

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Balanced