[MLAS] Add AVX512 (+VNNI) 2-bit weight CPU kernels by hariharans29 · Pull Request #29064 · microsoft/onnxruntime

hariharans29 · 2026-06-16T00:25:31Z

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), 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.

Copilot

Pull request overview

This PR adds a native AVX-512 (+VNNI when available) 2-bit-weight SQNBit GEMM path to MLAS (including packing + kernel dispatch) to avoid the dequant+SGEMM fallback for many MatMulNBits (W2) shapes, and fixes a missing fp16-activation W2 fallback path in the MatMulNBits contrib op.

Changes:

Integrate a new 2-bit SQNBit CompInt8 backend into MLAS: packed-B sizing/packing + AVX-512/VNNI kernel dispatch, with BlkLen support for 32/64/128.
Add/extend benchmarks and unit/operator tests to cover the new W2 paths and customer-representative shapes.
Fix MatMulNBits fp16-activation 2-bit fallback (ComputeBUnpacked) and prevent LUT-packed buffers from being corrupted by non-LUT repacking.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit.cpp	New unit tests validating W2 block-group pack/unpack helpers for BlkLen 32/64/128.
onnxruntime/test/mlas/bench/bench_qnbitgemm.cpp	Adds W2-focused benchmark argument grids and customer-shape benchmark coverage.
onnxruntime/test/mlas/bench/bench_lutgemm.cpp	Extends LUT GEMM compute bench to also measure dequant+SGEMM fallback when LUT is unavailable for a shape.
onnxruntime/test/contrib_ops/matmul_2bits_test.cc	Adds operator-level tests for fp16-activation W2 fallback and broad W2 coverage across BlkLen/accuracy grids.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512vnni.cpp	Wires W2 pack+kernel dispatch into the AVX-512-VNNI SQNBit dispatch table (via a BlkLen forwarder).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512.cpp	Wires W2 pack+kernel dispatch into the AVX-512BW SQNBit dispatch table (via a BlkLen forwarder).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit.h	New W2 pack/layout helpers and reference utilities for the AVX-512 block-group packed format.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit.cpp	New scalar pack implementation + scalar reference kernel for W2 block-group layout (including K-padding behavior).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit_blklen64.h	New AVX-512 W2 SIMD kernel implementation for BlkLen=64 (incl. K-tail/N-tail/odd-M handling + BlkSum correction).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit_blklen32.h	New AVX-512 W2 SIMD kernel implementation for BlkLen=32.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx512_2bit_blklen128.h	New AVX-512 W2 SIMD kernel implementation for BlkLen=128.
onnxruntime/core/mlas/lib/qnbitgemm.h	Extends MLAS_QNBIT_GEMM_DISPATCH with 2-bit CompInt8 dispatch surface; adjusts packed-B workspace sizing for W2 K-padding.
onnxruntime/core/mlas/lib/qnbitgemm.cpp	Adds SQ2BitGemm CompInt8 variant selection/availability/packing/compute integration (incl. effective K-block stride hook).
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Fixes fp16-activation W2 fallback dequant path and guards LUT vs non-LUT repacking interactions.
cmake/onnxruntime_mlas.cmake	Adds new W2 AVX-512 source files to the MLAS build.

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.

Copilot

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

tianleiwu

I found one availability-contract issue that should be fixed before this lands. The W2 implementation and tests look much stronger now around K tails, N tails, zero points, and the LUT/native routing, but MlasIsQNBitGemmAvailable currently reports support for W2 block sizes that the new pack/kernel path does not implement.

tianleiwu

Re-reviewed at 4d442ce. The availability-contract concern from my previous review is fully resolved: W2 variant selection in GetQNBitGemmVariant is now gated independently on BlkLen ∈ {32, 64, 128}, matching what Q2BitGemmPackQuantBDataSize_Avx512 can actually size, so MlasIsQNBitGemmAvailable no longer over-reports support for BlkLen 16/256. The new AvailabilityContract_BlkLens test locks in both the supported and unsupported block sizes plus the unimplemented compute types. No remaining concerns. LGTM.

tianleiwu

