diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index bcc9b41fd1d3f..9c7976b580b5a 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/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,10 +66,11 @@ 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). @@ -73,6 +78,12 @@ 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
full-N or chunked along N] + DQ[Dequantize blockwise 2b/4b/8b] --> GEMM[cuBLAS GEMM
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 diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h index 68f85ca3b58ed..0114e20182308 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h +++ b/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); diff --git a/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh b/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh index 27fbc48bdb88c..048eeb996b6f9 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh @@ -46,6 +46,18 @@ struct BlkQuantTraits { using ThreadBlk = Shape2D; }; +template +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 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(output, quant_data, scales_data, zero_points, reorder_idx, k, n, block_size, stream); + } else if (bits == 4) { return Dequantize4Bits(output, quant_data, scales_data, zero_points, reorder_idx, k, n, block_size, stream); } else { ORT_ENFORCE(bits == 8); diff --git a/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu b/onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu new file mode 100644 index 0000000000000..01ce31633aab1 --- /dev/null +++ b/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 +#include + +#include +#include + +#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). 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((zero_points[byte_index] >> ((block_index & 0x03) << 1)) & 0x03); +} + +template +__device__ __forceinline__ void StoreSixteen2b(const T* src, T* dst) { + constexpr int kVectors = static_cast(kElementsPerThread2b * sizeof(T) / sizeof(float4)); +#pragma unroll + for (int i = 0; i < kVectors; ++i) { + reinterpret_cast(dst)[i] = reinterpret_cast(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 +__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((values_quant >> (2 * i)) & 0x03u); + results[i] = static_cast(fmaf(code, scale_f, zp_adjust)); + } + StoreSixteen2b(results, output); +} + +template +__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(block_id) * block_size + + ((threadIdx.x * kElementsPerThread2b) & (block_size - 1)); + const uint32_t quant_value = *(reinterpret_cast(quant_data + element_offset / kElementsPerByte2b)); + const T scale = *(scale_data + block_id); + + float zero_point_value; + if constexpr (std::is_same_v) { + 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(n_idx) * zero_point_shape_x + kb_idx / kElementsPerByte2b, kb_idx); + } + zero_point_value = static_cast(zp); + } else { + zero_point_value = zero_points ? ToFloat2b(*(zero_points + block_id)) + : static_cast(kDefaultZeroPoint2b); + } + + DequantizeSixteen2b(quant_value, scale, zero_point_value, output + element_offset); +} + +} // namespace + +template +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<<>>( + 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 * 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 diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu new file mode 100644 index 0000000000000..ba3de13d9a9ab --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cuda_common.h" +#include "contrib_ops/cuda/quantization/matmul_2bits_batched.cuh" +#include "contrib_ops/cuda/quantization/matmul_2bits_m1.cuh" +#include "contrib_ops/cuda/quantization/matmul_nbits.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +bool TryMatMul2Bits( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + int device_sm, + size_t shared_mem_per_block, + cudaStream_t stream) { + constexpr int kColsPerThreadBlock = 8; + constexpr int kElementsPerThreadPerIteration = 16; + constexpr int kWarpSize = onnxruntime::cuda::GPU_WARP_SIZE; + + if constexpr (std::is_same_v) { + if (device_sm < 53) { + return false; + } + } + + // Above this m the dequantize + cuBLAS path wins, and it is the only correct path for a K that + // is not a multiple of block_size or an N that is not a multiple of the warp tile. + constexpr int kSmallMMax = 8; + if (m < 1 || m > kSmallMMax || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { + return false; + } + constexpr int kPerIter = kWarpSize * kElementsPerThreadPerIteration; + if (block_size < kElementsPerThreadPerIteration || kPerIter % block_size != 0 || k % block_size != 0) { + return false; + } + + if (m >= 2) { + return TryMatMul2BitsBatched(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, stream); + } + + return TryMatMul2BitsM1(output, a_data, b_data_quant, scales_data, zero_points, + n, k, block_size, shared_mem_per_block, stream); +} + +template bool TryMatMul2Bits( + float*, const float*, const uint8_t*, const float*, const uint8_t*, int, int, int, int, int, size_t, cudaStream_t); + +template bool TryMatMul2Bits( + half*, const half*, const uint8_t*, const half*, const uint8_t*, int, int, int, int, int, size_t, cudaStream_t); + +template bool TryMatMul2Bits( + nv_bfloat16*, const nv_bfloat16*, const uint8_t*, const nv_bfloat16*, const uint8_t*, int, int, int, int, + int, size_t, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched.cuh new file mode 100644 index 0000000000000..298c082200dc5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched.cuh @@ -0,0 +1,29 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +bool TryMatMul2BitsBatched( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_bfloat16.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_bfloat16.cu new file mode 100644 index 0000000000000..f13ca73db4619 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_bfloat16.cu @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsBatched( + nv_bfloat16*, const nv_bfloat16*, const uint8_t*, const nv_bfloat16*, const uint8_t*, + int, int, int, int, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_float.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_float.cu new file mode 100644 index 0000000000000..f62571620483b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_float.cu @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsBatched( + float*, const float*, const uint8_t*, const float*, const uint8_t*, + int, int, int, int, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_half.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_half.cu new file mode 100644 index 0000000000000..746818351965e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_half.cu @@ -0,0 +1,16 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsBatched( + half*, const half*, const uint8_t*, const half*, const uint8_t*, + int, int, int, int, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh new file mode 100644 index 0000000000000..5aadf72adf9df --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "contrib_ops/cuda/quantization/matmul_2bits_batched.cuh" +#include "contrib_ops/cuda/quantization/matmul_2bits_common.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { +namespace { + +constexpr int kSmallMMax2b = 8; + +// Register-tiled small-M kernel: each warp owns CtaN output channels and CtaM rows of A, so one +// dequantized weight chunk is reused across CtaM rows and one activation chunk across CtaN +// channels. Scales and zero points are read straight from global memory; unlike the m1 kernel +// there is no reuse across warps to justify staging them. +template +__global__ void __launch_bounds__(kWarpSize2b* kColsPerThreadBlock2b, 3) MatMulFloat2bKernelBatched( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int blocks_per_K) { + using Acc = typename Traits2b::Acc; + const int lane_id = threadIdx.x; + const int warp_id = threadIdx.y; + const int col_base = (blockIdx.x * kColsPerThreadBlock2b + warp_id) * CtaN; + const int m_base = blockIdx.y * CtaM; + const int valid = m - m_base; + constexpr int k_per_iter = kWarpSize2b * kElementsPerThreadPerIteration2b; + const int lane_offset = lane_id * kElementsPerThreadPerIteration2b; + const int zp_bytes_per_col = (blocks_per_K + kElementsPerByte2b - 1) / kElementsPerByte2b; + + const T* a_base = a_data + static_cast(m_base) * k + lane_offset; + const uint8_t* b_ptr[CtaN]; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + b_ptr[c] = b_data_quant + + static_cast(col_base + c) * blocks_per_K * (block_size / kElementsPerByte2b) + lane_offset / kElementsPerByte2b; + } + + Acc acc[CtaM][CtaN]; +#pragma unroll + for (int r = 0; r < CtaM; ++r) { +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + acc[r][c] = Acc{}; + } + } + + int k_id = 0; + for (; k_id + k_per_iter <= k; k_id += k_per_iter) { + const int blk = (lane_offset + k_id) / block_size; + typename Traits2b::Weights w[CtaN]; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + uint8_t zp = kDefaultZeroPoint2b; + if constexpr (has_zero_point) { + zp = UnpackZeroPoint2b(zero_points + static_cast(col_base + c) * zp_bytes_per_col, blk); + } + DequantizeSixteen2b(*reinterpret_cast(b_ptr[c] + k_id / kElementsPerByte2b), + scales_data[static_cast(col_base + c) * blocks_per_K + blk], zp, w[c]); + } +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r < valid) { + typename Traits2b::Acts av; + LoadSixteen2b(a_base + static_cast(r) * k + k_id, av); +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + DotAccum2b(w[c], av, acc[r][c]); + } + } + } + } + + if (lane_offset + k_id < k) { + const int blk = (lane_offset + k_id) / block_size; + typename Traits2b::Weights w[CtaN]; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + uint8_t zp = kDefaultZeroPoint2b; + if constexpr (has_zero_point) { + zp = UnpackZeroPoint2b(zero_points + static_cast(col_base + c) * zp_bytes_per_col, blk); + } + DequantizeSixteen2b(*reinterpret_cast(b_ptr[c] + k_id / kElementsPerByte2b), + scales_data[static_cast(col_base + c) * blocks_per_K + blk], zp, w[c]); + } +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r < valid) { + typename Traits2b::Acts av; + LoadSixteen2b(a_base + static_cast(r) * k + k_id, av); +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + DotAccum2b(w[c], av, acc[r][c]); + } + } + } + } + +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r >= valid) continue; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + float sum = HorizontalAdd2b(acc[r][c]); + for (int offset = kWarpSize2b / 2; offset > 0; offset /= 2) { + sum += onnxruntime::cuda::WARP_SHFL_DOWN(sum, offset); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + (col_base + c)] = static_cast(sum); + } + } + } +} + +} // namespace + +template +bool TryMatMul2BitsBatched( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + cudaStream_t stream) { + if (m < 2 || m > kSmallMMax2b) { + return false; + } + const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 + : 8; + const int cta_n = (n % (kColsPerThreadBlock2b * 2) == 0) ? 2 : 1; + const int blocks_per_K = k / block_size; + dim3 threads(onnxruntime::cuda::GPU_WARP_SIZE_HOST, kColsPerThreadBlock2b); + dim3 blocks(n / (kColsPerThreadBlock2b * cta_n), (m + cta_m - 1) / cta_m); + +#define MatMulFloat2bBatchedDispatch(bs, cm, cn) \ + if (nullptr != zero_points) { \ + MatMulFloat2bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, blocks_per_K); \ + } else { \ + MatMulFloat2bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, blocks_per_K); \ + } +#define MatMulFloat2bBatchedDispatchN(cm, cn) \ + if (16 == block_size) { \ + MatMulFloat2bBatchedDispatch(16, cm, cn) \ + } else if (32 == block_size) { \ + MatMulFloat2bBatchedDispatch(32, cm, cn) \ + } else if (64 == block_size) { \ + MatMulFloat2bBatchedDispatch(64, cm, cn) \ + } else if (128 == block_size) { \ + MatMulFloat2bBatchedDispatch(128, cm, cn) \ + } else if (256 == block_size) { \ + MatMulFloat2bBatchedDispatch(256, cm, cn) \ + } else { \ + return false; \ + } +#define MatMulFloat2bBatchedDispatchM(cn) \ + switch (cta_m) { \ + case 2: \ + MatMulFloat2bBatchedDispatchN(2, cn) break; \ + case 4: \ + MatMulFloat2bBatchedDispatchN(4, cn) break; \ + default: \ + MatMulFloat2bBatchedDispatchN(8, cn) break; \ + } + + if (cta_n == 2) { + MatMulFloat2bBatchedDispatchM(2) + } else { + MatMulFloat2bBatchedDispatchM(1) + } + +#undef MatMulFloat2bBatchedDispatchM +#undef MatMulFloat2bBatchedDispatchN +#undef MatMulFloat2bBatchedDispatch + return true; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_bfloat16.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_bfloat16.cu new file mode 100644 index 0000000000000..c2a23e8ec6e1e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_bfloat16.cu @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsM1( + nv_bfloat16*, const nv_bfloat16*, const uint8_t*, const nv_bfloat16*, const uint8_t*, int, int, int, size_t, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh new file mode 100644 index 0000000000000..6bfd5c6afd968 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// Shared geometry for the 2-bit GEMV kernels. One uint32 of the packed weight blob holds 16 +// consecutive codes, so a thread consumes 16 elements of K per iteration and a warp consumes 512. +// block_size is a power of two in [16, 256] for every shape these kernels accept, so a thread's +// 16 codes always belong to a single quantization block. +constexpr int kColsPerThreadBlock2b = 8; +constexpr int kElementsPerThreadPerIteration2b = 16; +constexpr int kWarpSize2b = onnxruntime::cuda::GPU_WARP_SIZE; +constexpr int kElementsPerByte2b = 4; +constexpr uint8_t kDefaultZeroPoint2b = 2; // 1 << (bits - 1) + +template +__device__ __forceinline__ float ToFloat2b(T v); +template <> +__device__ __forceinline__ float ToFloat2b(float v) { return v; } +template <> +__device__ __forceinline__ float ToFloat2b(half v) { return __half2float(v); } +template <> +__device__ __forceinline__ float ToFloat2b(nv_bfloat16 v) { return __bfloat162float(v); } + +// Reads the 2-bit zero point of quantization block `block_index` from one output channel's +// zero-point row. The uint8 row is bit-packed four blocks to a byte, low crumb first. +__device__ __forceinline__ uint8_t UnpackZeroPoint2b(const uint8_t* zero_points_row, int block_index) { + return static_cast((zero_points_row[block_index >> 2] >> ((block_index & 0x03) << 1)) & 0x03); +} + +// --------------------------------------------------------------------------------------------- +// Per-activation-type traits. A 16-element chunk of dequantized weights (Weights) and the matching +// activations (Acts) are held in whatever lane order makes the dot product cheapest, so callers +// never need to know which representation was chosen. +// +// Instruction count matters even though the GEMV is nominally memory bound: at 2 bits per weight +// there are 4 multiply-accumulates per byte of B, four times the 8-bit ratio, so a naive +// per-element integer-to-float conversion makes the kernel instruction bound. +// --------------------------------------------------------------------------------------------- +template +struct Traits2b { + struct Weights { + float v[kElementsPerThreadPerIteration2b]; + }; + struct Acts { + float v[kElementsPerThreadPerIteration2b]; + }; + using Acc = float; +}; + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +template <> +struct Traits2b { + // Lane j holds elements (j, j + 8). That pairing falls straight out of masking the shifted + // packed word with 0x00030003, which is what keeps the expansion at two instructions per lane. + struct Weights { + half2 v[8]; + }; + struct Acts { + half2 v[8]; + }; + using Acc = half2; +}; +#endif + +// Generic fp32 path, used for float and bfloat16 activations. +template +__device__ __forceinline__ void DequantizeSixteenGeneric2b(uint32_t values_quant, T scale, uint8_t zp, + float (&dq)[kElementsPerThreadPerIteration2b]) { + const float scale_f = ToFloat2b(scale); + const float zp_adjust = -scale_f * static_cast(zp); +#pragma unroll + for (int i = 0; i < kElementsPerThreadPerIteration2b; ++i) { + dq[i] = fmaf(static_cast((values_quant >> (2 * i)) & 0x03u), scale_f, zp_adjust); + } +} + +template +__device__ __forceinline__ void LoadSixteenGeneric2b(const T* a, float (&out)[kElementsPerThreadPerIteration2b]) { +#pragma unroll + for (int i = 0; i < kElementsPerThreadPerIteration2b; ++i) { + out[i] = ToFloat2b(a[i]); + } +} + +__device__ __forceinline__ void DotAccumGeneric2b(const float (&w)[kElementsPerThreadPerIteration2b], + const float (&a)[kElementsPerThreadPerIteration2b], float& acc) { +#pragma unroll + for (int i = 0; i < kElementsPerThreadPerIteration2b; ++i) { + acc = fmaf(w[i], a[i], acc); + } +} + +__device__ __forceinline__ void DequantizeSixteen2b(uint32_t values_quant, float scale, uint8_t zp, + Traits2b::Weights& w) { + DequantizeSixteenGeneric2b(values_quant, scale, zp, w.v); +} +__device__ __forceinline__ void LoadSixteen2b(const float* a, Traits2b::Acts& out) { + LoadSixteenGeneric2b(a, out.v); +} +__device__ __forceinline__ void DotAccum2b(const Traits2b::Weights& w, const Traits2b::Acts& a, + float& acc) { + DotAccumGeneric2b(w.v, a.v, acc); +} + +__device__ __forceinline__ void DequantizeSixteen2b(uint32_t values_quant, nv_bfloat16 scale, uint8_t zp, + Traits2b::Weights& w) { + DequantizeSixteenGeneric2b(values_quant, scale, zp, w.v); +} +__device__ __forceinline__ void LoadSixteen2b(const nv_bfloat16* a, Traits2b::Acts& out) { + LoadSixteenGeneric2b(a, out.v); +} +__device__ __forceinline__ void DotAccum2b(const Traits2b::Weights& w, + const Traits2b::Acts& a, float& acc) { + DotAccumGeneric2b(w.v, a.v, acc); +} + +__device__ __forceinline__ float HorizontalAdd2b(float acc) { return acc; } + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +// 0x6400 is half(1024). OR-ing a 2-bit code into its low mantissa bits yields half(1024 + code) +// exactly, so subtracting half(1024) recovers the code with no rounding. Masking the shifted word +// with 0x00030003 extracts codes j and j + 8 together, one LOP3 for the pair. +__device__ __forceinline__ void DequantizeSixteen2b(uint32_t values_quant, half scale, uint8_t zp, + Traits2b::Weights& w) { + constexpr uint32_t kImmLut = (0xf0 & 0xcc) | 0xaa; // (a & b) | c + constexpr uint32_t kCrumbMask = 0x00030003; + constexpr uint32_t kMagicNum = 0x64006400; + + const half2 scale_h2 = __half2half2(scale); + const half2 zp_adjust_h2 = __half2half2(__hneg(scale) * __ushort2half_rn(zp)); + + uint32_t* h = reinterpret_cast(w.v); +#pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t shifted = values_quant >> (2 * j); + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[j]) + : "r"(shifted), "n"(kCrumbMask), "n"(kMagicNum), "n"(kImmLut)); + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[j]) : "r"(h[j]), "r"(kMagicNum)); + w.v[j] = __hfma2(w.v[j], scale_h2, zp_adjust_h2); + } +} + +// Repacks 16 contiguous activations into the (j, j + 8) lane order used by Weights. +__device__ __forceinline__ void LoadSixteen2b(const half* a, Traits2b::Acts& out) { + constexpr uint32_t kLowHalf2 = 0x5410; + constexpr uint32_t kHighHalf2 = 0x7632; + const uint4 lo = *reinterpret_cast(a); + const uint4 hi = *reinterpret_cast(a + 8); + const uint32_t words[8] = {lo.x, lo.y, lo.z, lo.w, hi.x, hi.y, hi.z, hi.w}; + uint32_t* o = reinterpret_cast(out.v); +#pragma unroll + for (int i = 0; i < 4; ++i) { + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(o[2 * i]) : "r"(words[i]), "r"(words[i + 4]), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(o[2 * i + 1]) : "r"(words[i]), "r"(words[i + 4]), "r"(kHighHalf2)); + } +} + +__device__ __forceinline__ void DotAccum2b(const Traits2b::Weights& w, const Traits2b::Acts& a, + half2& acc) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + acc = __hfma2(w.v[j], a.v[j], acc); + } +} + +__device__ __forceinline__ float HorizontalAdd2b(half2 acc) { + return __half2float(acc.x) + __half2float(acc.y); +} +#else +__device__ __forceinline__ void DequantizeSixteen2b(uint32_t values_quant, half scale, uint8_t zp, + Traits2b::Weights& w) { + DequantizeSixteenGeneric2b(values_quant, scale, zp, w.v); +} +__device__ __forceinline__ void LoadSixteen2b(const half* a, Traits2b::Acts& out) { + LoadSixteenGeneric2b(a, out.v); +} +__device__ __forceinline__ void DotAccum2b(const Traits2b::Weights& w, const Traits2b::Acts& a, + float& acc) { + DotAccumGeneric2b(w.v, a.v, acc); +} +#endif + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_float.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_float.cu new file mode 100644 index 0000000000000..b1d310e0ba771 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_float.cu @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsM1( + float*, const float*, const uint8_t*, const float*, const uint8_t*, int, int, int, size_t, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_half.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_half.cu new file mode 100644 index 0000000000000..af5a47286d52b --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_half.cu @@ -0,0 +1,15 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template bool TryMatMul2BitsM1( + half*, const half*, const uint8_t*, const half*, const uint8_t*, int, int, int, size_t, cudaStream_t); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1.cuh new file mode 100644 index 0000000000000..ac83845e0526e --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1.cuh @@ -0,0 +1,30 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +bool TryMatMul2BitsM1( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int n, + int k, + int block_size, + size_t shared_mem_per_block, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh new file mode 100644 index 0000000000000..5b8c71b7c10ef --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh @@ -0,0 +1,156 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "contrib_ops/cuda/quantization/matmul_2bits_common.cuh" +#include "contrib_ops/cuda/quantization/matmul_2bits_m1.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// One warp per output channel, one row of A. Scales and zero points for the eight channels owned +// by the thread block are staged in shared memory; the zero points are unpacked to one byte per +// block there so the main loop indexes them without any bit arithmetic. +template +__global__ void __launch_bounds__(kWarpSize2b* kColsPerThreadBlock2b) MatMulFloat2bKernelM1( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int n, + int k, + int blocks_per_K) { + const int lane_id = threadIdx.x; + const int warp_id = threadIdx.y; + const int n_block_id = blockIdx.x; + const int n_id = n_block_id * kColsPerThreadBlock2b + warp_id; + const int meta_count = kColsPerThreadBlock2b * blocks_per_K; + + extern __shared__ char shared_buffer[]; + T* b_scale_vec = reinterpret_cast(shared_buffer); + uint8_t* b_zp_vec = reinterpret_cast(b_scale_vec + meta_count); + + const int64_t scale_offset = static_cast(n_block_id) * kColsPerThreadBlock2b * blocks_per_K; + for (int i = warp_id * kWarpSize2b + lane_id; i < meta_count; i += kColsPerThreadBlock2b * kWarpSize2b) { + b_scale_vec[i] = scales_data[scale_offset + i]; + } + if constexpr (has_zero_point) { + const int zp_bytes_per_col = (blocks_per_K + kElementsPerByte2b - 1) / kElementsPerByte2b; + for (int i = warp_id * kWarpSize2b + lane_id; i < meta_count; i += kColsPerThreadBlock2b * kWarpSize2b) { + const int col = i / blocks_per_K; + const int blk = i % blocks_per_K; + const uint8_t* zp_row = + zero_points + static_cast(n_block_id * kColsPerThreadBlock2b + col) * zp_bytes_per_col; + b_zp_vec[i] = UnpackZeroPoint2b(zp_row, blk); + } + } + __syncthreads(); + + const T* scale_thread = b_scale_vec + warp_id * blocks_per_K; + const uint8_t* zp_thread = b_zp_vec + warp_id * blocks_per_K; + const int lane_offset = lane_id * kElementsPerThreadPerIteration2b; + const uint8_t* b_column = + b_data_quant + static_cast(n_id) * blocks_per_K * (block_size / kElementsPerByte2b); + + typename Traits2b::Acc acc{}; + constexpr int k_per_iter = kWarpSize2b * kElementsPerThreadPerIteration2b; + + int k_id = 0; + for (; k_id + k_per_iter <= k; k_id += k_per_iter) { + const int k_offset = lane_offset + k_id; + const int blk = k_offset / block_size; + typename Traits2b::Weights w; + typename Traits2b::Acts av; + DequantizeSixteen2b(*reinterpret_cast(b_column + k_offset / kElementsPerByte2b), + scale_thread[blk], has_zero_point ? zp_thread[blk] : kDefaultZeroPoint2b, w); + LoadSixteen2b(a_data + k_offset, av); + DotAccum2b(w, av, acc); + } + if (lane_offset + k_id < k) { + const int k_offset = lane_offset + k_id; + const int blk = k_offset / block_size; + typename Traits2b::Weights w; + typename Traits2b::Acts av; + DequantizeSixteen2b(*reinterpret_cast(b_column + k_offset / kElementsPerByte2b), + scale_thread[blk], has_zero_point ? zp_thread[blk] : kDefaultZeroPoint2b, w); + LoadSixteen2b(a_data + k_offset, av); + DotAccum2b(w, av, acc); + } + + float sum = HorizontalAdd2b(acc); + for (int offset = kWarpSize2b / 2; offset > 0; offset /= 2) { + sum += onnxruntime::cuda::WARP_SHFL_DOWN(sum, offset); + } + if (lane_id == 0) { + output[n_id] = static_cast(sum); + } +} + +template +bool TryMatMul2BitsM1( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int n, + int k, + int block_size, + size_t shared_mem_per_block, + cudaStream_t stream) { + if (n % kColsPerThreadBlock2b != 0 || k % kElementsPerThreadPerIteration2b != 0) { + return false; + } + // A thread's 16 codes must lie inside one quantization block, and the trailing partial block of + // a K that is not a multiple of block_size is left to the dequantize + cuBLAS fallback. + constexpr int k_per_iter = kWarpSize2b * kElementsPerThreadPerIteration2b; + if (block_size < kElementsPerThreadPerIteration2b || k_per_iter % block_size != 0 || k % block_size != 0) { + return false; + } + + const int blocks_per_K = k / block_size; + const size_t shared_mem_size = (sizeof(T) + (zero_points != nullptr ? sizeof(uint8_t) : 0)) * + static_cast(blocks_per_K) * kColsPerThreadBlock2b; + if (shared_mem_size > shared_mem_per_block) { + return false; + } + + dim3 threads(onnxruntime::cuda::GPU_WARP_SIZE_HOST, kColsPerThreadBlock2b); + dim3 blocks(n / kColsPerThreadBlock2b, 1); + +#define MATMUL_FLOAT2B_M1_DISPATCH(bs) \ + if (zero_points != nullptr) { \ + MatMulFloat2bKernelM1<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, n, k, blocks_per_K); \ + } else { \ + MatMulFloat2bKernelM1<<>>( \ + output, a_data, b_data_quant, scales_data, nullptr, n, k, blocks_per_K); \ + } + + if (block_size == 16) { + MATMUL_FLOAT2B_M1_DISPATCH(16) + } else if (block_size == 32) { + MATMUL_FLOAT2B_M1_DISPATCH(32) + } else if (block_size == 64) { + MATMUL_FLOAT2B_M1_DISPATCH(64) + } else if (block_size == 128) { + MATMUL_FLOAT2B_M1_DISPATCH(128) + } else if (block_size == 256) { + MATMUL_FLOAT2B_M1_DISPATCH(256) + } else { + return false; + } +#undef MATMUL_FLOAT2B_M1_DISPATCH + return true; +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 43ca37ed82ee9..f6c6e5a64fb1d 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -813,6 +813,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); + ORT_RETURN_IF(nbits_ == 2 && reorder_idx != nullptr, + "CUDA MatMulNBits does not support g_idx (reorder_idx) for 2-bit weights."); const auto* a_data = a->Data(); const auto* reorder_idx_data = reorder_idx == nullptr ? nullptr : reorder_idx->Data(); @@ -980,7 +982,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { k, SafeInt(block_size_), GetDeviceProp().sharedMemPerBlock, - stream)) { + stream, + sm_)) { return Status::OK(); } @@ -1001,7 +1004,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { k, SafeInt(block_size_), GetDeviceProp().sharedMemPerBlock, - stream)) { + stream, + sm_)) { LaunchMatMulNBitsBiasAdd( reinterpret_cast(Y->MutableData()), reinterpret_cast(bias_data), @@ -1041,7 +1045,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { IAllocatorUniquePtr b_data_ptr = this->template GetScratchBuffer(scratch_n * K_padded, this->GetComputeStream(ctx)); auto* b_data = b_data_ptr.get(); - // Column-wise dequant helper: dispatches 8-bit / 4-bit × typed / uint8 zero-points. + // Column-wise dequant helper: dispatches 8/4/2-bit x typed / uint8 zero-points. // Used by both the full-N and chunked paths so the offset math stays in one place. const int64_t blocks_per_col = K_padded / block_size_; auto dequant_column_wise = [&](const uint8_t* chunk_blob, @@ -1049,31 +1053,18 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const void* chunk_zp, const int32_t* chunk_reorder_idx, int n_rows) -> Status { - if (nbits_ == 8) { - if (zero_points && zero_points->IsDataType()) { - return Dequantize8Bits( - reinterpret_cast(b_data), chunk_blob, chunk_scales, - static_cast(chunk_zp), chunk_reorder_idx, - SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); - } else { - return Dequantize8Bits( - reinterpret_cast(b_data), chunk_blob, chunk_scales, - static_cast(chunk_zp), chunk_reorder_idx, - SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); - } - } else { - if (zero_points && zero_points->IsDataType()) { - return Dequantize4Bits( - reinterpret_cast(b_data), chunk_blob, chunk_scales, - static_cast(chunk_zp), chunk_reorder_idx, - SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); - } else { - return Dequantize4Bits( - reinterpret_cast(b_data), chunk_blob, chunk_scales, - static_cast(chunk_zp), chunk_reorder_idx, - SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); - } + const bool typed_zero_points = zero_points && zero_points->IsDataType(); + if (typed_zero_points) { + return DequantizeNBits( + SafeInt(nbits_), reinterpret_cast(b_data), chunk_blob, chunk_scales, + static_cast(chunk_zp), chunk_reorder_idx, + SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); } + + return DequantizeNBits( + SafeInt(nbits_), reinterpret_cast(b_data), chunk_blob, chunk_scales, + static_cast(chunk_zp), chunk_reorder_idx, + SafeInt(K_padded), n_rows, SafeInt(block_size_), stream); }; // Skip full dequant when chunked path will handle it in the GEMM loop @@ -1102,6 +1093,9 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { stream)); } else { // row-wise block (4-bit) + ORT_RETURN_IF_NOT(nbits_ == 4, + "CUDA MatMulNBits row-wise quantization blocks are only implemented for bits = 4 or 8, but got bits = ", + nbits_); K_padded = K_; ORT_RETURN_IF_ERROR(DequantizeBlockwise4b( reinterpret_cast(b_data), @@ -1141,17 +1135,17 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { ORT_ENFORCE(column_wise_quant_blk_, "Chunked path requires column-wise quantization blocks"); // Compute per-chunk pointers into the column-wise-packed weight, scale, and ZP arrays. - const uint8_t* chunk_blob = blob_data + n_start * (nbits_ == 8 ? K_padded : K_padded / 2); + const int64_t elements_per_byte = 8 / nbits_; + const uint8_t* chunk_blob = blob_data + n_start * (K_padded / elements_per_byte); const auto* chunk_scales = reinterpret_cast(scales_data) + n_start * blocks_per_col; const void* chunk_zp = nullptr; if (zero_points_data) { if (zero_points && zero_points->IsDataType()) { chunk_zp = reinterpret_cast(zero_points_data) + n_start * blocks_per_col; - } else if (nbits_ == 8) { - chunk_zp = static_cast(zero_points_data) + n_start * blocks_per_col; } else { - // 4-bit ZP: packed, offset by n_start * ceil(blocks_per_col / 2) - chunk_zp = static_cast(zero_points_data) + n_start * ((blocks_per_col + 1) / 2); + // uint8 ZP rows are bit-packed: ceil(blocks_per_col / elements_per_byte) bytes per output channel. + chunk_zp = static_cast(zero_points_data) + + n_start * ((blocks_per_col + elements_per_byte - 1) / elements_per_byte); } } ORT_RETURN_IF_ERROR(dequant_column_wise( diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index 61eb62fef4a2b..71c57a85f07c4 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -8,6 +8,21 @@ namespace onnxruntime { namespace contrib { namespace cuda { +template +bool TryMatMul2Bits( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const uint8_t* zero_points, + int m, + int n, + int k, + int block_size, + int device_sm, + size_t shared_mem_per_block, + cudaStream_t stream); + template bool TryMatMul4Bits( T* output, @@ -51,7 +66,8 @@ bool TryMatMulNBits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream) { + cudaStream_t stream, + int device_sm) { if (bits == 8) { if (bias_data != nullptr) { return false; @@ -65,6 +81,14 @@ bool TryMatMulNBits( m, n, k, block_size, shared_mem_per_block, stream); } + if (bits == 2) { + if (bias_data != nullptr) { + return false; + } + return TryMatMul2Bits(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, device_sm, shared_mem_per_block, stream); + } + return false; } diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 4246f0c9634df..5638f03232e75 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -144,6 +144,12 @@ class MatMulNBits final : public CudaKernel { ORT_ENFORCE(Status::OK() == info.GetAttr("block_size", &block_size_)); ORT_ENFORCE(Status::OK() == info.GetAttr("bits", &nbits_)); ORT_ENFORCE(block_size_ > 0, "block_size must be greater than zero"); + // The op schema and matmul_nbits_helper::CheckInputs accept bits in {2, 4, 8}. Keep this in + // lockstep with the widths the CUDA kernels actually implement: without it a node with an + // unimplemented width is accepted and then read with the wrong stride, silently producing + // garbage instead of failing. + ORT_ENFORCE(nbits_ == 2 || nbits_ == 4 || nbits_ == 8, + "MatMulNBits on the CUDA execution provider supports bits = 2, 4 or 8, but got bits = ", nbits_); constexpr int kInputIndexScale = 2; constexpr int kInputIndexZeroPoints = 3; diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc index 8620745c0c5eb..4e825219faa1b 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc @@ -521,7 +521,8 @@ class KernelTestFixture : public ::testing::Test { reinterpret_cast(d_scales_->data()), static_cast(d_uint8_zeros.data()), static_cast(nullptr), - m_, n_, k_, block_size_, device_prop_.sharedMemPerBlock, s_); + m_, n_, k_, block_size_, device_prop_.sharedMemPerBlock, s_, + device_prop_.major * 10 + device_prop_.minor); }, warmup_, repeats_, s_); } diff --git a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc index 025c52b86f181..f9437b63aa67f 100644 --- a/onnxruntime/test/contrib_ops/matmul_2bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_2bits_test.cc @@ -5,6 +5,8 @@ #include #include +#include +#include #include "gtest/gtest.h" #include "gmock/gmock.h" @@ -53,11 +55,16 @@ struct TestOptions2Bits { int64_t accuracy_level{0}; bool has_zero_point{false}; + bool has_typed_zero_point{false}; bool has_g_idx{false}; bool has_bias{false}; + // Run the node on the CUDA EP instead of the default (CPU / WebGPU) set. + bool use_cuda{false}; + std::optional output_abs_error{}; std::optional output_rel_error{}; + std::optional expected_failure{}; }; [[maybe_unused]] std::ostream& operator<<(std::ostream& os, const TestOptions2Bits& opts) { @@ -65,13 +72,16 @@ struct TestOptions2Bits { << ", block_size:" << opts.block_size << ", accuracy_level:" << opts.accuracy_level << ", has_zero_point:" << opts.has_zero_point + << ", has_typed_zero_point:" << opts.has_typed_zero_point << ", has_g_idx:" << opts.has_g_idx - << ", has_bias:" << opts.has_bias; + << ", has_bias:" << opts.has_bias + << ", use_cuda:" << opts.use_cuda; } template void RunTest2Bits(const TestOptions2Bits& opts) { SCOPED_TRACE(opts); + ASSERT_FALSE(opts.has_typed_zero_point && !opts.has_zero_point); const int64_t M = opts.M, K = opts.K, @@ -101,7 +111,7 @@ void RunTest2Bits(const TestOptions2Bits& opts) { MlasQuantizeBlockwise( input1_vals.data(), scales.data(), - opts.has_zero_point ? zp.data() : nullptr, + opts.has_zero_point && !opts.has_typed_zero_point ? zp.data() : nullptr, input1_fp32_vals.data(), static_cast(opts.block_size), true, @@ -115,7 +125,7 @@ void RunTest2Bits(const TestOptions2Bits& opts) { input1_fp32_vals.data(), input1_vals.data(), scales.data(), - opts.has_zero_point ? zp.data() : nullptr, + opts.has_zero_point && !opts.has_typed_zero_point ? zp.data() : nullptr, static_cast(opts.block_size), true, static_cast(K), @@ -168,13 +178,29 @@ void RunTest2Bits(const TestOptions2Bits& opts) { } if (opts.has_zero_point) { - test.AddInput("zero_points", {N, static_cast(q_zp_size_in_bytes) / N}, zp, true); + if (opts.has_typed_zero_point) { + const std::vector typed_zp(q_scale_size, 2.0f); + if constexpr (std::is_same::value) { + test.AddInput("zero_points", {N, static_cast(q_scale_size) / N}, typed_zp, true); + } else if constexpr (std::is_same::value) { + test.AddInput("zero_points", {N, static_cast(q_scale_size) / N}, + FloatsToMLFloat16s(typed_zp), true); + } else if constexpr (std::is_same::value) { + test.AddInput("zero_points", {N, static_cast(q_scale_size) / N}, + FloatsToBFloat16s(typed_zp), true); + } + } else { + test.AddInput("zero_points", {N, static_cast(q_zp_size_in_bytes) / N}, zp, true); + } } else { test.AddOptionalInputEdge(); } - // Account for deprecated "g_idx" input - test.AddOptionalInputEdge(); + if (opts.has_g_idx) { + test.AddInput("g_idx", {K}, std::vector(static_cast(K), 0), false); + } else { + test.AddOptionalInputEdge(); + } if (bias.has_value()) { if constexpr (std::is_same::value) { @@ -205,6 +231,22 @@ void RunTest2Bits(const TestOptions2Bits& opts) { } std::vector> execution_providers; + if (opts.use_cuda) { +#ifdef USE_CUDA + execution_providers.emplace_back(DefaultCudaExecutionProvider()); + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.ConfigEps(std::move(execution_providers)); + if (opts.expected_failure.has_value()) { + test.Config(OpTester::ExpectResult::kExpectFailure, *opts.expected_failure); + } + test.RunWithConfig(); +#endif + return; + } + if constexpr (std::is_same::value) { #ifdef USE_WEBGPU execution_providers.push_back(DefaultWebGpuExecutionProvider()); @@ -1623,6 +1665,362 @@ TEST(MatMul2BitsWebGpu, Float32_ZeroPoint_DP4A) { #endif // defined(USE_WEBGPU) && !defined(ORT_USE_EP_API_ADAPTERS) +#ifdef USE_CUDA + +namespace { + +// The 2-bit CUDA kernels need at least Pascal-era half support; the GEMV path uses half2 FMA. +bool SkipIfNo2BitCudaDevice() { + return !HasCudaEnvironment(530); +} + +template +void RunCuda2BitsShapes(int64_t block_size, float abs_error, float rel_error) { + // M = 1 exercises the GEMV kernel, 3 the register-tiled small-M kernel with a partially filled + // row tile, 8 a full row tile, and 128 the dequantize + cuBLAS fallback. N = 24 is not a multiple + // of two warp tiles so it also covers the CtaN = 1 dispatch. + constexpr int64_t kRows[] = {1, 3, 8, 128}; + constexpr int64_t kShapes[][2] = {{32, 256}, {64, 512}, {24, 1024}}; + for (int64_t m : kRows) { + for (bool has_zero_point : {false, true}) { + for (const auto& nk : kShapes) { + TestOptions2Bits opts{}; + opts.M = m; + opts.N = nk[0]; + opts.K = nk[1]; + opts.block_size = block_size; + opts.has_zero_point = has_zero_point; + opts.use_cuda = true; + opts.output_abs_error = abs_error; + opts.output_rel_error = rel_error; + RunTest2Bits(opts); + } + } + } +} + +} // namespace + +TEST(MatMul2BitsCuda, Float16_BlkLen16) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(16, 0.1f, 0.02f); +} + +TEST(MatMul2BitsCuda, Float16_BlkLen32) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(32, 0.1f, 0.02f); +} + +TEST(MatMul2BitsCuda, Float16_BlkLen64) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(64, 0.1f, 0.02f); +} + +TEST(MatMul2BitsCuda, Float16_BlkLen128) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(128, 0.1f, 0.02f); +} + +TEST(MatMul2BitsCuda, Float16_BlkLen256) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(256, 0.1f, 0.02f); +} + +TEST(MatMul2BitsCuda, Float32_BlkLen32) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(32, 0.05f, 0.01f); +} + +TEST(MatMul2BitsCuda, Float32_BlkLen128) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + RunCuda2BitsShapes(128, 0.05f, 0.01f); +} + +TEST(MatMul2BitsCuda, BFloat16_BlkLen128) { + if (!CudaHasBF16Support()) GTEST_SKIP() << "CUDA device does not support BFloat16"; + RunCuda2BitsShapes(128, 0.3f, 0.05f); +} + +// A K that is not a multiple of block_size forces the GEMV kernels to decline and the +// dequantize + cuBLAS fallback to handle the padded trailing block. +TEST(MatMul2BitsCuda, Float16_UnalignedK) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + for (int64_t m : {int64_t{1}, int64_t{4}, int64_t{64}}) { + for (bool has_zero_point : {false, true}) { + TestOptions2Bits opts{}; + opts.M = m; + opts.N = 32; + opts.K = 320; // 2.5 blocks of 128 + opts.block_size = 128; + opts.has_zero_point = has_zero_point; + opts.use_cuda = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + } + } +} + +TEST(MatMul2BitsCuda, Float16_Bias) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + for (int64_t m : {int64_t{1}, int64_t{4}, int64_t{64}}) { + TestOptions2Bits opts{}; + opts.M = m; + opts.N = 32; + opts.K = 512; + opts.block_size = 128; + opts.has_zero_point = true; + opts.has_bias = true; + opts.use_cuda = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + } +} + +TEST(MatMul2BitsCuda, Float16_TypedZeroPointFallback) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + TestOptions2Bits opts{}; + opts.M = 128; + opts.N = 32; + opts.K = 256; + opts.block_size = 128; + opts.has_zero_point = true; + opts.has_typed_zero_point = true; + opts.use_cuda = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); +} + +TEST(MatMul2BitsCuda, RejectsGroupIndex) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + TestOptions2Bits opts{}; + opts.M = 1; + opts.N = 32; + opts.K = 256; + opts.block_size = 128; + opts.has_g_idx = true; + opts.use_cuda = true; + opts.expected_failure = "does not support g_idx"; + RunTest2Bits(opts); +} + +TEST(MatMul2BitsCuda, RejectsOversizedBlock) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + + constexpr int64_t M = 1, N = 1, K = 512, block_size = 512; + constexpr int64_t k_blocks = K / block_size; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", static_cast(QBits)); + test.AddAttribute("accuracy_level", static_cast(0)); + + test.AddInput("A", {M, K}, FloatsToMLFloat16s(std::vector(M * K, 0.5f)), false); + test.AddInput("B", {N, k_blocks, block_size / 4}, + std::vector(static_cast(N * k_blocks * block_size / 4), 0x55), true); + test.AddInput("scales", {N, k_blocks}, + FloatsToMLFloat16s(std::vector(static_cast(N * k_blocks), 0.25f)), true); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("Y", {M, N}, FloatsToMLFloat16s(std::vector(M * N, 0.0f))); + + std::vector> execution_providers; + execution_providers.emplace_back(DefaultCudaExecutionProvider()); + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.Config(OpTester::ExpectResult::kExpectFailure, "block_size must not exceed 256") + .ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + +TEST(MatMul2BitsCuda, ChunkedFallback) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{ + {"ORT_MATMULNBITS_FORCE_CHUNKED", "1"}, + {"ORT_MATMULNBITS_CHUNK_SIZE", "64"}}}; + + for (bool has_zero_point : {false, true}) { + TestOptions2Bits opts{}; + opts.M = 2; + opts.N = 130; + opts.K = 384; + opts.block_size = 128; + opts.has_zero_point = has_zero_point; + opts.use_cuda = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest2Bits(opts); + } +} + +namespace { + +// Ternary weights {-1, 0, +1} map exactly onto the affine 2-bit ABI as codes {0, 1, 2} with a +// uniform zero point of 1. MlasQuantizeBlockwise never emits that subset, so the node is built by +// hand here and checked against an exact float reference; every quantity is chosen to be +// representable in fp16 so the only error left is accumulation order. +void RunTernary2BitsCudaTest(int64_t M, int64_t N, int64_t K, int64_t block_size, + float abs_error = 0.05f, float rel_error = 0.01f) { + SCOPED_TRACE("ternary M:" + std::to_string(M) + " N:" + std::to_string(N) + + " K:" + std::to_string(K) + " block_size:" + std::to_string(block_size)); + ASSERT_EQ(K % block_size, 0) << "this helper assumes an aligned K"; + + const int64_t k_blocks = K / block_size; + const int64_t blob_size = block_size / 4; + const int64_t zp_bytes = (k_blocks * 2 + 7) / 8; + + // Codes cycle through {0, 1, 2} with a stride that is co-prime with the block size so every + // block sees all three symbols; code 3 never appears, as in a real ternary checkpoint. + std::vector b_packed(static_cast(N * k_blocks * blob_size), 0); + std::vector ternary(static_cast(N * K)); + for (int64_t n = 0; n < N; ++n) { + for (int64_t k = 0; k < K; ++k) { + const uint8_t code = static_cast((7 * n + 3 * k + (k / 5)) % 3); + ternary[static_cast(n * K + k)] = static_cast(static_cast(code) - 1); + const size_t byte_index = static_cast(n * K + k) / 4; + b_packed[byte_index] |= static_cast(code << (2 * (k % 4))); + } + } + + // Powers of two are exact in fp16, so dequantization introduces no error at all. + constexpr float kScaleTable[] = {0.0625f, 0.03125f, 0.015625f}; + std::vector scales(static_cast(N * k_blocks)); + for (size_t i = 0; i < scales.size(); ++i) { + scales[i] = kScaleTable[i % 3]; + } + + // Every 2-bit crumb is 1: 0b01010101. + std::vector zp(static_cast(N * zp_bytes), 0x55); + + // Multiples of 1/64 in [-1, 1) are exact in fp16. The values are zero-mean so the running dot + // product stays small and fp16 accumulation error cannot masquerade as a kernel bug. + uint32_t rng = 0x9E3779B9u; + std::vector a_vals(static_cast(M * K)); + for (size_t i = 0; i < a_vals.size(); ++i) { + rng = rng * 1664525u + 1013904223u; + a_vals[i] = static_cast(static_cast((rng >> 16) & 0x7F) - 64) / 64.0f; + } + + // Dequantize one output channel at a time and reuse it across all M rows: the naive + // m/n/k triple loop costs an integer division per element and is unusably slow at the + // 27B prefill shape (N = 10240, K = 5120). + std::vector expected(static_cast(M * N)); + std::vector w_row(static_cast(K)); + for (int64_t n = 0; n < N; ++n) { + for (int64_t kb = 0; kb < k_blocks; ++kb) { + const float s = scales[static_cast(n * k_blocks + kb)]; + for (int64_t j = 0; j < block_size; ++j) { + const int64_t k = kb * block_size + j; + w_row[static_cast(k)] = static_cast(ternary[static_cast(n * K + k)]) * s; + } + } + for (int64_t m = 0; m < M; ++m) { + const float* a_row = a_vals.data() + static_cast(m * K); + float sum = 0.0f; + for (int64_t k = 0; k < K; ++k) { + sum += a_row[k] * w_row[static_cast(k)]; + } + expected[static_cast(m * N + n)] = sum; + } + } + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", QBits); + test.AddAttribute("accuracy_level", static_cast(0)); + + test.AddInput("A", {M, K}, FloatsToMLFloat16s(a_vals), false); + test.AddInput("B", {N, k_blocks, blob_size}, b_packed, true); + test.AddInput("scales", {N, k_blocks}, FloatsToMLFloat16s(scales), true); + test.AddInput("zero_points", {N, zp_bytes}, zp, true); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("Y", {M, N}, FloatsToMLFloat16s(expected)); + test.SetOutputAbsErr("Y", abs_error); + test.SetOutputRelErr("Y", rel_error); + + std::vector> execution_providers; + execution_providers.emplace_back(DefaultCudaExecutionProvider()); + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.ConfigEps(std::move(execution_providers)); + test.RunWithConfig(); +} + +} // namespace + +TEST(MatMul2BitsCuda, TernaryZeroPointOne) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + for (int64_t block_size : {int64_t{16}, int64_t{32}, int64_t{64}, int64_t{128}}) { + for (int64_t m : {int64_t{1}, int64_t{2}, int64_t{8}, int64_t{96}}) { + RunTernary2BitsCudaTest(m, 32, 512, block_size); + } + } + // The shape the ternary 27B model actually runs at decode time, trimmed to keep the test fast. + RunTernary2BitsCudaTest(1, 512, 5120, 128, 0.2f, 0.02f); +} + +// The qkv_proj of the ternary 27B model: N = 10240, K = 5120, block_size = 128. M >= 32 is +// past the small-M kernels and therefore exercises the dequantize-to-fp16 + cuBLAS fallback; +// M = 32 covers the fallback boundary, while M = 1 covers GEMV at the same shape. +TEST(MatMul2BitsCuda, TernaryPrefillQkvProjShape) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + for (int64_t m : {int64_t{1}, int64_t{32}}) { + RunTernary2BitsCudaTest(m, 10240, 5120, 128, 0.2f, 0.02f); + } +} + +// Regression guard for the silent-corruption bug this support was built on top of: an +// unimplemented weight width used to be accepted by the shared validator, fall through to the +// 4-bit dequantize path, and return garbage. It must fail at session initialization instead. +TEST(MatMul2BitsCuda, UnsupportedBitWidthFailsLoudly) { + if (SkipIfNo2BitCudaDevice()) GTEST_SKIP() << "No CUDA device with the required architecture"; + + constexpr int64_t M = 1, N = 32, K = 256, block_size = 128; + constexpr int64_t k_blocks = K / block_size; + + OpTester test("MatMulNBits", 1, kMSDomain); + test.AddAttribute("K", K); + test.AddAttribute("N", N); + test.AddAttribute("block_size", block_size); + test.AddAttribute("bits", static_cast(3)); + test.AddAttribute("accuracy_level", static_cast(0)); + + test.AddInput("A", {M, K}, FloatsToMLFloat16s(std::vector(M * K, 0.5f)), false); + test.AddInput("B", {N, k_blocks, block_size * 3 / 8}, + std::vector(static_cast(N * k_blocks * block_size * 3 / 8), 0x55), true); + test.AddInput("scales", {N, k_blocks}, + FloatsToMLFloat16s(std::vector(static_cast(N * k_blocks), 0.25f)), true); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("Y", {M, N}, FloatsToMLFloat16s(std::vector(M * N, 0.0f))); + + std::vector> execution_providers; + execution_providers.emplace_back(DefaultCudaExecutionProvider()); + SessionOptions session_options; + session_options.use_per_session_threads = false; + ASSERT_STATUS_OK(session_options.config_options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1")); + test.Config(session_options); + test.Config(OpTester::ExpectResult::kExpectFailure, "bits") + .ConfigEps(std::move(execution_providers)) + .RunWithConfig(); +} + +#endif // USE_CUDA + } // namespace test } // namespace onnxruntime