[CUDA] Add 2-bit MatMulNBits support

docs/contrib_ops/cuda/matmul_nbits.md

@@ -6,7 +6,7 @@ chain, the fast / fallback / specialized kernels, the weight format they expect,
 and the environment variables that control routing.
 
 MatMulNBits computes `Y = A · dequant(B)ᵀ (+ bias)` where `B` is an `N × K`
-weight matrix quantized to 4 or 8 bits with block-wise (group) scales and
+weight matrix quantized to 2, 4, or 8 bits with block-wise (group) scales and
 optional zero points. It is the building block for weight-only quantized linear
 layers (including MoE routers and LM heads).
 
@@ -15,6 +15,8 @@ Source files:
 - [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc) — operator, `ComputeInternal`, dispatch chain, dequant+GEMM fallback.
 - [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h) — kernel class, constructor-time configuration, environment-variable parsing.
 - [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh) — `TryMatMulNBits` fast-path entry and bias-add launcher.
+- [onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu) — 2-bit fast-path dispatch.
+- [onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu) — generic 2-bit fallback dequantization.
 - [onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu) — 4-bit fast GEMV kernels (generic + router specialization).
 
 ---
@@ -26,7 +28,9 @@ Source files:
 3. [Dispatch Chain](#3-dispatch-chain)
 4. [Fast Path — Fused GEMV](#4-fast-path--fused-gemv)
    - [4.1 Generic 4-bit GEMV kernel](#41-generic-4-bit-gemv-kernel)
-   - [4.2 Router GEMV specialization](#42-router-gemv-specialization)
+  - [4.2 Small-M batched GEMV](#42-small-m-batched-gemv)
+  - [4.3 Router GEMV specialization](#43-router-gemv-specialization)
+  - [4.4 2-bit fused kernels](#44-2-bit-fused-kernels)
 5. [Fallback Path — Dequantize + GEMM](#5-fallback-path--dequantize--gemm)
 6. [fpA_intB_gemm Path (CUTLASS weight-only)](#6-fpa_intb_gemm-path-cutlass-weight-only)
 7. [Bias Handling](#7-bias-handling)
@@ -41,18 +45,18 @@ Source files:
 |-----------|---------|
 | `K` | Input feature dimension (columns of `A`, columns of the logical `B`). |
 | `N` | Output feature dimension (rows of the logical `B`). |
-| `bits` | Quantization bit width: `4` or `8`. |
-| `block_size` | Quantization group size along `K` (16 / 32 / 64 / 128). One scale (and optional zero point) per group. |
+| `bits` | Quantization bit width: `2`, `4`, or `8`. |
+| `block_size` | Power-of-two quantization group size along `K`, at least 16. Fast-path support varies by bit width. One scale (and optional zero point) per group. |
 | `accuracy_level` | Minimum accuracy level for internal handling of `A`; default `0` means unset. |
 | `weight_prepacked` | CUDA fpA_intB weight-layout selector. `0` (default): `B` is in standard MatMulNBits layout and may be runtime-prepacked. `1`: `B` is already prepacked in the CUDA SM80 fpA_intB layout. `2`: `B` is prepacked in the CUDA SM90 (Hopper) fpA_intB layout, consumed by the native SM90 kernel (requires an SM90 device and `block_size` in {64, 128}). The native SM90 kernel is not compiled on Windows/MSVC builds (CUDA 13 host stubs hit MSVC `C2719` with over-aligned TMA parameters — see [moe_qmoe.md §14.1](./moe_qmoe.md)); on those builds the default `0`/`1` layouts run the SM80 compatibility kernel on Hopper instead. |
 
 | Input | Index | Notes |
 |-------|-------|-------|
 | `A` | 0 | Activations, FP16 / BF16 / FP32. Shape `[M, K]`. |
-| `B` | 1 | Packed 4/8-bit weights. |
+| `B` | 1 | Packed 2/4/8-bit weights. |
 | `scales` | 2 | Per-group scales, same element type as `A`. |
 | `zero_points` | 3 | Optional. Packed integer (symmetric default) **or** same type as `A`. |
-| `g_idx` / `reorder_idx` | 4 | Optional group/reorder index (act-order). |
+| `g_idx` / `reorder_idx` | 4 | Optional group/reorder index (act-order). Not supported for 2-bit CUDA weights. |
 | `bias` | 5 | Optional `[N]` bias added to the output. |
 
 `M` is the (flattened) token count: `M = 1` is the decode / GEMV case that the
@@ -62,17 +66,24 @@ fast kernels target.
 
 ## 2. Weight Format
 
-For 4-bit, `B` is stored as `[N, ceil(K / block_size), block_size / 2]` bytes:
-each expert/output row `n` is a contiguous run of `K/2` bytes (two 4-bit weights
-per byte), preceded conceptually by `K / block_size` scales in the `scales`
-tensor. Quantization is **column-wise block** by default
+For `bits=b`, `B` is stored as
+`[N, ceil(K / block_size), block_size / (8 / b)]` bytes: each expert/output row
+`n` is a contiguous packed run with `8 / b` weights per byte, preceded
+conceptually by `K / block_size` scales in the `scales` tensor. Quantization is
+**column-wise block** by default
 (`column_wise_quant_blk_ = true`); row-wise layouts interleave `K` blocks across
 `N` and cannot be sliced along `N` (this disables the chunked fallback).
 
 Symmetric 4-bit weights store values `0..15` that dequantize to `(q − 8) ·
 scale`; the fast kernels hard-code the zero point of `8` when no `zero_points`
 input is present.
 
+For 2-bit weights, four codes are packed into each byte, low two-bit code first.
+Packed zero points use the same four-per-byte layout and default to `2` when the
+input is absent. The CUDA implementation supports power-of-two block sizes from
+16 through 256; fused kernels accept block sizes that divide their 512-element
+warp iteration. `g_idx` is rejected for 2-bit weights.
+
 ### 2.1 CUDA fpA_intB prepacked layout
 
 `weight_prepacked=1` means input `B` has the same tensor shape and byte count as
@@ -126,7 +137,7 @@ flowchart TD
   T1 -- fail due to bias --> T2[TryMatMulNBits without bias]
   T2 -- success --> BIAS[MatMulNBitsBiasAdd] --> R
   T2 -- fail --> DQ
-  DQ[Dequantize blockwise 4b/8b] --> GEMM[cuBLAS GEMM<br/>full-N or chunked along N]
+  DQ[Dequantize blockwise 2b/4b/8b] --> GEMM[cuBLAS GEMM<br/>full-N or chunked along N]
   GEMM --> BIAS2[optional bias add] --> R
 ```
 
@@ -142,6 +153,7 @@ dispatches by bit width:
 
 - `bits == 8` → `TryMatMul8Bits` (no bias support; returns `false` if bias set).
 - `bits == 4` → `TryMatMul4Bits`.
+- `bits == 2` → `TryMatMul2Bits` (no fused bias support; returns `false` if bias set).
 
 `TryMatMul4Bits` ([matmul_4bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu))
 first applies a guard common to all fused kernels:
@@ -214,6 +226,15 @@ Design notes:
 To add another router, extend `IsSupportedRouterGemvShape` with its `(N, K)`;
 no kernel change is required as long as `N % 8 == 0` and `K % block_size == 0`.
 
+### 4.4 2-bit fused kernels
+
+The 2-bit fast path handles `1 <= M <= 8`, `N % 8 == 0`, and `K % 16 == 0`.
+Each thread loads one aligned 32-bit word containing 16 codes. The M=1 kernel
+uses one warp per output column; M=2 through 8 use register-tiled batched
+kernels. The block size must be at least 16 and divide both `K` and the
+512-element warp iteration. FP16 half2 kernels require SM53 or newer; older
+devices decline the fused path and use generic dequantization plus cuBLAS.
+
 ---
 
 ## 5. Fallback Path — Dequantize + GEMM
@@ -222,8 +243,9 @@ When neither the fpA_intB path nor the fused GEMV applies (e.g. `M > 1`,
 `reorder_idx` present, typed zero points, or an unsupported shape), the operator
 dequantizes `B` into a scratch buffer and runs a dense cuBLAS GEMM:
 
-- `DequantizeBlockwise4b` / `DequantizeBlockwise8b` (or the column-wise helper)
-  expands packed weights to `T` into a `N × K_padded` scratch buffer, where
+- `DequantizeNBits` dispatches the column-wise 2/4/8-bit implementations; the
+  row-wise fallback uses `DequantizeBlockwise4b` / `DequantizeBlockwise8b`. These
+  expand packed weights to `T` into a `N × K_padded` scratch buffer, where
   `K_padded = ceil(K / block_size) · block_size`.
 - A single cuBLAS `GEMM` (`transb = true`) then produces `Y`.
 
@@ -292,7 +314,7 @@ Prepacked weights are intentionally strict:
 ## 7. Bias Handling
 
 Only the router specialization (§4.3) fuses bias inside the GEMV. For every other
-fast-path shape, `TryMatMul4Bits` / `TryMatMul8Bits` return `false` when bias is
+fast-path shape, `TryMatMul2Bits` / `TryMatMul4Bits` / `TryMatMul8Bits` return `false` when bias is
 present. `ComputeInternal` then:
 
 1. Retries `TryMatMulNBits` with `bias = nullptr`; on success it adds the bias
@@ -342,6 +364,10 @@ present. `ComputeInternal` then:
   fpA_intB prepacking for int4/int8 and GEMV/GEMM-shaped `M` values.
 - Constructor failure tests for unsupported prepacked configurations live in
   [onnxruntime/test/contrib_ops/matmul_4bits_test.cc](../../../onnxruntime/test/contrib_ops/matmul_4bits_test.cc).
+- CUDA 2-bit fused, fallback, chunked, and validation coverage lives in
+  [onnxruntime/test/contrib_ops/matmul_2bits_test.cc](../../../onnxruntime/test/contrib_ops/matmul_2bits_test.cc).
+  Run it from `onnxruntime_provider_test` with
+  `--gtest_filter=MatMul2BitsCuda.*`.
 - GEMV profiling baselines and methodology are recorded in
   [qmoe_gemv_experiments.md](qmoe_gemv_experiments.md).
 - To compare the router specialization against the generic path, run the same

onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h

@@ -63,6 +63,11 @@ Status CheckInputs(const T* /*activation*/,
                            "block_size must be a power of 2, and >= 16. Got ", block_size);
   }
 
+  if (bits == 2 && block_size > 256) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "2-bit MatMulNBits block_size must not exceed 256. Got ", block_size);
+  }
+
   const int64_t k_blocks = GetKBlocks(k, block_size);
   const int64_t blob_size = GetBlobSize(block_size, bits);
 

onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh

@@ -46,6 +46,18 @@ struct BlkQuantTraits {
   using ThreadBlk = Shape2D<QuantBlk::kRow * kPackSize, QuantBlk::kColumn>;
 };
 
+template <class T, typename ZeroT>
+Status Dequantize2Bits(
+    T* output,
+    const uint8_t* quant_data,
+    const T* scales_data,
+    const ZeroT* zero_points,
+    const int32_t* reorder_idx,
+    int k,
+    int n,
+    int block_size,
+    cudaStream_t stream);
+
 template <class T, typename ZeroT>
 Status Dequantize4Bits(
     T* output,
@@ -82,7 +94,9 @@ Status DequantizeNBits(
     int n,
     int block_size,
     cudaStream_t stream) {
-  if (bits == 4) {
+  if (bits == 2) {
+    return Dequantize2Bits<T, ZeroT>(output, quant_data, scales_data, zero_points, reorder_idx, k, n, block_size, stream);
+  } else if (bits == 4) {
     return Dequantize4Bits<T, ZeroT>(output, quant_data, scales_data, zero_points, reorder_idx, k, n, block_size, stream);
   } else {
     ORT_ENFORCE(bits == 8);

onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu

@@ -0,0 +1,158 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include <cstdint>
+#include <type_traits>
+
+#include "core/providers/cuda/cu_inc/common.cuh"
+#include "core/providers/cuda/cuda_common.h"
+#include "contrib_ops/cuda/quantization/dequantize_blockwise.cuh"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+namespace {
+
+using onnxruntime::cuda::CeilDiv;
+using onnxruntime::cuda::GridDim;
+
+// A 2-bit weight blob packs four codes per byte, code i at bit offset 2*(i%4) of byte i/4
+// (the layout produced by MlasQuantizeBlockwise<T, 2>). One uint32 load therefore covers
+// 16 consecutive codes, which is also the widest aligned store (32 B for fp16, 64 B for fp32).
+constexpr int kElementsPerThread2b = 16;
+constexpr int kElementsPerByte2b = 4;
+constexpr uint8_t kDefaultZeroPoint2b = 2;  // 1 << (bits - 1)
+
+__device__ __forceinline__ float ToFloat2b(float v) { return v; }
+__device__ __forceinline__ float ToFloat2b(half v) { return __half2float(v); }
+__device__ __forceinline__ float ToFloat2b(nv_bfloat16 v) { return __bfloat162float(v); }
+
+__device__ __forceinline__ uint8_t UnpackZeroPoint2b(const uint8_t* zero_points, int64_t byte_index, int block_index) {
+  return static_cast<uint8_t>((zero_points[byte_index] >> ((block_index & 0x03) << 1)) & 0x03);
+}
+
+template <class T>
+__device__ __forceinline__ void StoreSixteen2b(const T* src, T* dst) {
+  constexpr int kVectors = static_cast<int>(kElementsPerThread2b * sizeof(T) / sizeof(float4));
+#pragma unroll
+  for (int i = 0; i < kVectors; ++i) {
+    reinterpret_cast<float4*>(dst)[i] = reinterpret_cast<const float4*>(src)[i];
+  }
+}
+
+// Dequantizes 16 codes with a single (scale, zero point) pair. The arithmetic is done in fp32 and
+// rounded once on store, which is at least as accurate as the reference CPU dequantizer.
+template <class T>
+__device__ __forceinline__ void DequantizeSixteen2b(uint32_t values_quant, T scale, float zero_point, T* output) {
+  const float scale_f = ToFloat2b(scale);
+  const float zp_adjust = -scale_f * zero_point;
+  alignas(16) T results[kElementsPerThread2b];
+#pragma unroll
+  for (int i = 0; i < kElementsPerThread2b; ++i) {
+    const float code = static_cast<float>((values_quant >> (2 * i)) & 0x03u);
+    results[i] = static_cast<T>(fmaf(code, scale_f, zp_adjust));
+  }
+  StoreSixteen2b<T>(results, output);
+}
+
+template <class T, typename ZeroT>
+__global__ void Dequantize2BitsKernel(
+    T* output,
+    const uint8_t* quant_data,
+    const T* scale_data,
+    const ZeroT* zero_points,
+    int block_size,
+    int groups_per_K,
+    int groups_per_threadblock,
+    int total_groups) {
+  const int block_id = blockIdx.x * groups_per_threadblock + ((threadIdx.x * kElementsPerThread2b) / block_size);
+  if (block_id >= total_groups) {
+    return;
+  }
+  const int64_t element_offset = static_cast<int64_t>(block_id) * block_size +
+                                 ((threadIdx.x * kElementsPerThread2b) & (block_size - 1));
+  const uint32_t quant_value = *(reinterpret_cast<const uint32_t*>(quant_data + element_offset / kElementsPerByte2b));
+  const T scale = *(scale_data + block_id);
+
+  float zero_point_value;
+  if constexpr (std::is_same_v<ZeroT, uint8_t>) {
+    const int zero_point_shape_x = (groups_per_K + kElementsPerByte2b - 1) / kElementsPerByte2b;
+    const int kb_idx = block_id % groups_per_K;
+    const int n_idx = block_id / groups_per_K;
+    uint8_t zp = kDefaultZeroPoint2b;
+    if (zero_points) {
+      zp = UnpackZeroPoint2b(zero_points,
+                             static_cast<int64_t>(n_idx) * zero_point_shape_x + kb_idx / kElementsPerByte2b, kb_idx);
+    }
+    zero_point_value = static_cast<float>(zp);
+  } else {
+    zero_point_value = zero_points ? ToFloat2b(*(zero_points + block_id))
+                                   : static_cast<float>(kDefaultZeroPoint2b);
+  }
+
+  DequantizeSixteen2b<T>(quant_value, scale, zero_point_value, output + element_offset);
+}
+
+}  // namespace
+
+template <class T, typename ZeroT>
+Status Dequantize2Bits(
+    T* output,
+    const uint8_t* quant_data,
+    const T* scales_data,
+    const ZeroT* zero_points,
+    const int32_t* reorder_idx,
+    int k,
+    int n,
+    int block_size,
+    cudaStream_t stream) {
+  // k is padded and equal to groups_per_K * block_size. block_size is a power of two >= 16
+  // (enforced by matmul_nbits_helper::CheckInputs), so each thread's 16 codes always fall inside
+  // one quantization block and one aligned uint32 of the packed blob.
+  ORT_ENFORCE(k % block_size == 0, "k must be a multiplier of block_size");
+  ORT_ENFORCE(block_size >= kElementsPerThread2b && (block_size & (block_size - 1)) == 0,
+              "2-bit dequantization requires block_size to be a power of two >= 16, got ", block_size);
+  ORT_RETURN_IF(reorder_idx != nullptr, "CUDA 2-bit dequantization does not support g_idx (reorder_idx).");
+  ORT_ENFORCE(block_size <= 256, "2-bit dequantization block_size must not exceed 256, got ", block_size);
+
+  const int groups_per_K = k / block_size;
+  const int total_groups = n * groups_per_K;
+  const int groups_per_threadblock = GridDim::maxThreadsPerBlock * kElementsPerThread2b / block_size;
+  const int groups_per_grid = CeilDiv(total_groups, groups_per_threadblock);
+  dim3 grid_dim(groups_per_grid);
+  dim3 block_dim(GridDim::maxThreadsPerBlock);
+
+  Dequantize2BitsKernel<T, ZeroT><<<grid_dim, block_dim, 0, stream>>>(
+      output,
+      quant_data,
+      scales_data,
+      zero_points,
+      block_size,
+      groups_per_K,
+      groups_per_threadblock,
+      total_groups);
+
+  return CUDA_CALL(cudaGetLastError());
+}
+
+#define INSTANTIATE_DEQUANTIZE_2BITS(T, ZeroT)                                              \
+  template Status Dequantize2Bits<T, ZeroT>(T * output, const uint8_t* quant_data,          \
+                                            const T* scales_data, const ZeroT* zero_points, \
+                                            const int32_t* reorder_idx, int k, int n,       \
+                                            int block_size, cudaStream_t stream)
+
+INSTANTIATE_DEQUANTIZE_2BITS(float, uint8_t);
+INSTANTIATE_DEQUANTIZE_2BITS(half, uint8_t);
+INSTANTIATE_DEQUANTIZE_2BITS(nv_bfloat16, uint8_t);
+INSTANTIATE_DEQUANTIZE_2BITS(float, float);
+INSTANTIATE_DEQUANTIZE_2BITS(half, half);
+INSTANTIATE_DEQUANTIZE_2BITS(nv_bfloat16, nv_bfloat16);
+
+#undef INSTANTIATE_DEQUANTIZE_2BITS
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime