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

[blazingphoenix7]

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.

[blazingphoenix7]

@microsoft-github-policy-service agree

[Copilot]

Pull request overview

This PR extends MLAS’s SQNBitGemm (2-bit weight, CompInt8) CPU backend with native AVX2 and AVX2-VNNI compute kernels (BlkLen 32/64/128), wiring them into the AVX2/AVX2-VNNI dispatch tables and adding direct-kernel unit tests to cover both the VNNI and non-VNNI dot-product paths.

Changes:

• Adds AVX2/AVX2-VNNI W2 (2-bit) CompInt8 kernels for BlkLen 32/64/128 (header-only compute implementations).
• Wires W2 entries into MlasSQNBitGemmDispatchAvx2 and MlasSQNBitGemmDispatchAvx2vnni, including BlkLen-routing forwarders.
• Adds new unit tests to directly exercise the AVX2 and AVX2-VNNI W2 kernels (with/without zero points) using the existing test harness.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp	Adds AVX2/AVX2-VNNI direct-kernel tests for W2 across BlkLen 32/64/128, with/without zero points.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2.cpp	Includes new W2 AVX2 headers, adds BlkLen-routing forwarders, and populates W2 function pointers in AVX2/AVX2-VNNI dispatch tables.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit.h	Declares the AVX2/AVX2-VNNI W2 dispatch forwarders for use by dispatch tables and unit tests.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit_blklen32.h	Implements the AVX2/AVX2-VNNI W2 kernel for BlkLen=32 (R1xC4 and R2xC4 tiles + tails + BlkSum correction).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit_blklen64.h	Implements the AVX2/AVX2-VNNI W2 kernel for BlkLen=64 (R1xC4 and R2xC4 tiles + tails + BlkSum correction).
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit_blklen128.h	Implements the AVX2/AVX2-VNNI W2 kernel for BlkLen=128 (R1xC4 tile + tails + BlkSum correction).

[hariharans29]

Why is this change needed ? The previous state had it to just strip away avx512 vnni flags so that auto-vectorization (if the compiler fdoes do it) didn't end up using avx512 vnni but we can use avx512 ? Or does the code there get dispatched on avx2 only hosts with this change and hence necessaitates avx512 flags strip off ?

[blazingphoenix7]

Yes, the latter. With this PR the AVX2 and AVX2-VNNI dispatch tables point at the pack helpers in this TU, so it now runs at model load on hosts with no AVX-512 at all. The first CI round caught exactly that: the three Linux GPU pipelines run on AVX2-only CPUs, and onnxruntime_provider_test died with an illegal instruction there, since with -mavx512bw/dq/vl still on the file the compiler is free to autovectorize the scalar pack loops with EVEX encodings. It is the same failure mode the existing -mavx512vnni strip was guarding against, one ISA level down.

The TU is intrinsic-free scalar and the pack is one-time load work, so baseline codegen should not cost anything measurable, and ARM64 already builds this same file with no x86 flags. If you would rather keep avx512 codegen for the avx512 dispatch path, I can give the avx2 side its own copy in a separate TU instead, stripping the flags on the shared file just seemed like the smaller change. The second commit should turn those three jobs green once CI reruns.

[Copilot]

Pull request overview

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

[blazingphoenix7]

Added the explicit include in 4ab49f5.

[hariharans29]

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

Approve pending CI green on commit 2d86379. Well-scoped, well-documented PR that closes the W2 gap on AVX2 / AVX2-VNNI hosts (Alder Lake through Arrow Lake, Zen 1–3 on the AVX2 path) — the same gap that #29064 closed on AVX-512 and #29466 closed on ARM64. Design mirrors the AVX-512 sibling exactly: no packed-layout change, no operator change, no A-quantization change (QuantizeARow_CompInt8_avx2 is reused from the existing int8 path), and no cmake-listed source additions since the compute kernels are header-only. Adds 24 direct-kernel gtests that thoroughly cover BlkLen ∈ {32, 64, 128} × {AVX2, AVX2-VNNI} × {with, without zero-points} × {with, without bias} × 2 seeds × several M/N/K shapes. The two-line cmake change on the shared W2 pack-helper TU is a genuine bugfix that the first CI round caught.

What looks correct

