[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels (#29619)
### 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.