Follow-up review of the newly added sqnbitgemm_kernel_avx512_2bit_blklen32.h. The kernel is correct — these are non-blocking maintainability/perf nits only and do not change the prior approval. Verified the block-group/stride constants against sqnbitgemm_kernel_avx512_2bit.h (e.g. kBlockGroupBytes32 == kBlockGroupBlks * kBlkBytes32), so the dispatch-vs-tail pointer strides are consistent by definition.

tianleiwu

LGTM

### Description PR that introduced issue: #29064 Fixes in this PR: 1) Add relevant platform guard in some tests that was previously missing 2) Added the new AVX512 headers that host the kernels to their right location within the cmake file grouping - previously they were placed in AVX2 grouping (ultimately the TU that included those headers were compiled with AVX512 flags - so no harm was done). This fix is more pedantic than fixing a real issue. The lone .cpp file in that list didn't include any intrinsics manually but the compiler might use AVX512 now for auto-vectorization with the shuffling. Since that file contains only the pre-packing functions that are used in production, it is safe. The "scalar" kernel implementation in that file is mostly a test oracle - nothing else Sample failed run (before PR): https://aiinfra.visualstudio.com/Lotus/_build/results?buildId=1268723&view=results Sample successful run (with PR): https://aiinfra.visualstudio.com/Lotus/_build/results?buildId=1273196&view=results ### Motivation and Context Fix Python packaging pipeline

## Summary Adds a new ARM64 NEON implementation of the 2-bit `MatMulNBits` (`SQNBIT_CompInt8`, `BlkBitWidth=2`) path built on `SDOT` + in-kernel 2-bit unpack. Replaces the previous ARM64 W2 fallback, which dequantized the 2-bit weights to fp32 and dispatched them through the standard fp32 GEMM — correct but far below the achievable throughput for a native int8-dot kernel. Also fixes a correctness bug in `MatMulNBits` that mis-handled the MLFloat16 A-input variant on the `accuracy_level=4` (CompInt8) code path on ARM64. This is the ARM64 equivalent of the PR #29064 ## What's in the PR ### New ARM64 W2 DotProd kernel family `onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp` (new TU) - Native `SDOT`-based kernels for `BlkLen ∈ {32, 64, 128}`. - Templated `R{1,2} × C{1,4,8}` DotProd tile grid — one implementation covers the full `M × N` register-tile matrix instead of six copy-pasted variants. - In-kernel 2-bit → int8 unpack via `vsubq_s8(B, 2)`; B zero-point folded into the unpack so no post-kernel `BlkSum` SGEMM is needed. Same shape as the ARM64 W4 DotProd path. - Dispatch plumbing in `qnbitgemm.{cpp,h}` and `qnbitgemm_kernel_neon.{cpp,h}` routes `BlkBitWidth==2` through the new TU on ARM64 hosts with DotProd (Neoverse, Cortex-A76+, Snapdragon X, Apple M-series). ### `MatMulNBits` fp16 A-input CompInt8 fix `onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc` - The MLFloat16 specialization of `PrePack` did not follow the same eager-fold path as the fp32 specialization, so the ARM64 CompInt8 W2 dispatch could mis-select a kernel or use un-recomputed `QuantBBlkSum` values when the A-input is fp16 and `accuracy_level=4`. - Fixed by routing the MLFloat16 `PrePack` through a mirror of the fp32 CompInt8 pipeline, and by having the W2 zero-points callback recompute `QuantBBlkSum = -scale × zp` from `scales_fp32_` consistently for both element types. - A `packed_b_finalized_` flag now guards the single-shot CompInt8 packing so a re-entrant `PrePack` cannot double-pack. ### Docs - `docs/Arm64_w2_kernel_future_enhancements.md` — forward-looking notes on the two follow-ups (SMMLA-based I8MM TU; lane-indexed SDOT for the C8 tile). Not required for this PR to land; captures the design context so the next contributor doesn't re-derive it. ## Performance **Throughput comparison (tok/s):** | seq_len | W4 (1.27 release) | W2 (1.27 release, DeQuant + fp32 GEMM) | **W2 (PR)** | W2 PR vs W2 1.27 | W2 PR vs W4 1.27 | |---|---|---|---|---|---| | 32 | 878.3 | 130.7 | 757.2 | 5.8× | 0.86× | | 64 | 904.4 | 220.3 | 953.4 | 4.3× | 1.05× | | 128 | 797.7 | 318.2 | 1032.0 | 3.2× | 1.29× | ## Notes / follow-ups (not in this PR) - SMMLA-based W2 I8MM TU for `FEAT_I8MM` hosts (R2+ tiles, ~2× dot throughput). - Lane-indexed SDOT for the C8 tile (A-load reduction). - Both are described in `docs/Arm64_w2_kernel_future_enhancements.md` §1 and §2.

