[MLAS] Add ARM64 2-bit weight CPU kernels

cmake/onnxruntime_mlas.cmake

@@ -112,6 +112,15 @@ function(setup_mlas_source_for_windows)
         ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp
         ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp
         ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
+        # Portable W2 pack helpers + scalar reference kernel. Misleadingly
+        # Avx512-named because the AVX-512 W2 path was the first consumer, but
+        # the TU contains no x86 intrinsics (see sqnbitgemm_kernel_avx512_2bit.cpp).
+        # ARM64 W2 dispatch reuses these for pack-size / pack / layout; the
+        # native compute kernel is the NEON DotProd TU listed just below.
+        ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.h
+        ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp
+        # W2 CompInt8 DotProd kernel (NEON FEAT_DotProd backend).
+        ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp
         ${MLAS_SRC_DIR}/cast_kernel_neon.cpp
         ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp
         ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16_8bit.cpp
@@ -517,6 +526,15 @@ else()
           ${MLAS_SRC_DIR}/qnbitgemm_kernel_neon.cpp
           ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_fp32.cpp
           ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
+          # Portable W2 scalar pack / reference kernel. See ARM64 (Windows) branch
+          # above for the rationale; this is the matching entry for non-Windows
+          # ARM64 builds (Linux, macOS).
+          ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.h
+          ${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp
+          # W2 CompInt8 DotProd kernel (NEON FEAT_DotProd backend). Compiled
+          # with -march=armv8.2-a+dotprod on this branch -- see flag override
+          # further below alongside the W4/W8 dotprod TU.
+          ${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp
           ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.h
           ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.cpp
           ${MLAS_SRC_DIR}/qkv_quant_kernel.h
@@ -551,6 +569,8 @@ else()
         endif()
         set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp
                                     PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod")
+        set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_2bit.cpp
+                                    PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod")
         set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8_i8mm.cpp
 				    PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ")
 

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

@@ -420,7 +420,14 @@ Status MatMulNBits<T1>::PrePack(const Tensor& tensor, int input_idx, /*out*/ All
 #if defined(MLAS_TARGET_AMD64_IX86)
       return true;
 #else
-      return (nbits_ == 8);
+      // W2 on ARM64 (NEON DotProd) uses the same 3-call prepack pattern as
+      // AVX-512: B in the input_idx==B call, scales+ZP in the separate
+      // input_idx==scales / input_idx==zero_points calls. Without this, the
+      // ZP tensor is silently dropped and QuantBBlkSum bakes in the symmetric
+      // default ZP=2 for every block, producing wrong outputs whenever the
+      // model's real ZPs ≠ 2. W4 stays out because it goes through the
+      // KleidiAI scale-baked path above and does not use the 3-call pattern.
+      return (nbits_ == 2 || nbits_ == 8);
 #endif
     }();
 
@@ -549,12 +556,14 @@ Status MatMulNBits<T1>::PrePack(const Tensor& tensor, int input_idx, /*out*/ All
       }
     }
 
-    // Pack zero_points separately only for 8-bit (matching standard SQNBIT_CompInt8 behavior).
-    // For 4-bit, zero_points are passed directly in data params or handled via KleidiAI BZpCorr.
-    if (input_idx == InputIndex::zero_points && packed_b_ != nullptr && nbits_ == 8 && !packed_b_finalized_) {
+    // Pack zero_points separately for W2; W4/W8 are folded into packed_b_ during the B PrePack.
+    // Pass scales_fp32_ so QuantBBlkSum = -scale*zp is recomputed (otherwise it keeps the symmetric
+    // default ZP=2 from the B-pack call and silently produces wrong outputs when real ZPs != 2).
+    if (input_idx == InputIndex::zero_points && packed_b_ != nullptr && !packed_b_finalized_ && nbits_ == 2) {
       auto zptr = tensor.Data<uint8_t>();
-      MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(), nullptr,
-                                  has_zp_input_, zptr, nullptr, &mlas_backend_kernel_selector_config_);
+      MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(),
+                                  scales_fp32_.get(), has_zp_input_, zptr, nullptr,
+                                  &mlas_backend_kernel_selector_config_);
       is_packed = false;
     }
 