1. Dispatch wiring is symmetric with the AVX-512 side. Both new dispatch tables register the same four W2 entries (Q2BitGemmPackQuantBDataSize, SQ2BitGemmPackQuantBDataAndBlkSum, SQ2BitGemmKernel_BlkSum_CompInt8, Q2BitGemmEffectiveBlockCountK), and they correctly reuse onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512 and SQ2BitGemmPackQuantBDataAndBlkSum_Scalar — those are portable scalar C++ helpers so cross-namespace reuse is fine and matches what the ARM64 W2 dispatch already does. The Q2BitGemmEffectiveBlockCountK lambda MlasDivRoundup(BlockCountK, kBlockGroupBlks) * kBlockGroupBlks is copy-pasted verbatim from the AVX-512 tables. The static_assert on kBlockGroupBlks == kSq2BitAvx512WeightKBlockGroup is duplicated in both new tables, so the invariant is guarded regardless of which table you look at.
2. BlkLen-routing forwarders match the AVX-512 pattern. The SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch / SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch pair is a straightforward if (BlkLen == 128) ... if (BlkLen == 32) ... else BlkLen==64 cascade. The 128 and 32 branches use their per-BlkLen impls; the fallthrough calls the BlkLen=64 impl by dropping the BlkLen/QuantBZeroPoint/CountK arguments the per-BlkLen kernels don't take. That's the exact same shape the AVX-512 forwarders use. Nothing surprising.
3. BlkLen=128 tail handling in accumulate_w2_blklen128_r1c1blk4 is safe. dot_one_block_w2_blklen128<Shift=0> is always evaluated (nblk ≥ 1 by preconditions on the caller), so it always reads a_base + 0..127, which is real A memory. Blocks 1/2/3 are gated by (nblk > k) ternaries, so a_base + k*128 is never dereferenced for k ≥ nblk. scale_a_safe[k] for k ≥ TailBlocks is set to 0.0f before the call — so the FMA d_k * s_a[k] * s_b[k] yields 0.0f as long as s_b[k] is finite. The AVX-512 sibling has been shipping the same tail pattern, which implies the pack helper (SQ2BitGemmPackQuantBDataAndBlkSum_Scalar) already zero-initializes (or at least keeps finite) the padded scale slots — if it didn't, 0.0f * NaN = NaN would have contaminated AVX-512 outputs already.
4. VNNI vs non-VNNI dot-product dispatch inside dot_one_block_w2_blklen128. The #if !defined(__GNUC__) || (__GNUC__ > 10) guard degrades cleanly to the _mm256_maddubs_epi16 + _mm256_madd_epi16 accumulator when the compiler is GCC ≤10 (which lacks _mm256_dpbusds_avx_epi32). On such compilers both the VNNI and non-VNNI dispatch instances end up calling the maddubs path, which is correct behavior even if the VNNI dispatch table was selected at runtime — you just don't get the VNNI speedup. MSVC and modern Clang have the intrinsic, so no cost on production toolchains.
5. Zero-init and unaligned-load choice. _mm256_loadu_si256 is used throughout for both A and B, and _mm256_setzero_si256() for the base accumulator. The packed-B layout guarantees 32-byte alignment on the group boundaries as long as the packer respects it, so a mixed loadu/loada choice would probably be cheap-perf-wise fine — but loadu on modern Intel/AMD costs no more than loada on aligned data, so this is a safe default.
6. Tile-shape decision is well-justified. BlkLen 32 and 64 use an R2xC4 main tile with an R1xC4 tail for an odd trailing row (share one B block-group load and unpack across two rows), and BlkLen 128 stays R1xC4 because the R2 variant would need 4 B registers + 8 accumulators live simultaneously, which doesn't fit the 16-YMM AVX2 register file. Author states the R2 variant measured 3–5% slower at BlkLen=128; this matches what register pressure alone would predict. M=1 always takes R1, so decode workloads are unaffected by the tiling choice at either BlkLen. This is the right analysis.
7. cmake flag strip on sqnbitgemm_kernel_avx512_2bit.cpp is the right fix for the CI failure. The TU is pure C++ / intrinsic-free (only calls scalar pack helpers), and the AVX2 W2 dispatch now points at those helpers, so on AVX2-only hosts the TU is entered at model load. If any of -mavx512bw / -mavx512dq / -mavx512vl remain enabled, the compiler is free to autovectorize the scalar packing loops into EVEX-encoded instructions, which will SIGILL on hosts with no AVX-512 — which is exactly what the three Linux GPU pipelines saw on commit 359378a. Stripping all AVX-512 flags forces baseline codegen. Author's clarification in the PR conversation is convincing.
8. Test additions. 12 non-VNNI + 12 VNNI tests, each sweeping kSimdShapes* for the respective BlkLen (32/64/128) × 2 seeds × {bias on, bias off}. Non-VNNI tests gate on Avx2Supported_ so they also run on AVX-512 CI hosts and exercise the maddubs+maddwd fallback there — that's clever because it grows real coverage without needing an AVX2-only CI leg. VNNI tests gate on QNBitGemmDispatch == &MlasSQNBitGemmDispatchAvx2vnni so they only run when AVX2-VNNI is the actual selected dispatch (avoids double-exercising the same path when AVX-512 outranks). The direct-kernel RunW2Case* harness is reused verbatim from the AVX-512 tests, so all the invariants those tests check (float64 reference, tolerance) apply here too.

