onnxruntime
4b728dd1 - Batch QMoE CPU QNBit expert GEMMs into grouped MLAS dispatches (#32668)

Commit
58 days ago
Batch QMoE CPU QNBit expert GEMMs into grouped MLAS dispatches (#32668) ### 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).
Author
Parents
Loading