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