perf: run block-wise int4/int8 QMoE CPU experts on the MLAS QNBit GEMM kernels

[Yuri Khrustalev (ykhrustalev)]

Description

Problem
On CPU, QMoE runs block-wise 4-bit and 8-bit experts slowly. Before every matrix multiply it converts the expert weights back to fp32, and it gives each expert its own thread. During decode only top_k experts run, so most threads sit idle. On an 8B-A1B model that means about 3 tokens/s. The existing MLAS Q4 fast path doesn't solve this: it only runs on AVX-512, is off by default, and re-quantizes the weights.

Solution

• Pack weights once at load time. Block-wise 4-bit and 8-bit experts are packed for the same MLAS kernels that MatMulNBits uses. The weights aren't re-quantized, and no fp32 copy is kept. Zero points work too, because QMoE stores them the same way MatMulNBits does.
• Use the threads better. When fewer experts are active than there are threads, experts run one after another and each matrix multiply gets the whole thread pool.
• New accuracy_level attribute on QMoE, with the same meaning as in MatMulNBits:
  • 4-bit: the default keeps fp32 activations. accuracy_level=4 switches to int8 activations, which is faster but slightly less accurate.
  • 8-bit: MLAS has no fp32 kernel, so int8 activations are used at every level.
• Environment override: ORT_QMOE_CPU_QNBIT_GEMM takes precedence over the attribute.
  • fp32 forces fp32 activations, so 8-bit experts go back to the old path.
  • int8 forces int8 activations.
  • 0 turns the new path off.
  • Any other value is an error.
• Shared pre-packed weights across sessions. The packed buffer is tagged with how it was packed. A buffer packed differently is rejected instead of being used.
• Fallback warning. If a block-wise node can't use the new path, a warning with the reason is logged once, because the old path is about 10x slower for decode.
• Unchanged: row-wise scales, 2-bit weights, and other execution providers.

Testing

• New MoETest.QMoETest_CPU_Int4_BlockWise_* and QMoETest_CPU_Int8_BlockWise_* tests compare against an fp32 reference. They cover decode, prefill, bias with block size 64, fp16, accuracy_level=4, zero points, each environment variable setting, and shared pre-packed weights across two sessions. All MoETest.* and MatMulNBits* tests pass in onnxruntime_provider_test on Linux x64 (AVX-512).

• Model benchmark:

  • Model: LFM2.5-8B-A1B (32 experts, top-4, int4 with block size 32), exported with feat: add LFM2-MoE support onnxruntime-genai#2575.
  • Machine: Xeon 8358, 30 threads, same build for every row.
  • Decode: genai greedy generation.
  • Prefill: a single InferenceSession.run.
  • Accuracy: KL divergence against the Hugging Face fp32 logits on an 81-token prompt.

  ORT_QMOE_CPU_QNBIT_GEMM	decode tok/s	prefill 256 tok/s	prefill 16 tok/s	KL vs HF fp32	top-1 agree
  0 (old path)	3.2	147	27	0.096	0.901
  unset (fp32, the default)	55.2	346	148	0.098	0.889
  int8 (accuracy_level=4)	60.7	364	103	0.106	0.889

  An earlier int8-only run on an M3 Ultra gave:

  • Decode: 4 → 125 tok/s.
  • Prefill (256 tokens): 190 → 855 tok/s.
  • KL divergence: 0.112 → 0.101.

  The fp32 default hasn't been measured on Apple Silicon yet. There it uses the NEON fp32 kernels, because KleidiAI only provides int8.

Motivation and Context

This makes block-wise quantized MoE models fast enough for CPU decode by reusing the MatMulNBits kernels instead of converting weights back to fp32 on every call. fp32 activations stay the default where MLAS supports them because int8 loses a little accuracy. Models can opt in to int8 with accuracy_level=4, as with MatMulNBits. This pairs with microsoft/onnxruntime-genai#2575, which exports LFM2-MoE models in this format.

