Fix MLAS LUT GEMM to support 2-bit MatMulNBits with N not divisible by 128#28528
Draft
vraspar wants to merge 1 commit into
Draft
Fix MLAS LUT GEMM to support 2-bit MatMulNBits with N not divisible by 128#28528vraspar wants to merge 1 commit into
vraspar wants to merge 1 commit into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes a CPU performance regression for 2-bit MatMulNBits where certain N sizes (e.g., N=192, divisible by 64 but not 128) incorrectly bypass the MLAS LUT GEMM (T-MAC) path and fall back to slower dequantization + FP32 GEMM.
Changes:
- Extend the LUT GEMM tile-size search space to include
bm=128, enabling valid tiling forNvalues like 192. - Relax LUT GEMM availability for 2-bit from requiring
N % 128 == 0toN % 64 == 0. - Add MLAS unit tests covering the newly supported
Nvalues (includingN=192andN=64).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/core/mlas/lib/qlutgemm.cpp | Enables LUT GEMM for 2-bit cases where N is 64-aligned (not necessarily 128-aligned) and adds bm=128 so the kernel config can tile such shapes. |
| onnxruntime/test/mlas/unittest/test_sqlutgemm.cpp | Adds unit tests targeting the previously failing shapes (e.g., N=192). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+221
to
+222
| count += RegisterSingleTest(1, 192, 64, with_threadpool, symmetric); | ||
| count += RegisterSingleTest(1, 64, 64, with_threadpool, symmetric); |
hariharans29
added a commit
that referenced
this pull request
Jun 17, 2026
### Description Today the only high-performance MLAS kernels for 2-bit weights are the LUT-based kernels. They under-utilize AVX-512 (e.g. they don't issue VNNI dot-product instructions), come with shape constraints (most notably `N` must be a multiple of 128), and — most importantly — under-perform on prefill regimes where the GEMM is compute-bound rather than weight-bandwidth-bound. Some of the shape limitations are being chipped away at incrementally (e.g. [#28528](#28528)), but the underlying performance gap remains. This PR addresses the gap directly by adding a "traditional" non-LUT 2-bit weight AVX-512 (+VNNI) GEMM kernel to MLAS, modeled on the existing 4-bit and 8-bit kernels in the same family. It supports `BlkLen` ∈ {32, 64, 128} and runs on any shape (no `N % 128 == 0` constraint). A real customer model heavily benefits from this change. The customer's model has 120 W2 `MatMulNBits` nodes. ~80 of them have shapes the LUT kernel supports (`N % 128 == 0`); the remaining ~40 do not. On `main`, those 40 nodes always fell through to the slow fp32 dequant + SGEMM fallback regardless of the LUT session option. This PR replaces that fallback with a native AVX-512(+VNNI) W2 kernel and adds a fast non-LUT path for all 120 nodes (the new kernel beats W2 LUT on the overlapping 80 nodes at prefill shapes and roughly matches it on decode shapes). ### Per-node routing | Branch | LUT mode | LUT-supported nodes (~80) | LUT-unsupported nodes (~40) | |---|---|---|---| | main | LUT off (default) | fp32 dequant + SGEMM | fp32 dequant + SGEMM | | main | LUT on | LUT kernel | fp32 dequant + SGEMM | | This PR | LUT off (default) | new W2 AVX-512(+VNNI) kernel | new W2 AVX-512(+VNNI) kernel | | This PR | LUT on | LUT kernel | new W2 AVX-512(+VNNI) kernel | ### End-to-end throughput (tokens/sec, higher is better) | Branch | LUT mode | seq=32 | seq=64 | seq=128 | |---|---|---:|---:|---:| | main | LUT on (80 LUT + 40 fp32 fallback) | 587 | 538 | 554 | | This PR | LUT on (80 LUT + 40 new kernel) | 681 | 722 | 764 | | This PR | LUT off (120 new kernel) | 1179 | 1338 | 1487 | | W4 baseline | n/a | 1245 | 1445 | 1647 | **Headline:** With this PR + the non-LUT path, W2 throughput is now within ~5–10% of W4. The remaining gap is the W2 "unpack tax" inherent to 2-bit weight kernels (more bytes-per-output-element to unpack vs. 4-bit). ## Behavior with `mlas.use_lut_gemm` opt-in User-facing dispatch precedence is unchanged. When the user opts into LUT via the session config, LUT is still preferred for the shapes where it is available (`MlasIsLutGemmAvailable` returns true). The new W2 native kernel takes over in two cases: * **Default sessions (no LUT opt-in)** — this is the headline behavior change. Previously these sessions had no fast 2-bit path on AVX-512 and would fall back to dequant + SGEMM; now they run the new native W2 kernel. * **LUT-enabled sessions on shapes where LUT is unavailable** (e.g. `N` not a multiple of 128). Previously these shapes hit the dequant+SGEMM fallback inside the LUT path; now they run the much faster native W2 kernel. Net effect: LUT-enabled users get a faster fallback path, default users get a fast default path, and there is no regression in either case. ## fp16 fallback path for 2-bit weights Semi-related addition in `onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc`: * `MatMulNBits<MLFloat16>::ComputeBUnpacked` previously only handled `nbits_ == 4`, with everything else falling into an `ORT_ENFORCE(nbits_ == 8)` branch — meaning fp16-activation 2-bit models would fail at runtime. Added the missing `nbits_ == 2` arm that calls `MlasDequantizeBlockwise<float, 2>` and then runs standard fp32 MatMul (same dequant-then-SGEMM pattern as the existing 4-bit/8-bit fp16 paths). * Added a `!prefer_lut_gemm_` guard in `MatMulNBits<T1>::PrePack` so we don't re-pack scales/zero-points in the W2 layout on top of an already-LUT-packed `packed_b_` buffer (which corrupts the LUT layout and crashes the LUT compute path). The guard is gated to `T1 == float` via the constructor, so the fp16 path is unaffected. ### Motivation and Context Adds a native AVX-512 (+VNNI) 2-bit weight CPU GEMM kernel to MLAS so that 2-bit quantized models have a competitive non-LUT path — closing the prefill perf gap versus 4-bit and removing LUT's shape-multiple constraints — and fixes a missing fp16 fallback so fp16-activation 2-bit models load and run on CPU.
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
Fixes a performance regression where 2-bit quantized models with MatMulNBits nodes having N=192 (or any N divisible by 64 but not 128) fall back to slow FP32 dequantization on CPU EP instead of using the MLAS LUT GEMM kernel.
Root cause: The LUT GEMM availability check (\MlasIsLutGemmAvailable) required \N % 128 == 0\ for 2-bit, and the tile-size search space (\�ms) had no value where \�m/2\ divides 192. This caused 40 out of 120 MatMulNBits nodes in the affected model to bypass the LUT kernel entirely.
Fix:
_div\ for 2-bit from 128 to 64 in \MlasIsLutGemmAvailable\
Note: The existing penalty heuristic may now select \�m=128\ for some shapes that previously used \�m=256\ (e.g., N=384, N=1024). This reflects better thread utilization. Benchmarking across representative shapes is recommended.
Motivation and Context
A 2-bit model benchmarked on AMD Ryzen AI 7 PRO 350 showed ~10x slower inference on CPU EP (225ms vs 21ms at seq=32) compared to the equivalent 4-bit model, because the LUT kernel rejected N=192 nodes.
Testing