perf: run block-wise int4/int8 QMoE CPU experts on the MLAS QNBit GEMM kernels - #32644
Open
Yuri Khrustalev (ykhrustalev) wants to merge 5 commits into
Open
Yuri Khrustalev (ykhrustalev) wants to merge 5 commits into
Yuri Khrustalev (ykhrustalev) wants to merge 5 commits into
Conversation
…M kernels The CPU QMoE kernel dequantized block-wise experts to fp32 and ran SGEMM, one expert per thread, so decode (top_k active experts) left most of the pool idle. Experts are now pre-packed once for the MatMulNBits kernels (no re-quantization) and each expert GEMM runs on them, sequentially with the whole pool when fewer experts are active than threads. ORT_QMOE_CPU_QNBIT_GEMM selects fp32 (default) or int8 activations, or disables the path.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Copilot started reviewing on behalf of
Yuri Khrustalev (ykhrustalev)
September 16, 2026 15:09
View session
Contributor
There was a problem hiding this comment.
🟡 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.
| const size_t nbits = static_cast<size_t>(expert_weight_bits_); | ||
| const size_t blk = static_cast<size_t>(block_size_); | ||
| const size_t per_expert = packed.packed_size_per_expert; | ||
| const size_t total_packed_size = per_expert * static_cast<size_t>(num_experts); |
Comment on lines
+966
to
+970
| } else if (allow_int8_compute && MlasIsQNBitGemmAvailable(nbits, blk, SQNBIT_CompInt8)) { | ||
| qnbit_compute_type_ = SQNBIT_CompInt8; | ||
| } else { | ||
| use_qnbit_gemm_ = false; | ||
| } |
| // Experts pack independently into disjoint regions, so spread them over a load-time pool | ||
| // (the session pool is not reachable from PrePack; same approach as MatMulNBits::PrePack). | ||
| OrtThreadPoolParams pack_tp_params; | ||
| pack_tp_params.thread_pool_size = Env::Default().GetNumPhysicalCpuCores(); |
Make int8 activations opt-in (ORT_QMOE_CPU_QNBIT_GEMM=int8), matching MatMulNBits without accuracy_level=4; 8-bit experts keep the dequantize path by default. Record the compute type in the tagged prepacked shape buffer and reject shared buffers packed for another compute type. Share one shape-buffer writer across the three prepack layouts, return the packed size from the eligibility check instead of recomputing it, pass fp32 scales to the GEMM lambda by expert index, and only spin up the load-time pack pool when there is more than one expert. Split the 8-bit test into default (dequantize) and int8-activation cases. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ng on the CPU QNBit path Add an accuracy_level attribute to QMoE with the MatMulNBits meaning so a model can opt into int8 activations (4) instead of relying on the ORT_QMOE_CPU_QNBIT_GEMM environment variable, which now only overrides the attribute. Pack constant block-wise zero points into the QNBit layout (QMoE's per-expert zero point layout is the MatMulNBits one) so asymmetric checkpoints take the same kernels. Log once when a block-wise node cannot use the path and why, since the dequantize fallback is an order of magnitude slower for decode. Record the zero point flag next to the compute type in the prepacked shape buffer, and cover accuracy_level and zero points (4-bit fp32 and int8, 8-bit int8) in the tests. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…Bit prepack Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
mirounga
self-requested a review
September 16, 2026 16:58
kunal-vaishnavi
requested review from
Akshay Sonawane (apsonawane),
Edward Chen (edgchen1),
kunal-vaishnavi and
Tianlei Wu (tianleiwu)
September 16, 2026 18:48
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Problem
The CPU QMoE kernel dequantizes block-wise int4/int8 experts to fp32 and runs SGEMM, threading one expert per core. For decode only
top_kexperts are active, so most of the pool sits idle (3 tok/s on an 8B-A1B model). The existing MLAS Q4 fast path is AVX-512 only, off by default, and re-quantizes the weights to a different grid.Solution
MatMulNBitskernels) at load time, oneMlasQNBitGemmPackQuantBDataper expert, with no re-quantization and no unpacked fp32/uint8 copies. QMoE's per-expert zero point layout is theMatMulNBitsone, so asymmetric checkpoints take the same kernelsMlasQNBitGemmBatch; when fewer experts are active than there are threads, experts run sequentially and each GEMM gets the whole poolaccuracy_levelattribute toQMoEwith theMatMulNBitsmeaning: the default keeps fp32 activations where MLAS has that kernel (4-bit; 8-bit has no fp32 variant and keeps the dequantize path),4allows int8 activations.ORT_QMOE_CPU_QNBIT_GEMM(fp32,int8,0) overrides the attributeTesting
New
MoETest.QMoETest_CPU_Int4_BlockWise_*/QMoETest_CPU_Int8_BlockWise_*cases with an fp32 reference: decode, prefill, bias + block 64, fp16,accuracy_level=4, zero points (4-bit fp32 and int8, 8-bit int8), and two-session shared pre-packed weights. AllMoETest.*andMatMulNBits*cases pass inonnxruntime_provider_teston Linux x64 (AVX-512).LFM2.5-8B-A1B (32 experts, top-4, int4 block 32, signed-scale export from feat: add LFM2-MoE support onnxruntime-genai#2575), Xeon 8358 x30, same build for all rows (
0= previous dequantize path). Decode is genai greedy generation; prefill is a singleInferenceSession.run; KL is against HF fp32 logits on an 81-token chat-template prompt:0(before)int8(accuracy_level=4)Earlier int8-only measurements on an M3 Ultra: decode 4 → 125 tok/s, prefill 256 tokens 190 → 855 tok/s, KL 0.112 → 0.101. KleidiAI only serves the int8 compute type, so the fp32 default on Apple Silicon uses the NEON fp32 kernels and has not been measured yet.
Motivation and Context
Makes block-wise quantized MoE models usable for CPU decode by reusing the
MatMulNBitskernel set instead of dequantizing per call. fp32 activations are the default because int8 is measurably lossier;accuracy_level=4opts in per model, matchingMatMulNBits. Pairs with microsoft/onnxruntime-genai#2575, which exports LFM2-MoE models with this encoding.