🤖 Generated with Claude Code

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

The 8-bit default contradicts the documented behavior, and packed-buffer sizing can overflow.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds an MLAS QNBit GEMM fast path for block-wise quantized CPU QMoE experts.

Changes:

• Pre-packs int4/int8 expert weights.
• Adds QNBit execution, threading, and sharing support.
• Adds CPU correctness and sharing tests.

File summaries

File	Description
moe_quantization_cpu.h	Declares QNBit packed state and helpers.
moe_quantization_cpu.cc	Implements packing, execution, policy, and sharing.
moe_test.cc	Adds block-wise QMoE tests.

Review details

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

---

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

[Yuri Khrustalev (ykhrustalev)]

Thanks — addressed all three.

• 8-bit default: correct, MLAS has no SQNBIT_CompFp32 variant for 8 bits (only SQ8BitGemmVariant_CompInt8), so the old policy silently kept 8-bit experts on dequantize+SGEMM. The selection now falls back to int8 when the bit width has no fp32 kernel and fp32 was not explicitly forced; ORT_QMOE_CPU_QNBIT_GEMM=fp32 still pins the dequantize path. QMoETest_CPU_Int8_BlockWise_SwiGLU now exercises the QNBit path (its tolerance moved to 0.05, matching the other int8-activation cases — the previous 0.01 was exceeded by 0.0003), and a new ..._ForcedFp32FallsBackToDequantize case pins the fallback at the tighter tolerance.
• Packed-buffer sizing: now SafeInt<size_t>, matching the packed-cache path.
• Pack pool: sized by min(num_experts, physical cores), so it is skipped entirely for a single expert.

All 17 MoETest.* cases pass on macOS arm64 (M3 Ultra).

[Tianlei Wu (tianleiwu)]

Reviewed at d026ecb0. The QNBit packing, compute, zero-point, and shared-buffer plumbing follows the established MatMulNBits contracts. The earlier packed-size overflow, int8-default routing, and pack-thread-count concerns are fixed at this head. I left two non-blocking inline suggestions for route coverage and comment accuracy.

[Tianlei Wu (tianleiwu)]

Reviewed at 61173ff78c76a4d2072c2555e8b8dde297a8f830. Requesting changes because the new accuracy_level=4 policy still selects fp32 when both 4-bit kernels are available, so the per-model int8 opt-in does not work as documented. A second inline comment is a non-blocking suggestion to keep per-call QNBit scratch arena-backed.

The packing/finalization, shared-buffer metadata, bias handling, and non-nested thread-pool usage otherwise follow the existing contracts. The earlier 8-bit route-coverage and zero-point-comment suggestions are addressed at this head.

Validation: source review and a source-verified dispatch-policy model; no local C++ tests or benchmarks were run because no provider-test executable was available. CI was still running at review time.

[Yuri Khrustalev (ykhrustalev)]

Thanks, both addressed in acebeebc8d:

• accuracy_level=4: int8 is now selected ahead of fp32 when available, fp32 remains the fallback, and the explicit fp32 override still pins it. QMoETest_CPU_Int4_BlockWise_SwiGLU_AccuracyLevel4SelectsInt8 asserts level-4 outputs are bitwise identical to forced int8 and differ from forced fp32; it fails against the previous fp32-first ordering.
• Arena: the QNBit scratch allocation no longer uses reserve.

All MoETest.* pass locally on macOS arm64.

[Tianlei Wu (tianleiwu)]

Reviewed at bc5f446b. The earlier findings on int8 route proof, zero-point documentation, accuracy_level=4 selection, and arena-backed QNBit workspace are addressed. The packing, shared-buffer compatibility, overflow checks, and non-nested thread-pool use follow the established MatMulNBits/MLAS contracts, and I found no new actionable issues.

Validation: rebuilt onnxruntime_provider_test from this head and ran all 16 QMoETest_CPU_Int4_BlockWise* / QMoETest_CPU_Int8_BlockWise* tests; all passed.