### Description This adds native (non-LUT) 2-bit weight CompInt8 kernels to MLAS for AVX2 and AVX-VNNI, BlkLen 32/64/128, modeled on the AVX-512 kernels from #29064 and the ARM64 kernels from #29466. On these hosts a W2 MatMulNBits currently has the LUT path (opt-in, and only when N is a multiple of 128) or the fp32 dequant + SGEMM fallback. The packed layout is untouched. The block-group packer, scale layout and BlkSum machinery in sqnbitgemm_kernel_avx512_2bit.{h,cpp} are portable scalar C++ and the ARM64 kernels already reuse them, so this only adds the 256-bit compute kernels and the dispatch wiring. The four W2 entries are populated in MlasSQNBitGemmDispatchAvx2 and MlasSQNBitGemmDispatchAvx2vnni with BlkLen routing forwarders like the AVX-512 ones, and A quantization reuses the QuantizeARow_CompInt8_avx2 those tables already register. No cmake changes (the kernels are header-only and the avx2 source list already carries -mavxvnni where the compiler supports it) and no operator changes. Per-node routing: | Host | LUT mode | before | after | |---|---|---|---| | AVX2 / AVX-VNNI | off (default) | fp32 dequant + SGEMM | native W2 kernel | | AVX2 / AVX-VNNI | on, N % 128 == 0 | LUT kernel | LUT kernel (unchanged) | | AVX2 / AVX-VNNI | on, N % 128 != 0 | fp32 dequant + SGEMM | native W2 kernel | Tile shapes are per BlkLen and I picked them by measuring. BlkLen 32 and 64 use an R2xC4 main tile (one B block-group load and unpack shared across two rows) with an R1xC4 tail for an odd trailing row. BlkLen 128 stays R1xC4: the R2 variant measured 3 to 5% slower there, which tracks with register pressure, since a 128-byte group needs four B registers live plus eight accumulators and that does not fit sixteen YMM. M=1 always takes the R1 path, so decode is unaffected by the tiling choice either way. Hosts with AVX2 but no AVX-VNNI use the vpmaddubsw + vpmaddwd fallback, guarded the same way as the existing int8 kernels. Testing: - new direct-kernel tests in test_sqnbitgemm_2bit_gemm.cpp mirroring the AVX-512 set, four per BlkLen, with and without zero points. The non-VNNI variants guard on Avx2Supported_ so they also run on the AVX-512 CI hosts and cover the maddubs path there; the VNNI variants gate on the AVX2-VNNI dispatch being the active one. - the existing MatMul2Bits operator tests exercise the new path automatically on AVX2 hosts at accuracy_level 4. - validated the kernels against a float64 reference over 120+ cases (N tails, K tails, odd M, both dot paths) on packing produced by the production 3-call pack sequence. ### Motivation and Context #29064 closed this gap on AVX-512 and #29466 on ARM64, but AVX2/AVX-VNNI without AVX-512 covers most client x86 (Alder Lake through Arrow Lake, plus Zen 1-3 on the plain AVX2 path) and those hosts still land on dequant + SGEMM by default. Kernel-level numbers from a Core Ultra 5 225 (Arrow Lake, AVX2 + AVX-VNNI, no AVX-512), single thread, interleaved arms, min of 9 rounds on rdtsc: | BlkLen | cycles/MAC at prefill (shipped tile) | R2xC4 vs R1xC4 at prefill | |---|---|---| | 32 | ~0.048 | R2 5 to 6% faster (shipped) | | 64 | ~0.031 | R2 11 to 15% faster (shipped) | | 128 | ~0.027 | R2 3 to 5% slower (kept R1) | These are tile-level microbenchmarks, not end-to-end model numbers; my dev box has no MSVC so my validation is kernel-level and the MSVC build rides on CI. Happy to run whatever end-to-end comparison you want on top of this, and happy to restructure the tiles if you would rather keep all three BlkLens on the same shape.

