[MLAS] Add ARM64 2-bit weight CPU kernels#29466
Merged
Merged
Conversation
Checkpoint covering Commits 1 and 2 of the ARM64 W2 non-LUT port plan.
Commit 1 - Plumbing:
- qnbitgemm_kernel_neon.cpp: wire W2 SQNBIT_CompInt8 dispatch fields (size, pack, kernel, EffectiveBlockCountK) inside GetMlasQNBitGemmDispatchNeon, gated on InitializeWithDotSupport || InitializeWithI8MMSupport. Reuses the portable scalar pack helpers from sqnbitgemm_kernel_avx512_2bit.{h,cpp} (the file is misleadingly named: the pack and scalar-oracle kernel are pure C++ with no x86 intrinsics).
- cmake/onnxruntime_mlas.cmake: compile sqnbitgemm_kernel_avx512_2bit.{h,cpp} on Windows ARM64 and Linux/macOS ARM64.
Commit 2 - DotProd BlkLen=64 native kernel:
- sqnbitgemm_kernel_neon_int8_2bit.cpp (NEW): SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd. R1xC1 inner kernel for BlkLen=64 using 4x vdotq_s32 + vaddvq_s32 per K-block. UnpackBlockGroupSliceBlkLen64_DotProd<BlkInGroup> reads the 4-K-block packed-B layout (byte b holds weight b of all 4 K-blocks at bit positions {0..1,2..3,4..5,6..7}) via vshrq_n_u8 + mask. BlkLen=32/128 delegate to the scalar oracle (native specializations land in Commit 3).
- qnbitgemm_kernel_neon.h: declare SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd.
- qnbitgemm_kernel_neon.cpp: InitializeWithDotSupport branch overrides d.SQ2BitGemmKernel_BlkSum_CompInt8 to the native kernel; i8mm-only hosts keep the scalar fallback (i8mm-native kernel lands in Commit 4).
- cmake/onnxruntime_mlas.cmake: add new TU to both ARM64 source lists; -march=armv8.2-a+dotprod on Linux/macOS.
- qnbitgemm.cpp InitializeWorkspace_CompInt8<float> bug fix: when both A-quantizers are registered (e.g. NEON DotProd registers both QuantizeARow and QuantizeARowComputeBlkSum), the W4 branch picked QuantizeARow which produces an interleaved (scale-inside-Q8Blk) layout incompatible with the W2 kernel. Force W2 down the QuantizeARowComputeBlkSum branch. Bug is platform-specific: never surfaced on AVX-512 because that platform only registers QuantizeARowComputeBlkSum.
- test_sqnbitgemm_2bit_gemm.cpp: hoist BlkLen=64 helpers out of the MLAS_TARGET_AMD64 gate (pure scalar C++, now cross-arch). Add three ARM64 tests (gated on HasArmNeonDot): NeonDotProd_BlkLen64 (48 shapes), NeonDotProd_BlkLen64_WithZeroPoints (24 shapes), NeonDotProd_DispatchE2E_BlkLen64 (drives MlasQNBitGemmBatch end-to-end through the production dispatch).
Build clean on Windows ARM64; all 13 *Sq2Bit* tests pass.
Extend the ARM NEON FEAT_DotProd W2 SQNBIT_CompInt8 inner kernel TU with native specializations for BlkLen=32 and BlkLen=128, replacing the placeholder scalar-fallback delegations from the previous checkpoint.
Layout: identical bit-position scheme {0..1, 2..3, 4..5, 6..7} per byte across all three BlkLens; only bytes-per-K-block scales (8 / 16 / 32 for BlkLen 32 / 64 / 128). Block-group still 4 K-blocks. K-tail handler unchanged in structure (tail_blocks 1..3 path).
Per BlkLen the K-block dot uses 2 / 4 / 8 vdotq_s32 (BlkLen 32 / 64 / 128) followed by vaddvq_s32 reduction. Tile shape remains R1xC1 -- wider tiles are a follow-up perf checkpoint.
Tests: hoist the BlkLen=128 and BlkLen=32 helper namespaces in test_sqnbitgemm_2bit_gemm.cpp out of the MLAS_TARGET_AMD64 gate so the cross-arch helpers (RunW2Case_BlkLen128, RunW2Case_BlkLen32) can drive direct-call kernel coverage on ARM64. Add 4 new ARM64 NeonDotProd tests gated on HasArmNeonDot().
All 17 *Sq2Bit* tests pass on Windows ARM64 (build2/Release): 9 pack/unpack + 7 NeonDotProd + 1 availability contract.
W2 DotProd kernel: refactor from per-BlkLen R1xC1 flat kernels into a
single templated tile kernel Q2Int8GemmRxC_DotProd<BlkLen, NRows, NCols>
plus per-BlkLen dispatchers that split the (M, N) plane into up to six
sub-rectangles (R{1,2} x C{1,4,8}). The R2 variants amortise W2's
2-bit unpack (8 weights/byte, 2 shifts+masks per chunk) and the
per-K-block scale / BlkSum gather across an M-pair. Deviates from
the W4/W8 ARM convention of R1-only on DotProd (R2 reserved for
I8MM) because the unpack tax is larger for W2 than for W4/W8.
Decision: do not ship a W2 I8MM TU. FEAT_I8MM always implies FEAT_DotProd
on conformant ARMv8.2+ hardware, and USDOT and SDOT have identical
throughput on every core that implements both. A USDOT-based I8MM TU
would duplicate the DotProd path with no measurable speedup. The real
I8MM-only win lives in SMMLA (R2-tile 2x8 x 8x2 matrix-multiply, 2x
the dots/cycle of SDOT), which is left as future work along with
lane-indexed SDOT for the C8 tile.
Test extension: kSimdShapes and kSimdShapes_BlkLen{32,128} each gain 12
R*xC4 tile-coverage shapes (M = {1,2,3}, N = {4,5,12,13,21}) so the
existing direct-call gtests exercise all six tile variants. The six
NeonI8MM_* tests are removed alongside the dropped TU.
All *Sq2Bit* tests (17 total) pass on Windows ARM64.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an ARM64 NEON (FEAT_DotProd) native 2-bit (BlkBitWidth=2) MatMulNBits / MLAS QNBitGEMM CompInt8 kernel path, replacing the slow dequant-to-fp32 fallback on ARM64 and addressing a correctness issue in the MLFloat16 accuracy_level=4 prepack pipeline.
Changes:
- Introduces a new SDOT-based ARM64 W2 inner-kernel TU with a templated RxC tile grid covering BlkLen {32,64,128}.
- Updates MLAS QNBitGEMM dispatch/plumbing (including a W2-specific A-quantization override hook) and CMake source wiring for ARM64 builds.
- Adds/extends unit and operator tests for the new ARM64 W2 kernels and the
MLFloat16+ uint8-ZPaccuracy_level=4bugfix; includes a forward-looking design doc.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp | Refactors test gating to be cross-arch and adds direct-call ARM64 NEON DotProd W2 kernel tests + broader shape coverage. |
| onnxruntime/test/contrib_ops/matmul_2bits_test.cc | Adds an operator-level regression test for MLFloat16 + uint8 ZP on accuracy_level=4 to guard the fixed prepack path. |
| onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8_2bit.cpp | New ARM64 NEON DotProd (SDOT) W2 CompInt8 inner-kernel implementation for BlkLen {32,64,128}. |
| onnxruntime/core/mlas/lib/qnbitgemm.h | Extends dispatch struct with an optional W2-specific QuantizeARowComputeBlkSum_CompInt8 override. |
| onnxruntime/core/mlas/lib/qnbitgemm.cpp | Forces W2 down the ComputeBlkSum quantization layout and uses the new W2-specific quantizer override when present. |
| onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h | Declares the new W2 NEON DotProd kernel entry point. |
| onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp | Wires up W2 pack/pack-size to portable helpers and compute to the new NEON DotProd kernel; sets W2 signed-A quantization override. |
| onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc | Adjusts ARM64 W2 prepack behavior to ensure scales/ZPs flow correctly and QuantBBlkSum is recomputed consistently for the fp16 activation path. |
| docs/Arm64_w2_kernel_future_enhancements.md | New design notes on potential future SMMLA (I8MM) and lane-indexed SDOT improvements. |
| cmake/onnxruntime_mlas.cmake | Adds the new ARM64 W2 kernel TU (and portable W2 helpers) to MLAS build sources; applies dotprod compile flag on non-Windows ARM64. |
tianleiwu
approved these changes
Jul 2, 2026
Member
Author
|
The Mac python packaging pipeline failures are unrelated to this change - they will be fixed by the XCode update PR #29468 - hence merging this |
hariharans29
pushed a commit
that referenced
this pull request
Jul 17, 2026
### 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.
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.
Summary
Adds a new ARM64 NEON implementation of the 2-bit
MatMulNBits(SQNBIT_CompInt8,BlkBitWidth=2) path built onSDOT+ 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 inMatMulNBitsthat mis-handled the MLFloat16 A-input variant on theaccuracy_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)SDOT-based kernels forBlkLen ∈ {32, 64, 128}.R{1,2} × C{1,4,8}DotProd tile grid — one implementation covers the fullM × Nregister-tile matrix instead of six copy-pasted variants.vsubq_s8(B, 2); B zero-point folded into the unpack so no post-kernelBlkSumSGEMM is needed. Same shape as the ARM64 W4 DotProd path.qnbitgemm.{cpp,h}andqnbitgemm_kernel_neon.{cpp,h}routesBlkBitWidth==2through the new TU on ARM64 hosts with DotProd (Neoverse, Cortex-A76+, Snapdragon X, Apple M-series).MatMulNBitsfp16 A-input CompInt8 fixonnxruntime/contrib_ops/cpu/quantization/matmul_nbits.ccPrePackdid 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-recomputedQuantBBlkSumvalues when the A-input is fp16 andaccuracy_level=4.PrePackthrough a mirror of the fp32 CompInt8 pipeline, and by having the W2 zero-points callback recomputeQuantBBlkSum = -scale × zpfromscales_fp32_consistently for both element types.packed_b_finalized_flag now guards the single-shot CompInt8 packing so a re-entrantPrePackcannot 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):
Notes / follow-ups (not in this PR)
FEAT_I8MMhosts (R2+ tiles, ~2× dot throughput).docs/Arm64_w2_kernel_future_enhancements.md§1 and §2.