@@ -706,7 +715,7 @@ Status MatMulNBits<MLFloat16>::PrePack(const Tensor& tensor, int input_idx, /*ou
 #if defined(MLAS_TARGET_AMD64_IX86)
       return true;
 #else
-      return (nbits_ == 8);
+      return (nbits_ == 2 || nbits_ == 8);
 #endif
     }();
 

onnxruntime/core/mlas/lib/qnbitgemm.cpp

@@ -1153,9 +1153,24 @@ InitializeWorkspace_CompInt8<float>(
             }
         });
     } else {
-        // TODO(hasesh): Clean-up the following logic so that it is clean AND it works as expected on all platforms
+        // TODO(hasesh): The (BlkBitWidth x A-layout x A-signedness) matrix below is
+        // resolved by a hand-rolled cascade because the dispatch exposes both an
+        // interleaved-scale (W4-style) and a separate-scale (W8-style) quantize fn,
+        // plus a W2-specific signed-A override. Ideally the dispatch itself would
+        // expose a single "correct quantize fn for this (bit-width, kernel)" pointer
+        // so this call-site would not have to reason about layout/sign compatibility.
         if (BlkBitWidth == 4 || BlkBitWidth == 2) {
-            if (QuantizeARow) {
+            // W2 requires the W8-style separate-scale layout produced by
+            // QuantizeARowComputeBlkSum_CompInt8 because the W2 SQ2BitGemm
+            // kernel reads QuantData (M*K int8 flat) and QuantScale
+            // (M*BlockCountK float) as independent buffers. The W4-style
+            // interleaved layout produced by QuantizeARow_CompInt8 packs the
+            // scale inside each Q8Blk and is therefore incompatible. On hosts
+            // that register both (e.g. NEON FEAT_DotProd, which uses
+            // QuantizeARow for the W4 path and QuantizeARowComputeBlkSum for
+            // W8), force W2 down the W8-compatible branch.
+            const bool prefer_compute_blksum = (BlkBitWidth == 2);
+            if (QuantizeARow && !prefer_compute_blksum) {
                 MlasTrySimpleParallel(ThreadPool, BatchN, [&](ptrdiff_t gemm_idx) {
                     const auto& data = DataParams[gemm_idx];
 
@@ -1169,6 +1184,14 @@ InitializeWorkspace_CompInt8<float>(
                     }
                 });
             } else if (QuantizeARow2) {
+                // W2 needs SIGNED int8 A (NEON DotProd-only hosts wire the
+                // shared QuantizeARowComputeBlkSum_CompInt8 to the UNSIGNED
+                // u8 = i8+128 W8-compatible variant). Prefer the W2-specific
+                // override when set so the W2 kernel always sees signed A.
+                const auto QuantizeARow2_W2 =
+                    (BlkBitWidth == 2 && GetMlasPlatform().QNBitGemmDispatch->QuantizeARowComputeBlkSum_CompInt8_W2 != nullptr)
+                        ? GetMlasPlatform().QNBitGemmDispatch->QuantizeARowComputeBlkSum_CompInt8_W2
+                        : QuantizeARow2;
                 MlasTrySimpleParallel(ThreadPool, BatchN, [&](ptrdiff_t gemm_idx) {
                     const auto& data = DataParams[gemm_idx];
                     const float* ARowPtr = data.A;
@@ -1179,7 +1202,7 @@ InitializeWorkspace_CompInt8<float>(
                     float* QuantARowScalePtr = quant_a_data.QuantScale;
                     float* QuantARowBlkSum = quant_a_data.BlockSum;
                     for (size_t m = 0; m < M; ++m) {
-                        QuantizeARow2(BlkLen, ARowPtr, K, QuantARowPtr, QuantARowScalePtr, QuantARowBlkSum);
+                        QuantizeARow2_W2(BlkLen, ARowPtr, K, QuantARowPtr, QuantARowScalePtr, QuantARowBlkSum);
                         ARowPtr += data.lda;
                         QuantARowPtr += BlockCountK * BlkLen;
                         QuantARowScalePtr += BlockCountK;

onnxruntime/core/mlas/lib/qnbitgemm.h

@@ -679,6 +679,13 @@ struct MLAS_QNBIT_GEMM_DISPATCH {
     );
     QuantizeARowComputeBlkSum_CompInt8_Fn* QuantizeARowComputeBlkSum_CompInt8 = nullptr;
 
+    // Optional W2-specific override of QuantizeARowComputeBlkSum_CompInt8.
+    // The W2 NEON DotProd kernel assumes SIGNED int8 A, but on DotProd-only hosts
+    // the shared field is wired to the UNSIGNED (u8 = i8+128) W8-compatible variant.
+    // When set, the W2 dispatch in InitializeWorkspace_CompInt8 uses this in
+    // preference to the shared field; otherwise it falls back to the shared one.
+    QuantizeARowComputeBlkSum_CompInt8_Fn* QuantizeARowComputeBlkSum_CompInt8_W2 = nullptr;
+
     /**
      * @brief Multiply fp16 matrix A rows with fp16 matrix B columns.
      *        Results are written to fp16 matrix C.

onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp

@@ -27,6 +27,13 @@ Module Name:
 #include "qnbitgemm.h"
 #include "sqnbitgemm_q8_block.h"
 
+// W2 block-group pack helpers and scalar reference kernel declared in the
+// Avx512-named header are pure C++ (no x86 intrinsics). They serve as the
+// cross-arch layout authority for W2 and are reused on ARM64 for pack-size /
+// pack / layout; compute is provided by the native NEON DotProd kernel in
+// sqnbitgemm_kernel_neon_int8_2bit.cpp.
+#include "sqnbitgemm_kernel_avx512_2bit.h"
+
 #ifdef USE_KLEIDIAI
 #include "kai/kai_common.h"
 #include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0.h"
@@ -655,6 +662,47 @@ GetMlasQNBitGemmDispatchNeon(
             d.SQ8BitGemmKernel_BlkSum_CompInt8 = sqnbitgemm_neon::SQ8BitGemmKernel_BlkSum_CompInt8<false>;
         }
 
+        // W2 native CompInt8 path.
+        //
+        // Pack-size, pack-and-blksum, and EffectiveBlockCountK use the
+        // portable AVX-512-namespaced helpers (the file is misleadingly
+        // named -- the TU contains no x86 intrinsics) which serve as the
+        // cross-arch layout authority for W2.
+        //
+        // SQ2BitGemmKernel_BlkSum_CompInt8 is wired to the native NEON
+        // DotProd kernel when FEAT_DotProd is available; the kernel
+        // handles BlkLen ∈ {32, 64, 128} natively via SDOT inner loops.
+        // The 2-bit weights are unpacked to unsigned values in [0, 3] and
+        // fed directly to SDOT; the B zero-point correction is fused into
+        // the accumulator via the ABlockSum × QuantBBlkSum term (where
+        // QuantBBlkSum bakes in -scale × zp per block), so no post-kernel
+        // zero-point correction SGEMM is needed.
+        //
+        // FEAT_I8MM hosts also take this path: FEAT_I8MM always implies
+        // FEAT_DotProd per the ARM spec, and USDOT and SDOT have identical
+        // throughput on every core that implements both -- a separate I8MM
+        // TU using USDOT would be pure duplication. The real I8MM-only
+        // throughput win lives in SMMLA (R2-tile 2×2 matrix-multiply,
+        // 2× the dots/cycle of SDOT), which is future work.
+        //
+        // The Scalar fall-back is defensive: an I8MM-without-DotProd host
+        // is unreachable on conformant ARMv8.2+ hardware.
+        if (InitializeWithDotSupport || InitializeWithI8MMSupport) {
+            d.Q2BitGemmPackQuantBDataSize       = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512;
+            d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar;
+            d.SQ2BitGemmKernel_BlkSum_CompInt8  = InitializeWithDotSupport
+                ? sqnbitgemm_neon::SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd
+                : onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmKernel_BlkSum_CompInt8_Scalar;
+            d.Q2BitGemmEffectiveBlockCountK     = [](size_t BlockCountK) {
+                return MlasDivRoundup(BlockCountK, kSq2BitAvx512WeightKBlockGroup) * kSq2BitAvx512WeightKBlockGroup;
+            };
+            // W2 NEON DotProd kernel uses vdotq_s32 over A and so requires
+            // SIGNED int8 A. The shared QuantizeARowComputeBlkSum is wired
+            // to the UNSIGNED W8 variant on DotProd-only hosts; route W2
+            // through the signed variant explicitly.
+            d.QuantizeARowComputeBlkSum_CompInt8_W2 = sqnbitgemm_neon::QuantizeARowComputeBlkSum_CompInt8<false>;
+        }
+
 #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64)
         d.HQ4BitGemmPackQuantBData = sqnbitgemm_neon::HQ4BitGemmPackQuantBData_CompFp16;
         d.HQ4BitBlkDequantBForHgemm_CompFp16 = sqnbitgemm_neon::HQ4BitBlkDequantBForHgemm_CompFp16;

onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h

@@ -195,6 +195,28 @@ SQ4BitGemmKernel_CompInt8(
     const float* Bias
 );
 
+// W2 CompInt8 kernel entry point (DotProd backend). Implemented in
+// sqnbitgemm_kernel_neon_int8_2bit.cpp. Signature matches the
+// SQ4BitGemmKernel_BlkSum_CompInt8_Fn typedef in qnbitgemm.h.
+size_t
+SQ2BitGemmKernel_BlkSum_CompInt8_NeonDotProd(
+    size_t BlkLen,
+    const std::byte* QuantA,
+    const float* QuantAScale,
+    const std::byte* QuantBData,
+    const float* QuantBScale,
+    const std::byte* QuantBZeroPoint,
+    float* C,
+    size_t CountM,
+    size_t CountN,
+    size_t CountK,
+    size_t BlockCountK,
+    const float* Bias,
+    size_t ldc,
+    const float* ABlockSum,
+    const float* QuantBBlkSum
+);
+
 #ifdef USE_KLEIDIAI
 void
 QuantizeA_Packed_CompInt8(