github-advanced-security AI found potential problems Jun 16, 2026

View reviewed changes

Comment thread onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit.cpp Fixed

Comment thread onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp Fixed

hariharans29 requested a review from Copilot June 16, 2026 01:21

Copilot started reviewing on behalf of hariharans29 June 16, 2026 01:22 View session

hariharans29 added 11 commits June 15, 2026 18:27

Stage

8c36d54

AVX512 VNNI kernels

c3062b6

Stage 2

1406960

Stage 3

5b26abf

Add sqnbitgemm avx512vnni 2bit blklen64 kernel

1d5b4eb

W2 unification

b9f085a

LUT fix

a8c932f

W2-v1+ Cleanup

3711489

Support BlkLen 32 and 128

b8ef59d

Missed cmake file

268b97e

Fixes

60132f3

Copilot AI reviewed Jun 16, 2026

View reviewed changes

Cosmetic cleanups

1587c2b

hariharans29 force-pushed the hari/asg_perf_1 branch from 7ff4df4 to 1587c2b Compare June 16, 2026 01:55

Copilot comments

2701ba2

hariharans29 requested a review from Copilot June 16, 2026 21:53

Copilot started reviewing on behalf of hariharans29 June 16, 2026 21:54 View session

Copilot AI reviewed Jun 16, 2026

View reviewed changes

Copilot comments - 2

61c844a

hariharans29 requested a review from Copilot June 16, 2026 22:18

Copilot started reviewing on behalf of hariharans29 June 16, 2026 22:18 View session

Copilot AI reviewed Jun 16, 2026

View reviewed changes

Comment thread onnxruntime/core/mlas/lib/qnbitgemm.cpp Outdated

Copilot comments - 3

847222f

tianleiwu requested changes Jun 16, 2026

View reviewed changes

Comment thread onnxruntime/core/mlas/lib/qnbitgemm.cpp Outdated

Fix availability contract

4d442ce

tianleiwu previously approved these changes Jun 17, 2026

View reviewed changes

tianleiwu reviewed Jun 17, 2026

View reviewed changes

Tianlei PR comments

b40dfb2

hariharans29 dismissed tianleiwu’s stale review via b40dfb2 June 17, 2026 04:53

hariharans29 requested a review from tianleiwu June 17, 2026 07:44

tianleiwu approved these changes Jun 17, 2026

View reviewed changes

hariharans29 enabled auto-merge (squash) June 17, 2026 18:40

hariharans29 merged commit f7c2c2d into main Jun 17, 2026
89 of 90 checks passed

hariharans29 deleted the hari/asg_perf_1 branch June 17, 2026 18:41

hariharans29 mentioned this pull request Jun 20, 2026

[Pipeline] Fix Python packaging pipeline failures #29200

Merged

hariharans29 mentioned this pull request Jul 1, 2026

[MLAS] Add ARM64 2-bit weight CPU kernels #29466

Merged

blazingphoenix7 mentioned this pull request Jul 8, 2026

[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels #29619

Merged

Uh oh!

Conversation

hariharans29 commented Jun 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Per-node routing

End-to-end throughput (tokens/sec, higher is better)

Behavior with mlas.use_lut_gemm opt-in

fp16 fallback path for 2-bit weights

Motivation and Context

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

tianleiwu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tianleiwu left a comment

Choose a reason for hiding this comment

Uh oh!

tianleiwu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tianleiwu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

hariharans29 commented Jun 16, 2026 •

edited

Loading

Behavior with `mlas.use_lut_gemm` opt-in