Non-blocking observations

1. Sharing the pack-helper TU trades AVX-512 codegen for less duplication. With the flag strip, the pack helper runs the same baseline codegen on AVX-512 hosts too — the AVX-512 W2 dispatch loses whatever AVX-512 auto-vectorization the compiler was previously giving it in this TU. Since the pack is one-time model-load work, not hot per-token compute, this is imperceptible. If Hariharan (or a future perf regression bisect) prefers to keep AVX-512 codegen for the AVX-512 pack path, the author has already volunteered to split the TU in a follow-up. I'd accept the current single-TU form and only split if a perf need emerges. Documenting this trade-off in a comment above the set_source_files_properties(... COMPILE_FLAGS "") line would help the next reader:

   # NOTE: this TU is intrinsic-free scalar helpers reached from both the AVX2 and
   # AVX-512 W2 dispatch tables at model load. Flags must not enable any ISA the
   # AVX2-only host lacks, or the compiler may autovectorize the scalar loops into
   # EVEX instructions and SIGILL. If the AVX-512 W2 pack ever becomes a perf hot
   # spot, split this file: keep the scalar pack in a flag-less TU and put an
   # AVX-512-optimized variant in a separate one only used by the AVX-512 tables.

2. Duplicated Q2BitGemmEffectiveBlockCountK lambda. The same one-liner MlasDivRoundup(BlockCountK, kBlockGroupBlks) * kBlockGroupBlks is now copy-pasted in four places (AVX-512, AVX-512-VNNI, AVX2, AVX2-VNNI). Consider hoisting a constexpr function alongside kBlockGroupBlks in sqnbitgemm_kernel_avx512_2bit.h and referencing it from all four dispatch tables. Non-blocking, but tidies up.
3. Copy-paste between the two new dispatch table initialization blocks. The AVX2 and AVX2-VNNI setup blocks (MlasSQNBitGemmDispatchAvx2 and MlasSQNBitGemmDispatchAvx2vnni at the bottom of sqnbitgemm_kernel_avx2.cpp) differ by a single line — the choice of SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch vs SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch. A small SetupW2AvxDispatch(d, kernel_fn) helper would remove the duplication and make it visually clearer that "AVX2 and AVX2-VNNI configs are identical except for the compute kernel". Non-blocking.
4. Perf numbers. The Arrow Lake table in the PR body compares R2 vs R1 tiles at BlkLen 32/64/128 for VNNI. Two things would strengthen the perf story if easy to gather:
   • AVX2 (non-VNNI) numbers on the same host to confirm the maddubs fallback still beats the fp32-dequant+SGEMM baseline the PR is replacing. On plain AVX2 without VNNI, dequant+SGEMM using FMA is not obviously slower than one-integer-multiply-per-4-elements fake-VNNI, and Zen 1–3 users are the intended beneficiaries there.
   • End-to-end numbers on a real W2 model (e.g., a MatMulNBits-heavy int2 Llama variant). Author says CI-side MSVC exercises the code and end-to-end is offered as a follow-up if reviewers want it — one small model comparison would tie a bow on the perf claim.
5. __GNUC__ > 10 fallback breadth. GCC 10 is EOL upstream (May 2023), but Debian 11 (bullseye, GCC 10.2) is still on LTS through mid-2026, and some corporate builds pin to those toolchains. The current design gracefully degrades those builds to the maddubs path, which is correct but silently disables the VNNI speedup for AVX2-VNNI dispatch. Worth a one-line comment above the guard noting that: // GCC ≥11 needed for _mm256_dpbusds_avx_epi32; older toolchains fall through to the vpmaddubsw+vpmaddwd path even on AVX-VNNI hardware. Non-blocking.
6. assert(CountN % kNCols4 == 0) in Q2Int8GemmR1xC4BlkLen128Avx2. The caller SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen128_Impl computes NMain = (CountN / kNCols4) * kNCols4 and passes that as CountN, so the assert holds. Fine. But assert is a debug-only check — a stray future caller could pass a non-multiple. The Q2Int8GemmRMxC_Tail_BlkLen128Avx2 companion has assert(TailN >= 1 && TailN <= 3). Same story. Consider swapping these to ORT_ENFORCE or MLAS_ASSERT (whichever MLAS convention prefers for kernel-internal preconditions) to keep the guard live in Release. Non-blocking; existing MLAS kernels are inconsistent about this.
7. BiasPtr += BiasPtr != nullptr ? kNCols4 : 0; (~line 199). Cute but slightly obscure — a plain if (BiasPtr != nullptr) BiasPtr += kNCols4; reads more clearly. Style nit.
8. std::byte vs uint8_t for A/B pointers. The new code uses std::byte consistently, which is what the AVX-512 sibling and the dispatch signature already use. Good — no drift.
9. Where's sqnbitgemm_kernel_avx2_2bit_blklen32.h / blklen64.h? Both are listed as +619 and +603 lines but collapsed in the diff view. I didn't line-by-line them. The BlkLen=128 file follows a very clean pattern (dot helper → r1c1blk4 accumulator → main+tail kernel driver → VNNI/non-VNNI thin wrappers), and I'd expect the 32/64 variants to follow the same shape with an added R2 main tile. If the reviewer wants exhaustive verification of those two files, the pattern I'd check for is: (a) dot_one_block_w2_blklen{32,64} correctness for the R2 tile (one B unpack shared across two A rows), (b) tail-block scale zero-padding identical to the BlkLen=128 case, (c) the R1 tail path for the odd-M case at M % 2 == 1. Since the author validated against a float64 reference over 120+ cases and the tests are green on the first commit's CI (83/86), I'd trust the sibling files without a line-by-line pass unless a specific concern surfaces.
10. Bias != nullptr scan at BiasTail = (Bias != nullptr) ? Bias + NMain : nullptr;. Correct null-passthrough — the tail path preserves the null-bias case. Good.
11. PR body's "Testing" section already lists what was validated at the kernel level (float64 reference, 120+ cases, N tails, K tails, odd M, both dot paths, production 3-call pack sequence). That's the right diligence for MLAS work. Nice.

CI

Author noted first commit reached 83/86 checks OK (three failures on AVX2-only Linux GPU pipelines from the AVX-512 flag issue). Second commit (2d86379) currently 1/68 — CI is still running. Once the second-commit CI turns green, no other blockers from my side. If any of those 68 still fail on 2d86379, please share the specific failure and I'll re-review.

Summary

Clean, focused MLAS kernel PR that closes a real gap for the AVX2/AVX-VNNI-only client x86 audience. Design mirrors the AVX-512 sibling exactly, test coverage is thorough (24 new tests, non-VNNI variants deliberately also run on AVX-512 hosts to exercise the maddubs fallback), and the CI-caught cmake fix is well-reasoned. Approve once commit 2d86379's CI is green; non-blocking items (comment on the shared TU, factor the duplicated dispatch setup, assert → MLAS_ASSERT, GCC-10 fallback comment) are all follow-up polish.

[hariharans29]

Thanks for the contribution. Could you please take a look at the lone copilot comment so that we can merge it ? :)

If optionally, you'd like to address the non-blocking comments above, that would be nice as well.

[blazingphoenix7]

Pushed the fixes. The copilot comment is addressed, sqnbitgemm_kernel_avx2.cpp now includes sqnbitgemm_kernel_avx512_2bit.h directly instead of leaning on the transitive include.

From the non-blocking list I also picked up 1, 2 and 5: the cmake note above the flag strip (used your wording, thanks), the effective-BlockCountK one-liner is now a constexpr helper next to kBlockGroupBlks in sqnbitgemm_kernel_avx512_2bit.h with all four dispatch tables referencing it, and there is a comment above the GCC guard in each kernel header now. I left 3, 6 and 7 alone on purpose: the two table blocks mirror how sqnbitgemm_kernel_avx512.cpp lays out its own pair, the asserts match the AVX-512 sibling's, and the BiasPtr ternary is the idiom the existing int8 blklen kernels use, so changing those here felt like drift from the files this one is modeled on. Happy to revisit any of them if you disagree.

On perf: grabbed AVX2-only (maddubs path) numbers on the same host. Prefill kernel efficiency with the shipped tile config, single thread:

BlkLen	VNNI c/MAC	maddubs c/MAC	maddubs/VNNI
32	~0.051	~0.066	1.26x
64	~0.030	~0.047	1.52x
128	~0.027	~0.043	1.56x

Even the maddubs path sits at or under the fp32 SGEMM ceiling on this core (2x256-bit FMA = 16 MAC/cycle = 0.0625 c/MAC at perfect utilization) before you count the dequant pass itself or the 16x weight bytes the fp32 path has to stream, so Zen 1-3 class hosts should still come out well ahead of dequant+SGEMM, and decode is the lopsided case since it is bandwidth bound. An end-to-end model comparison is something I would like to do as a follow-up once I have a proper W2 model set up locally.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.

[blazingphoenix7]

Adding 0 to a null pointer is well defined in C++, [expr.add]/4 makes it a null pointer; it is the C rule that makes it UB. This is also the exact BiasPtr idiom the AVX-512 sibling and the existing int8 blklen kernels use (sqnbitgemm_kernel_avx512_2bit_blklen32.h does the same thing), so I kept it for consistency. Same call as non-blocking item 7 in the review above.