diff --git a/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md b/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md new file mode 100644 index 0000000000000..046f17f81fa31 --- /dev/null +++ b/docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md @@ -0,0 +1,127 @@ +# MatMulNBits Small-M GEMV Profiling Experiments + +This file records CUDA `MatMulNBits` 4-bit (M = 2..16) and 8-bit (M = 2..5) +profiling results for the small-M range (e.g. multi-row decode or short +prefill) so future kernel and dispatch changes can be compared against a +stable baseline. + +> **Note**: These are **point-in-time** measurements captured on the specific GPU, driver, CUDA +> toolkit, and ORT build noted in the section header. Treat the numbers as a historical baseline +> for regression comparison, not as current performance guidance — re-run the benchmark script on +> your own hardware before drawing conclusions. + +## 2026-06-30 Baseline: A100 (SM80), Warmup 25, Repeat 200 + +### Setup + +- Machine/GPU: A100 (SM80). +- CUDA toolkit: 13.0. ONNX Runtime build: `build/Release`. +- Benchmark script: `onnxruntime/test/python/transformers/profile_matmul_nbits.py`. +- Warmup: 25 ORT runs before the measured `benchmark` NVTX range. Repeat: best of 10 trials x 200 runs. +- Weights: 4-bit, block_size 32, asymmetric (with zero points), per the standard `[N, blocks, blob]` layout. +- Matrices sized after a Qwen3-8B-class dense decoder + lm_head. + +Command template: + +```bash +# Host-timing table across all matrices: +python onnxruntime/test/python/transformers/profile_matmul_nbits.py --warmup 25 --repeat 200 + +# Single case: +python onnxruntime/test/python/transformers/profile_matmul_nbits.py --k 4096 --n 12288 --m 8 + +# Kernel-level via nsys + the repo parser: +nsys profile -t cuda,nvtx -o mnb --export=sqlite \ + python onnxruntime/test/python/transformers/profile_matmul_nbits.py --k 4096 --n 12288 --m 8 +python onnxruntime/test/python/transformers/parse_nsys.py mnb.sqlite --nvtx-range benchmark --pattern '%' +``` + +### Before / After (4-bit) + +`before` is the prior dispatch (M=1 single-row GEMV; M>1 falls back to weight dequantization + cuBLAS +GEMM, which is M-independent and dequantizes the full weight regardless of M). `after` is the batched +small-M GEMV (`MatMulFloat4BatchedKernel`) covering M=2..16. Values are average op latency in microseconds +(lower is better). + +Before (dequant + cuBLAS for M>1): + +| matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | +|---------|-------|--------|-------|--------|--------|--------|--------| +| qkv | 4096 | 4096 | 25.9 | 77.2 | 69.3 | 69.5 | 69.9 | +| o_proj | 4096 | 4096 | 23.2 | 69.2 | 69.1 | 69.3 | 69.7 | +| gate_up | 4096 | 12288 | 39.3 | 172.0 | 172.1 | 172.1 | 172.4 | +| down | 12288 | 4096 | 38.8 | 174.8 | 175.1 | 175.4 | 175.6 | +| lm_head | 4096 | 151936 | 301.6 | 1868.0 | 1871.1 | 1877.8 | 1885.1 | + +After (batched small-M GEMV for M=2..16): + +| matrix | K | N | M=1 | M=2 | M=4 | M=8 | M=16 | +|---------|-------|--------|-------|-------|-------|-------|--------| +| qkv | 4096 | 4096 | 25.4 | 30.1 | 32.3 | 43.2 | 70.0 | +| o_proj | 4096 | 4096 | 22.6 | 26.4 | 28.1 | 36.7 | 58.7 | +| gate_up | 4096 | 12288 | 37.5 | 43.5 | 49.5 | 71.9 | 125.9 | +| down | 12288 | 4096 | 37.9 | 47.4 | 52.9 | 76.8 | 150.1 | +| lm_head | 4096 | 151936 | 300.7 | 329.9 | 424.1 | 635.3 | 1226.9 | + +Speedup (before / after, >1 means the batched GEMV is faster): + +| matrix | M=2 | M=4 | M=8 | M=16 | +|---------|-------|-------|-------|-------| +| qkv | 2.56x | 2.15x | 1.61x | 1.00x | +| o_proj | 2.62x | 2.46x | 1.89x | 1.19x | +| gate_up | 3.95x | 3.48x | 2.39x | 1.37x | +| down | 3.69x | 3.31x | 2.28x | 1.17x | +| lm_head | 5.66x | 4.41x | 2.96x | 1.54x | + +### Before / After (8-bit) + +The 8-bit batched GEMV (`MatMulFloat8bKernelBatched`) covers M=2..5. 8-bit weights are twice the bytes of +4-bit and the GEMV runs on CUDA cores, so it crosses over to the dequantize + cuBLAS (tensor-core) fallback +at a lower M than the 4-bit path; M>=6 keeps the fallback. Values are average op latency in microseconds. + +Before (dequant + cuBLAS for M>1): + +| matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | +|---------|-------|--------|-------|--------|--------|--------|--------| +| qkv | 4096 | 4096 | 36.2 | 80.6 | 72.9 | 72.8 | 73.3 | +| o_proj | 4096 | 4096 | 31.1 | 73.2 | 72.7 | 73.6 | 73.4 | +| gate_up | 4096 | 12288 | 63.5 | 184.4 | 184.2 | 184.1 | 184.3 | +| down | 12288 | 4096 | 67.8 | 187.3 | 187.4 | 187.8 | 188.0 | +| lm_head | 4096 | 151936 | 535.9 | 2025.0 | 2025.9 | 2028.1 | 2029.8 | + +After (batched small-M GEMV for M=2..5): + +| matrix | K | N | M=1 | M=2 | M=3 | M=4 | M=5 | +|---------|-------|--------|-------|-------|-------|--------|--------| +| qkv | 4096 | 4096 | 36.2 | 46.1 | 48.7 | 57.9 | 67.0 | +| o_proj | 4096 | 4096 | 31.6 | 39.5 | 48.6 | 57.9 | 67.1 | +| gate_up | 4096 | 12288 | 63.4 | 80.9 | 104.6 | 128.9 | 152.6 | +| down | 12288 | 4096 | 68.0 | 96.2 | 119.3 | 146.4 | 172.0 | +| lm_head | 4096 | 151936 | 536.3 | 647.0 | 896.9 | 1157.1 | 1420.1 | + +Speedup (before / after, >1 means the batched GEMV is faster): + +| matrix | M=2 | M=3 | M=4 | M=5 | +|---------|-------|-------|-------|-------| +| qkv | 1.75x | 1.50x | 1.26x | 1.09x | +| o_proj | 1.85x | 1.50x | 1.27x | 1.09x | +| gate_up | 2.28x | 1.76x | 1.43x | 1.21x | +| down | 1.95x | 1.57x | 1.28x | 1.09x | +| lm_head | 3.13x | 2.26x | 1.75x | 1.43x | + +### Observations + +- The previous M>1 path dequantizes the entire weight matrix to a temporary buffer and calls cuBLAS, so + its latency is flat across M (e.g. `gate_up` ~172 us, `lm_head` ~1.9 ms even at M=2). The batched small-M + GEMV reads the quantized weight once and scales with M, giving 2.6-5.7x at M=2 (4-bit) / 1.8-3.1x at M=2 + (8-bit). +- 4-bit stays at or above parity through M=16; 8-bit wins through M=5 and falls back to dequant + cuBLAS + for M>=6, where the tensor-core GEMM beats the CUDA-core GEMV on the heavier 8-bit weights. +- M=1 decode is unchanged (same single-row GEMV in both builds). +- No prepacking is used, so there is no extra resident weight memory and no GEMM tactic profiling at + session init. + +### Next Experiments + +- Sweep block_size {16, 32, 64, 128} and bf16 activations. +- Add kernel-level (nsys) breakdowns to separate compute from launch/dispatch overhead. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index b24e445401b99..1554cc2aecbea 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -5,6 +5,7 @@ #include #include #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" @@ -441,6 +442,631 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt } } +// ===== Small-M batched GEMV (short prefill / small batch) ===== +// The single-row MatMulFloatInt4Kernel launches one block per output row (grid.y = m), so each row +// independently re-reads and re-dequantizes all of B; weight traffic and dequant work scale with M. +// For 2 <= M <= cap we instead dequantize each packed weight word once and accumulate it against +// CtaM activation rows held in registers, cutting weight traffic to ceil(M/CtaM)x. This is the same +// design used by TensorRT-LLM weightOnlyBatchedGemv / AWQ / llama.cpp MMVQ for small batch. +// +// Upper bound on M is kSmallMMax for all dtypes (measured on A100 vs the dequantize+cuBLAS fallback). +// half/bf16 run the register-tiled batched kernel, which stays ahead of the tensor-core GEMM through +// M<=16. float has no tensor-core GEMM fallback, so it uses the shared-memory small-M kernel over the +// same range. +constexpr int kSmallMMax = 16; +template +__host__ __device__ constexpr int SmallMCap() { + return kSmallMMax; +} + +// Holds the 8 dequantized weights produced from one packed 4-bit word, in the layout each +// AccumulateRow expects for its dtype. Splitting dequantization from the per-row multiply lets one +// dequantized weight feed all CtaM rows. +template +struct DequantizedEight; + +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +template <> +struct DequantizedEight { + half2 v[4]; // order [04, 15, 26, 37], matching Convert8xInt4To8xHalfs + the prmt below +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, half scale, uint8_t zp, DequantizedEight& d) { + half2 scale_half2 = {scale, scale}; + half zp_adjust = -scale * __short2half_rn(zp); + half2 zp_adjust2 = {zp_adjust, zp_adjust}; + half2 elements[4]; + Convert8xInt4To8xHalfs(values_quant, elements); + d.v[0] = elements[0] * scale_half2 + zp_adjust2; + d.v[1] = elements[1] * scale_half2 + zp_adjust2; + d.v[2] = elements[2] * scale_half2 + zp_adjust2; + d.v[3] = elements[3] * scale_half2 + zp_adjust2; +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const half* a, half* sums) { + uint4 vec_a = *(reinterpret_cast(a)); + constexpr uint32_t kLowHalf2 = 0x5410; + constexpr uint32_t kHighHalf2 = 0x7632; + uint4 vp; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.x) : "r"(vec_a.x), "r"(vec_a.z), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.y) : "r"(vec_a.x), "r"(vec_a.z), "r"(kHighHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.z) : "r"(vec_a.y), "r"(vec_a.w), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.w) : "r"(vec_a.y), "r"(vec_a.w), "r"(kHighHalf2)); + half2* s = reinterpret_cast(sums); + s[0] = s[0] + d.v[0] * (*(reinterpret_cast(&vp.x))); + s[1] = s[1] + d.v[1] * (*(reinterpret_cast(&vp.y))); + s[2] = s[2] + d.v[2] * (*(reinterpret_cast(&vp.z))); + s[3] = s[3] + d.v[3] * (*(reinterpret_cast(&vp.w))); +} +#else +template <> +struct DequantizedEight { + half v[8]; +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, half scale, uint8_t zp, DequantizedEight& d) { + half zp_adjust = -scale * __short2half_rn(zp); +#pragma unroll + for (int i = 0; i < 8; i++) { + d.v[i] = __uint2half_rn((values_quant >> (4 * i)) & 0xF) * scale + zp_adjust; + } +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const half* a, half* sums) { +#pragma unroll + for (int i = 0; i < 8; i++) { + sums[i] += d.v[i] * a[i]; + } +} +#endif + +template <> +struct DequantizedEight { + float v[8]; +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, float scale, uint8_t zp, DequantizedEight& d) { + float zp_adjust = -scale * zp; +#pragma unroll + for (int i = 0; i < 8; i++) { + d.v[i] = float((values_quant >> (4 * i)) & 0xF) * scale + zp_adjust; + } +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const float* a, float* sums) { + float4 a0 = *(reinterpret_cast(a)); + float4 a1 = *(reinterpret_cast(a + 4)); + sums[0] += d.v[0] * a0.x; + sums[1] += d.v[1] * a0.y; + sums[2] += d.v[2] * a0.z; + sums[3] += d.v[3] * a0.w; + sums[4] += d.v[4] * a1.x; + sums[5] += d.v[5] * a1.y; + sums[6] += d.v[6] * a1.z; + sums[7] += d.v[7] * a1.w; +} + +template <> +struct DequantizedEight { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + __nv_bfloat162 v[4]; +#else + nv_bfloat16 v[8]; +#endif +}; +__device__ __forceinline__ void DequantizeEight(uint32_t values_quant, nv_bfloat16 scale, uint8_t zp, DequantizedEight& d) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + __nv_bfloat162 scale_bf162 = __bfloat162bfloat162(scale); + nv_bfloat16 zp_adjust = -scale * __uint2bfloat16_rn(zp); + __nv_bfloat162 zp_adjust2 = __bfloat162bfloat162(zp_adjust); + __nv_bfloat162 elements[4]; + Convert8xInt4To8xBF16s(values_quant, elements); + d.v[0] = __hfma2(elements[0], scale_bf162, zp_adjust2); + d.v[1] = __hfma2(elements[1], scale_bf162, zp_adjust2); + d.v[2] = __hfma2(elements[2], scale_bf162, zp_adjust2); + d.v[3] = __hfma2(elements[3], scale_bf162, zp_adjust2); +#endif +} +__device__ __forceinline__ void AccumulateRow(const DequantizedEight& d, const nv_bfloat16* a, nv_bfloat16* sums) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint4 vec_a = *(reinterpret_cast(a)); + constexpr uint32_t kLowHalf2 = 0x5410; + constexpr uint32_t kHighHalf2 = 0x7632; + uint4 vp; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.x) : "r"(vec_a.x), "r"(vec_a.z), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.y) : "r"(vec_a.x), "r"(vec_a.z), "r"(kHighHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.z) : "r"(vec_a.y), "r"(vec_a.w), "r"(kLowHalf2)); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(vp.w) : "r"(vec_a.y), "r"(vec_a.w), "r"(kHighHalf2)); + __nv_bfloat162* s = reinterpret_cast<__nv_bfloat162*>(sums); + s[0] = __hfma2(d.v[0], *reinterpret_cast<__nv_bfloat162*>(&vp.x), s[0]); + s[1] = __hfma2(d.v[1], *reinterpret_cast<__nv_bfloat162*>(&vp.y), s[1]); + s[2] = __hfma2(d.v[2], *reinterpret_cast<__nv_bfloat162*>(&vp.z), s[2]); + s[3] = __hfma2(d.v[3], *reinterpret_cast<__nv_bfloat162*>(&vp.w), s[3]); +#endif +} + +// ---- Small-M batched GEMV (half/bf16): CtaM x CtaN register tiling -------------------------- +// 2-wide accumulator type (half2 / bf162). +template +struct Acc2; +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +template <> +struct Acc2 { + using type = half2; +}; +#endif +template <> +struct Acc2 { + using type = __nv_bfloat162; +}; + +// Four 2-wide weight lanes in NATURAL element order [01,23,45,67], so a naturally-loaded activation +// (uint4 reinterpreted as four half2) can be multiply-accumulated with no per-activation permute. +template +struct WPack { + typename Acc2::type v[4]; +}; + +// DequantizeEight emits [04,15,26,37] (the order of Convert8xInt4To8xHalfs); repack to natural order +// once per column. Doing the prmt on the (CtaN) weights instead of the (CtaM) activations cuts the +// permute count by CtaM/CtaN, which dominates at small M. +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530) && !defined(__HIPCC__) +__device__ __forceinline__ WPack PackNatural(const DequantizedEight& d) { + uint32_t d0 = *reinterpret_cast(&d.v[0]); + uint32_t d1 = *reinterpret_cast(&d.v[1]); + uint32_t d2 = *reinterpret_cast(&d.v[2]); + uint32_t d3 = *reinterpret_cast(&d.v[3]); + constexpr uint32_t kLo = 0x5410; // (x0,x1) of two half2 -> elements 0,1 + constexpr uint32_t kHi = 0x7632; // (y0,y1) -> elements 4,5 + WPack w; + uint32_t t; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kLo)); + w.v[0] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kLo)); + w.v[1] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kHi)); + w.v[2] = *reinterpret_cast(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kHi)); + w.v[3] = *reinterpret_cast(&t); + return w; +} +__device__ __forceinline__ void DotAccum(const WPack& w, const half2* a4, half2& acc) { + acc = __hfma2(w.v[0], a4[0], acc); + acc = __hfma2(w.v[1], a4[1], acc); + acc = __hfma2(w.v[2], a4[2], acc); + acc = __hfma2(w.v[3], a4[3], acc); +} +__device__ __forceinline__ float HorizontalAdd(half2 acc) { + return static_cast(acc.x) + static_cast(acc.y); +} +#endif + +__device__ __forceinline__ WPack PackNatural(const DequantizedEight& d) { + WPack w; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint32_t d0 = *reinterpret_cast(&d.v[0]); + uint32_t d1 = *reinterpret_cast(&d.v[1]); + uint32_t d2 = *reinterpret_cast(&d.v[2]); + uint32_t d3 = *reinterpret_cast(&d.v[3]); + constexpr uint32_t kLo = 0x5410; + constexpr uint32_t kHi = 0x7632; + uint32_t t; + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kLo)); + w.v[0] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kLo)); + w.v[1] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d0), "r"(d1), "r"(kHi)); + w.v[2] = *reinterpret_cast<__nv_bfloat162*>(&t); + asm volatile("prmt.b32 %0, %1, %2, %3;\n" : "=r"(t) : "r"(d2), "r"(d3), "r"(kHi)); + w.v[3] = *reinterpret_cast<__nv_bfloat162*>(&t); +#endif + return w; +} +__device__ __forceinline__ void DotAccum(const WPack& w, const __nv_bfloat162* a4, __nv_bfloat162& acc) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + acc = __hfma2(w.v[0], a4[0], acc); + acc = __hfma2(w.v[1], a4[1], acc); + acc = __hfma2(w.v[2], a4[2], acc); + acc = __hfma2(w.v[3], a4[3], acc); +#endif +} +__device__ __forceinline__ float HorizontalAdd(__nv_bfloat162 acc) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return static_cast(acc.x) + static_cast(acc.y); +#else + return 0.f; +#endif +} + +// Each warp computes CtaN output columns x CtaM rows. Lanes split K (8 elems/lane/iter); per-row +// activations are loaded once (uint4) and reused across CtaN columns, and the int4->half order permute +// is applied once per column weight (not per row). A single 2-wide accumulator per (row,column) keeps +// CtaM=m up to 16 in registers with the weight streamed exactly once. The launch bound pins >=3 blocks +// per SM so the CtaN=4 tiling (which minimizes activation L2 traffic) keeps enough occupancy to hide +// memory latency. Standard MatMulNBits [N, blocks, blob] layout, no prepacking; scales/zp from global. +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat4BatchedKernel( + 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 AccT = typename Acc2::type; + const int lane_id = threadIdx.x; + const int warp_id = WarpUniform(threadIdx.y); + const int col_base = (blockIdx.x * kColsPerThreadBlock + warp_id) * CtaN; + const int m_base = blockIdx.y * CtaM; + const int valid = m - m_base; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; // 256 + const int zp_blocks = (blocks_per_K + 1) / 2; + + const T* a_base = a_data + static_cast(m_base) * k + (lane_id << 3); + 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 / 2) + lane_id * 4; + } + + AccT acc[CtaM][CtaN]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { +#pragma unroll + for (int c = 0; c < CtaN; c++) { + acc[r][c] = AccT{}; + } + } + + int k_id = 0; + int t_meta_k = lane_id * 8 / block_size; + constexpr int kWork = CtaM * CtaN; + constexpr int kMainUnroll = (kWork >= 20) ? 1 : (kWork >= 12) ? 2 + : 4; + +#define BATCHED_BODY(i) \ + do { \ + WPack w[CtaN]; \ + const int bk = t_meta_k + k_per_iter / block_size * (i); \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + uint32_t value = *(reinterpret_cast(b_ptr[c] + k_per_iter / 2 * (i))); \ + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; \ + uint8_t zp = 8; \ + if constexpr (has_zero_point) { \ + uint8_t zpb = zero_points[static_cast(col_base + c) * zp_blocks + (bk >> 1)]; \ + zp = (bk & 1) ? (zpb >> 4) : (zpb & 0x0f); \ + } \ + DequantizedEight d; \ + DequantizeEight(value, scale, zp, d); \ + w[c] = PackNatural(d); \ + } \ + _Pragma("unroll") for (int r = 0; r < CtaM; r++) { \ + if (r < valid) { \ + AccT a4[4]; \ + *reinterpret_cast(a4) = *reinterpret_cast( \ + a_base + static_cast(r) * k + k_id + (i) * k_per_iter); \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + DotAccum(w[c], a4, acc[r][c]); \ + } \ + } \ + } \ + } while (false) + +#define BATCHED_UNROLL(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ + BATCHED_BODY(i); \ + } \ + _Pragma("unroll") for (int c = 0; c < CtaN; c++) { \ + b_ptr[c] += k_per_iter / 2 * kUnroll; \ + } \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + BATCHED_UNROLL(kMainUnroll); + BATCHED_UNROLL(1); +#undef BATCHED_UNROLL + + if (k_id + lane_id * 8 < k) { + WPack w[CtaN]; + const int bk = t_meta_k; +#pragma unroll + for (int c = 0; c < CtaN; c++) { + uint32_t value = *(reinterpret_cast(b_ptr[c])); + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; + uint8_t zp = 8; + if constexpr (has_zero_point) { + uint8_t zpb = zero_points[static_cast(col_base + c) * zp_blocks + (bk >> 1)]; + zp = (bk & 1) ? (zpb >> 4) : (zpb & 0x0f); + } + DequantizedEight d; + DequantizeEight(value, scale, zp, d); + w[c] = PackNatural(d); + } +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r < valid) { + AccT a4[4]; + *reinterpret_cast(a4) = *reinterpret_cast(a_base + static_cast(r) * k + k_id); +#pragma unroll + for (int c = 0; c < CtaN; c++) { + DotAccum(w[c], a4, acc[r][c]); + } + } + } + } +#undef BATCHED_BODY + +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r >= valid) continue; +#pragma unroll + for (int c = 0; c < CtaN; c++) { + float sum = HorizontalAdd(acc[r][c]); + for (int i = kWarpSize / 2; i > 0; i = i / 2) { + sum += WARP_SHFL_DOWN(sum, i); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + (col_base + c)] = static_cast(sum); + } + } + } +} + +// Batched GEMV: block computes CtaM rows x kColsPerThreadBlock columns. Grid is +// (ceil(N/kColsPerThreadBlock), ceil(M/CtaM)). Mirrors MatMulFloatInt4Kernel's shared-memory scale/zp +// staging and packed-weight indexing; the only change is looping CtaM rows per dequantized word. +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4KernelSmallM( + 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) { + const int n_block_id = blockIdx.x; + const int m_base = blockIdx.y * CtaM; + const int lane_id = threadIdx.x; + const int warp_id = WarpUniform(threadIdx.y); + const int n_id = n_block_id * kColsPerThreadBlock + warp_id; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; + + extern __shared__ char shared_buffer[]; + T* b_scale_vec = (T*)shared_buffer; + int offset = n_block_id * kColsPerThreadBlock * blocks_per_K; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * blocks_per_K; i += kColsPerThreadBlock * kWarpSize) { + b_scale_vec[i] = scales_data[offset + i]; + } + + uint8_t* b_zp_vec; + (void)b_zp_vec; + if constexpr (has_zero_point) { + b_zp_vec = reinterpret_cast(b_scale_vec + kColsPerThreadBlock * blocks_per_K); + const int b_zp_k = (blocks_per_K + 1) / 2; + int zp_offset = n_block_id * kColsPerThreadBlock * b_zp_k; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * b_zp_k; i += kColsPerThreadBlock * kWarpSize) { + b_zp_vec[2 * i] = (zero_points[zp_offset + i] & 0x0f); + b_zp_vec[2 * i + 1] = (zero_points[zp_offset + i] >> 4); + } + b_zp_vec += warp_id * b_zp_k * 2; + } + __syncthreads(); + + const int valid = m - m_base; + const T* a_row[CtaM]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { + a_row[r] = a_data + static_cast(m_base + r) * k + (lane_id << 3); + } + b_scale_vec += warp_id * blocks_per_K; + + T sums[CtaM][8]; +#pragma unroll + for (int r = 0; r < CtaM; r++) { +#pragma unroll + for (int j = 0; j < 8; j++) { + sums[r][j] = static_cast(0); + } + } + + int k_id = 0; + int t_meta_k = lane_id * 8 / block_size; + b_data_quant += n_id * blocks_per_K * (block_size / 2) + lane_id * 4; + +#define SmallMUnRoll(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ + uint32_t value = *(reinterpret_cast(b_data_quant + k_per_iter / 2 * i)); \ + T scale = b_scale_vec[t_meta_k + k_per_iter / block_size * i]; \ + uint8_t zp = 8; \ + if constexpr (has_zero_point) { \ + zp = b_zp_vec[t_meta_k + k_per_iter / block_size * i]; \ + } \ + DequantizedEight d; \ + DequantizeEight(value, scale, zp, d); \ + _Pragma("unroll") for (int r = 0; r < CtaM; r++) { \ + if (r < valid) AccumulateRow(d, a_row[r] + k_id + i * k_per_iter, sums[r]); \ + } \ + } \ + b_data_quant += k_per_iter / 2 * kUnroll; \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + SmallMUnRoll(16); + SmallMUnRoll(4); + SmallMUnRoll(1); +#undef SmallMUnRoll + + if (k_id + lane_id * 8 < k) { + uint32_t value = *(reinterpret_cast(b_data_quant)); + T scale = b_scale_vec[t_meta_k]; + uint8_t zp = 8; + if constexpr (has_zero_point) { + zp = b_zp_vec[t_meta_k]; + } + DequantizedEight d; + DequantizeEight(value, scale, zp, d); +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r < valid) AccumulateRow(d, a_row[r] + k_id, sums[r]); + } + } + +#pragma unroll + for (int r = 0; r < CtaM; r++) { + if (r >= valid) continue; + float sum = static_cast(sums[r][0] + sums[r][1] + sums[r][2] + sums[r][3] + + sums[r][4] + sums[r][5] + sums[r][6] + sums[r][7]); + for (int i = kWarpSize / 2; i > 0; i = i / 2) { + sum += WARP_SHFL_DOWN(sum, i); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + n_id] = sum; + } + } +} + +// Launches the batched small-M kernel for 2 <= m <= SmallMCap(). Returns false if m is out of range +// so the caller can fall back to the single-row (m==1) or dequant+GEMM (large m) paths. +template +bool TryMatMulSmallM4Bits( + 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, + size_t shared_mem_size, + cudaStream_t stream) { + if (m < 2 || m > SmallMCap()) { + return false; + } + const int cta_m = (m <= 2) ? 2 : 4; + dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, (m + cta_m - 1) / cta_m); + +#define SmallMDispatch(BS, CM) \ + if (nullptr != zero_points) { \ + MatMulFloatInt4KernelSmallM<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } else { \ + MatMulFloatInt4KernelSmallM<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } +#define SmallMDispatchBlock(CM) \ + if (16 == block_size) { \ + SmallMDispatch(16, CM) \ + } else if (32 == block_size) { \ + SmallMDispatch(32, CM) \ + } else if (64 == block_size) { \ + SmallMDispatch(64, CM) \ + } else if (128 == block_size) { \ + SmallMDispatch(128, CM) \ + } else { \ + return false; \ + } + + if (cta_m == 2) { + SmallMDispatchBlock(2) + } else { + SmallMDispatchBlock(4) + } + +#undef SmallMDispatchBlock +#undef SmallMDispatch + return true; +} + +// Small-M launcher (half/bf16): picks CtaM >= m from {2,4,8,16} so a single block streams the +// weight once, and CtaN columns/warp (largest of {4,2,1} dividing N/kCols) to reuse each activation +// load across columns. Returns false for float or out-of-range m so the caller falls back. +template +bool TryMatMulBatched4Bits( + 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, + size_t /*shared_mem_size*/, + cudaStream_t stream) { + if constexpr (std::is_same::value) { + return false; + } else { + if (m < 2 || m > kSmallMMax) { + return false; + } + // CtaM = smallest of {2,4,8,16} >= m, streaming the weight once per block row. Non-power-of-2 CtaM + // (10/12/14) compile to materially slower code, so the row-tile is rounded up (M>8 uses CtaM=16). + // CtaN = 2 columns/warp where N allows (halves activation L2 traffic); CtaN=4 lost to register + // pressure so it is not used by default. + const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 + : (m <= 8) ? 8 + : 16; + const int cta_n = (n % (kColsPerThreadBlock * 2) == 0) ? 2 : 1; + dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + dim3 blocks(n / (kColsPerThreadBlock * cta_n), (m + cta_m - 1) / cta_m); + +#define BatchedDispatch(BS, CM, CN) \ + if (nullptr != zero_points) { \ + MatMulFloat4BatchedKernel<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } else { \ + MatMulFloat4BatchedKernel<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, (k + BS - 1) / BS); \ + } +#define BatchedDispatchN(CM, CN) \ + if (16 == block_size) { \ + BatchedDispatch(16, CM, CN) \ + } else if (32 == block_size) { \ + BatchedDispatch(32, CM, CN) \ + } else if (64 == block_size) { \ + BatchedDispatch(64, CM, CN) \ + } else if (128 == block_size) { \ + BatchedDispatch(128, CM, CN) \ + } else { \ + return false; \ + } +#define BatchedDispatchM(CN) \ + switch (cta_m) { \ + case 2: \ + BatchedDispatchN(2, CN) break; \ + case 4: \ + BatchedDispatchN(4, CN) break; \ + case 8: \ + BatchedDispatchN(8, CN) break; \ + default: \ + BatchedDispatchN(16, CN) break; \ + } + + if (cta_n == 2) { + BatchedDispatchM(2) + } else { + BatchedDispatchM(1) + } + +#undef BatchedDispatchM +#undef BatchedDispatchN +#undef BatchedDispatch + return true; + } +} + template bool TryMatMul4Bits( T* output, @@ -455,7 +1081,7 @@ bool TryMatMul4Bits( int block_size, size_t shared_mem_per_block, cudaStream_t stream) { - if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > 1) { + if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > SmallMCap()) { return false; } @@ -477,6 +1103,15 @@ bool TryMatMul4Bits( return false; } + // The register-tiled batched path (half/bf16, 2 <= m <= cap) launches with no shared memory, so try + // it before the shared-memory budget gate that only constrains the shared-memory kernels below. + if (m >= 2) { + if (TryMatMulBatched4Bits(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, 0, stream)) { + return true; + } + } + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, m); dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); int blocks_per_K = (k + block_size - 1) / block_size; @@ -486,6 +1121,14 @@ bool TryMatMul4Bits( return false; } + // Float, and any half/bf16 shape the batched path rejected, falls back to the shared-memory small-M + // kernel. m == 1 uses the single-row kernel below; larger m used the dequantize + cuBLAS path + // (returned false above). + if (m >= 2) { + return TryMatMulSmallM4Bits(output, a_data, b_data_quant, scales_data, zero_points, + m, n, k, block_size, shared_mem_size, stream); + } + #define MatMulFloatInt4KernelDispatch(block_size) \ if (nullptr != zero_points) { \ MatMulFloatInt4Kernel<<>>( \ diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu index 2788675765b62..495922185c939 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu @@ -175,6 +175,185 @@ __device__ __forceinline__ void AccumulateEightElements8b( #endif } +// ===== Small-M batched GEMV (8-bit) ===== +// Same idea as the 4-bit batched kernel: each warp computes CtaN columns x CtaM rows. Lanes split K +// (8 elems/lane/iter via one uint64 weight load); each column's 8 weights are dequantized once per row +// group, per-row activations are loaded once and reused across CtaN columns, and a single float +// accumulator per (row, column) keeps register pressure low while the weight is streamed once. The +// launch bound pins >=3 blocks per SM. Standard [N, blocks, blob] layout, no prepacking; the 8-bit path +// accumulates in float, so this stays dtype-agnostic (float/half/bf16). +constexpr int kSmallMMax8 = 5; +template +__host__ __device__ constexpr int SmallMCap8() { + return kSmallMMax8; +} + +template +__device__ __forceinline__ float ToFloat8b(T v); +template <> +__device__ __forceinline__ float ToFloat8b(float v) { return v; } +template <> +__device__ __forceinline__ float ToFloat8b(half v) { return __half2float(v); } +template <> +__device__ __forceinline__ float ToFloat8b(nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ void DequantizeEight8b(uint64_t values_quant, T scale, uint8_t zp, float dq[8]) { + float scale_f = ToFloat8b(scale); + float zp_f = static_cast(zp); +#pragma unroll + for (int i = 0; i < 8; ++i) { + uint8_t q = (values_quant >> (i * 8)) & 0xFF; + dq[i] = (static_cast(q) - zp_f) * scale_f; + } +} + +// Load 8 consecutive activations (one lane's iteration tile) into float[8]. +template +__device__ __forceinline__ void LoadEightActivations8b(const T* a, float out[8]) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + out[j] = ToFloat8b(a[j]); + } +} + +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock, 3) MatMulFloat8bKernelBatched( + 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) { + const int lane_id = threadIdx.x; + const int warp_id = threadIdx.y; + const int col_base = (blockIdx.x * kColsPerThreadBlock + warp_id) * CtaN; + const int m_base = blockIdx.y * CtaM; + const int valid = m - m_base; + constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; // 256 + const int lane_offset = lane_id * kElementsPerThreadPerIteration; // lane_id * 8 + + 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 + lane_offset; + } + + float acc[CtaM][CtaN]; +#pragma unroll + for (int r = 0; r < CtaM; ++r) { +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + acc[r][c] = 0.0f; + } + } + + int k_id = 0; + int t_meta_k = lane_offset / block_size; + constexpr int kWork = CtaM * CtaN; + constexpr int kMainUnroll = (kWork >= 20) ? 1 : (kWork >= 12) ? 2 + : 4; + +#define BATCHED8_BODY(i) \ + do { \ + float dq[CtaN][8]; \ + const int bk = t_meta_k + k_per_iter / block_size * (i); \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + uint64_t value = *reinterpret_cast(b_ptr[c] + k_per_iter * (i)); \ + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; \ + uint8_t zp = kDefaultZeroPoint; \ + if constexpr (has_zero_point) { \ + zp = zero_points[static_cast(col_base + c) * blocks_per_K + bk]; \ + } \ + DequantizeEight8b(value, scale, zp, dq[c]); \ + } \ + _Pragma("unroll") for (int r = 0; r < CtaM; ++r) { \ + if (r < valid) { \ + float av[8]; \ + LoadEightActivations8b(a_base + static_cast(r) * k + k_id + (i) * k_per_iter, av); \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + float s = 0.0f; \ + _Pragma("unroll") for (int j = 0; j < 8; ++j) { \ + s = fmaf(av[j], dq[c][j], s); \ + } \ + acc[r][c] += s; \ + } \ + } \ + } \ + } while (false) + +#define BATCHED8_UNROLL(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * k_per_iter; \ + const int k_unroll_bound = k - k % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; ++i) { \ + BATCHED8_BODY(i); \ + } \ + _Pragma("unroll") for (int c = 0; c < CtaN; ++c) { \ + b_ptr[c] += kUnrollStep; \ + } \ + t_meta_k += k_per_iter / block_size * kUnroll; \ + } \ + } while (false) + + BATCHED8_UNROLL(kMainUnroll); + BATCHED8_UNROLL(1); +#undef BATCHED8_UNROLL + + if (lane_offset + k_id < k) { + float dq[CtaN][8]; + const int bk = t_meta_k; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + uint64_t value = *reinterpret_cast(b_ptr[c]); + T scale = scales_data[static_cast(col_base + c) * blocks_per_K + bk]; + uint8_t zp = kDefaultZeroPoint; + if constexpr (has_zero_point) { + zp = zero_points[static_cast(col_base + c) * blocks_per_K + bk]; + } + DequantizeEight8b(value, scale, zp, dq[c]); + } +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r < valid) { + float av[8]; + LoadEightActivations8b(a_base + static_cast(r) * k + k_id, av); +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + float s = 0.0f; +#pragma unroll + for (int j = 0; j < 8; ++j) { + s = fmaf(av[j], dq[c][j], s); + } + acc[r][c] += s; + } + } + } + } +#undef BATCHED8_BODY + +#pragma unroll + for (int r = 0; r < CtaM; ++r) { + if (r >= valid) continue; +#pragma unroll + for (int c = 0; c < CtaN; ++c) { + float sum = acc[r][c]; + for (int off = kWarpSize / 2; off > 0; off /= 2) { + sum += WARP_SHFL_DOWN(sum, off); + } + if (lane_id == 0) { + output[static_cast(m_base + r) * n + (col_base + c)] = static_cast(sum); + } + } + } +} + // --- CUDA Kernel: MatMulFloat8bKernel (Optimized for m=1) --- // Computes C(1, N) = A(1, K) x B(K, N) // B(K, N) is quantized with 8 bits and block_size bs, stored as [N, K/bs, bs] @@ -346,17 +525,17 @@ bool TryMatMul8Bits( const uint8_t* b_data_quant, // Input B Quantized [N, K/bs, bs] const T* scales_data, // Scales [N, K/bs] const uint8_t* zero_points, // Zero Points [N, K/bs] (can be nullptr) - int m, // Rows of A and C (MUST be 1) + int m, // Rows of A and C (1 single-row, 2..cap batched) int n, // Columns of B and C int k, // Columns of A / Rows of B int block_size, // Quantization block size for B size_t shared_mem_per_block, // Available shared memory per block cudaStream_t stream) { // Constraints Check - // m must be 1 (since this kernel is optimized for m=1) // N must be a multiple of kColsPerThreadBlock (8) for warps to align with columns. // K must be a multiple of kElementsPerThreadPerIteration (8) for full uint64_t reads/processing. - if (m != 1 || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { + // m up to the small-M cap is handled (m==1 single-row, 2..cap batched); larger m falls back. + if (m < 1 || m > SmallMCap8() || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { return false; } @@ -380,7 +559,61 @@ bool TryMatMul8Bits( // Calculate K / block_size (no rounding needed due to k % block_size == 0 check) int blocks_per_K = k / block_size; - // --- Shared Memory Calculation --- + // 2 <= m <= cap: batched GEMV. CtaM = smallest of {2,4,8} >= m streams the weight once per block row + // (m=5 uses CtaM=8 and skips the unused rows); CtaN = 2 columns/warp where N allows (reuses each + // activation load across columns). One float accumulator per (row, column) keeps register pressure + // low. 8-bit weights are twice the bytes of 4-bit and the GEMV runs on CUDA cores, so it crosses over + // to the dequantize + cuBLAS (tensor-core) fallback at a lower M than the 4-bit path; the cap keeps + // only the row counts where it is faster on every matrix shape. This path launches with no shared + // memory, so it runs before the shared-memory budget gate that only constrains the m==1 kernel below. + if (m >= 2) { + const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 + : 8; + const int cta_n = (n % (kColsPerThreadBlock * 2) == 0) ? 2 : 1; + dim3 batched_blocks(n / (kColsPerThreadBlock * cta_n), (m + cta_m - 1) / cta_m); +#define MatMulFloat8bBatchedDispatch(bs, cm, cn) \ + if (nullptr != zero_points) { \ + MatMulFloat8bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, zero_points, m, n, k, blocks_per_K); \ + } else { \ + MatMulFloat8bKernelBatched<<>>( \ + output, a_data, b_data_quant, scales_data, nullptr, m, n, k, blocks_per_K); \ + } +#define MatMulFloat8bBatchedDispatchN(cm, cn) \ + if (16 == block_size) { \ + MatMulFloat8bBatchedDispatch(16, cm, cn) \ + } else if (32 == block_size) { \ + MatMulFloat8bBatchedDispatch(32, cm, cn) \ + } else if (64 == block_size) { \ + MatMulFloat8bBatchedDispatch(64, cm, cn) \ + } else if (128 == block_size) { \ + MatMulFloat8bBatchedDispatch(128, cm, cn) \ + } else if (256 == block_size) { \ + MatMulFloat8bBatchedDispatch(256, cm, cn) \ + } else { \ + return false; \ + } +#define MatMulFloat8bBatchedDispatchM(cn) \ + switch (cta_m) { \ + case 2: \ + MatMulFloat8bBatchedDispatchN(2, cn) break; \ + case 4: \ + MatMulFloat8bBatchedDispatchN(4, cn) break; \ + default: \ + MatMulFloat8bBatchedDispatchN(8, cn) break; \ + } + if (cta_n == 2) { + MatMulFloat8bBatchedDispatchM(2) + } else { + MatMulFloat8bBatchedDispatchM(1) + } +#undef MatMulFloat8bBatchedDispatchM +#undef MatMulFloat8bBatchedDispatchN +#undef MatMulFloat8bBatchedDispatch + return true; + } + + // --- Shared Memory Calculation (m == 1) --- // Memory for scales + optional zero points for the columns handled by the block size_t scale_zp_shared_mem = (sizeof(T) + (zero_points != nullptr ? sizeof(uint8_t) : 0)) * static_cast(blocks_per_K) * kColsPerThreadBlock; @@ -395,7 +628,7 @@ bool TryMatMul8Bits( return false; } - // --- Kernel Launch --- + // --- Kernel Launch (m == 1) --- // Macro to simplify dispatching based on block size and presence of zero_points #define MatMulFloat8bKernelM1Dispatch(bs) \ if (nullptr != zero_points) { \ diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index bedf035d320f8..4d5a1c5c17c0b 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -869,6 +869,47 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint) { } } +// Exercises the CUDA small-M batched GEMV tiles: CtaM in {2,4,8,16} (with M values that are not a +// multiple of CtaM so the row-skip path runs) and CtaN in {1,2} (N divisible / not divisible by 16). +TEST(MatMulNBits, Fp16_Int4_SmallMBatchedTiles) { + constexpr float abs_error = 0.1f; + constexpr bool zp_is_4bit = true; + for (auto block_size : {32, 128}) { + for (auto m : {3, 4, 5, 8, 12, 16}) { + for (auto has_zeropoint : {false, true}) { + RunTest(m, 256, 1024, block_size, has_zeropoint, zp_is_4bit, abs_error); // N % 16 == 0 -> CtaN=2 + RunTest(m, 24, 1024, block_size, has_zeropoint, zp_is_4bit, abs_error); // N % 16 != 0 -> CtaN=1 + } + } + } +} + +TEST(MatMulNBits, BFloat16_Int4_SmallMBatchedTiles) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "Skipping BFloat16 tests on CUDA < 8.0"; + } + + constexpr float abs_error = 0.1f; + for (auto block_size : {32, 128}) { + for (auto m : {3, 4, 5, 8, 12, 16}) { + for (auto n : {256, 24}) { // N=256 -> CtaN=2, N=24 -> CtaN=1 + for (auto has_zeropoint : {false, true}) { + TestOptions opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_4bit = true; + opts.output_abs_error = abs_error; + opts.output_rel_error = 0.02f; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } + } + } + } +} + TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { constexpr float abs_error = 0.1f; diff --git a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc index 411e83536c190..c3cd8b388dd8e 100644 --- a/onnxruntime/test/contrib_ops/matmul_8bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_8bits_test.cc @@ -682,6 +682,54 @@ TEST(MatMulNBits, BFloat16_Int8_Chunked_BFloat16ZeroPoint) { RunTest8Bits(opts); } } + +// Exercises the CUDA small-M batched GEMV tiles for 8-bit: CtaM in {2,4,8} (M=3,5 hit the row-skip +// path) and CtaN in {1,2} (N divisible / not divisible by 16). 8-bit caps the batched path at M=5. +TEST(MatMulNBits, Fp16_Int8_SmallMBatchedTiles) { + constexpr float abs_error = 0.1f; + constexpr float rel_error = 0.02f; + for (auto block_size : {32, 128}) { + for (auto m : {2, 3, 4, 5}) { + for (auto n : {256, 24}) { // N=256 -> CtaN=2, N=24 -> CtaN=1 + for (auto has_zeropoint : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_typed = false; + opts.output_abs_error = abs_error; + opts.output_rel_error = rel_error; + RunTest8Bits(opts); + } + } + } + } +} + +TEST(MatMulNBits, BFloat16_Int8_SmallMBatchedTiles) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "Skipping BFloat16 tests on CUDA < 8.0"; + } + + constexpr float abs_error = 0.1f; + constexpr float rel_error = 0.02f; + for (auto block_size : {32, 128}) { + for (auto m : {2, 3, 4, 5}) { + for (auto n : {256, 24}) { + for (auto has_zeropoint : {false, true}) { + TestOptions8Bits opts{}; + opts.M = m, opts.N = n, opts.K = 1024; + opts.block_size = block_size; + opts.has_zero_point = has_zeropoint; + opts.zp_is_typed = false; + opts.output_abs_error = abs_error; + opts.output_rel_error = rel_error; + RunTest8Bits(opts); + } + } + } + } +} #endif #if !defined(USE_CUDA) && !defined(USE_WEBGPU) diff --git a/onnxruntime/test/python/transformers/profile_matmul_nbits.py b/onnxruntime/test/python/transformers/profile_matmul_nbits.py new file mode 100644 index 0000000000000..75edc892ede1f --- /dev/null +++ b/onnxruntime/test/python/transformers/profile_matmul_nbits.py @@ -0,0 +1,206 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +""" +Profiling script for the CUDA MatMulNBits (4-bit / 8-bit weight-only) op. + +Times the op across representative decoder weight matrices and row counts M +(M=1 decode and M>1 small-batch). Mirrors profile_qmoe_gemv.py: +warmup + measured runs wrapped in NVTX ranges so the results can be parsed +with parse_nsys.py for kernel-level timing, while also printing host-observed +average latency for a quick end-to-end view. + +Usage: + # Host-timing table across all cases: + python profile_matmul_nbits.py --warmup 25 --repeat 200 + + # Single case: + python profile_matmul_nbits.py --k 4096 --n 4096 --m 8 --block-size 32 --bits 4 --dtype fp16 + + # Kernel-level via nsys + the repo parser: + nsys profile -t cuda,nvtx -o mnb --export=sqlite \ + python profile_matmul_nbits.py --k 4096 --n 4096 --m 8 + python parse_nsys.py mnb.sqlite --nvtx-range benchmark --pattern '%' +""" + +import argparse +import json +import time +from contextlib import nullcontext + +import numpy as np +import torch +from onnx import TensorProto, helper + +import onnxruntime + +try: + import ml_dtypes + + bfloat16 = ml_dtypes.bfloat16 +except ImportError: + bfloat16 = None + +try: + import nvtx + + has_nvtx = True +except ImportError: + has_nvtx = False + nvtx = None + +RESULT_PREFIX = "MATMUL_NBITS_RESULT " + +_OT = {"fp16": TensorProto.FLOAT16, "bf16": TensorProto.BFLOAT16} +_TT = {"fp16": torch.float16, "bf16": torch.bfloat16} +_ELEM = {torch.float16: TensorProto.FLOAT16, torch.bfloat16: TensorProto.BFLOAT16} + +# Representative decoder weight matrices (K = in features, N = out features), +# sized after a Qwen3-8B-class dense decoder + lm_head. +DEFAULT_CASES = [ + ("qkv", 4096, 4096), + ("o_proj", 4096, 4096), + ("gate_up", 4096, 12288), + ("down", 12288, 4096), + ("lm_head", 4096, 151936), +] +DEFAULT_MS = [1, 2, 4, 8, 16] + + +def _nvtx_range(name, color="green"): + if not has_nvtx: + return nullcontext() + return nvtx.annotate(name, color=color) + + +def build_model(k, n, block_size, bits, onnx_dtype, with_zero_point=True): + rng = np.random.default_rng(0) + n_blocks = (k + block_size - 1) // block_size + blob = block_size // (8 // bits) + b = rng.integers(0, 256, size=(n, n_blocks, blob), dtype=np.uint8) + scales_f32 = rng.random(n * n_blocks).astype(np.float32) * 0.02 + 0.01 + if onnx_dtype == TensorProto.FLOAT16: + scales = scales_f32.astype(np.float16) + elif onnx_dtype == TensorProto.BFLOAT16: + if bfloat16 is None: + raise RuntimeError("ml_dtypes is required for bf16 (pip install ml_dtypes)") + scales = scales_f32.astype(bfloat16) + else: + scales = scales_f32 + inits = [ + helper.make_tensor("B", TensorProto.UINT8, list(b.shape), b.tobytes(), raw=True), + helper.make_tensor("scales", onnx_dtype, list(scales.shape), scales.tobytes(), raw=True), + ] + inputs = ["A", "B", "scales"] + if with_zero_point: + # 4-bit zero points are packed two per byte; 8-bit zero points are one byte per block. + zp_count = n * ((n_blocks + 1) // 2) if bits == 4 else n * n_blocks + zp = rng.integers(0, 256, size=(zp_count,), dtype=np.uint8) + inits.append(helper.make_tensor("zero_points", TensorProto.UINT8, list(zp.shape), zp.tobytes(), raw=True)) + inputs.append("zero_points") + node = helper.make_node( + "MatMulNBits", + inputs, + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=bits, + block_size=block_size, + accuracy_level=0, + ) + graph = helper.make_graph( + [node], + "mnb", + [helper.make_tensor_value_info("A", onnx_dtype, ["M", k])], + [helper.make_tensor_value_info("Y", onnx_dtype, ["M", n])], + inits, + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("com.microsoft", 1), helper.make_opsetid("", 17)] + ) + return model.SerializeToString() + + +def make_session(model): + so = onnxruntime.SessionOptions() + so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + so.log_severity_level = 3 + return onnxruntime.InferenceSession(model, so, providers=["CUDAExecutionProvider"]) + + +def run_case(name, m, k, n, block_size, bits, dtype, warmup, repeat): + onnx_dtype = _OT[dtype] + torch_dtype = _TT[dtype] + sess = make_session(build_model(k, n, block_size, bits, onnx_dtype)) + a = np.random.default_rng(m).random((m, k)).astype(np.float32) * 0.02 - 0.01 + at = torch.from_numpy(a).to(torch_dtype).cuda().contiguous() + y = torch.empty((m, n), dtype=torch_dtype, device="cuda") + io = sess.io_binding() + io.bind_input("A", "cuda", 0, _ELEM[torch_dtype], list(at.shape), at.data_ptr()) + io.bind_output("Y", "cuda", 0, _ELEM[torch_dtype], list(y.shape), y.data_ptr()) + + with _nvtx_range("warmup", "yellow"): + for _ in range(warmup): + sess.run_with_iobinding(io) + torch.cuda.synchronize() + + best = float("inf") + trials = 10 + with _nvtx_range("benchmark", "green"): + for _ in range(trials): + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(repeat): + sess.run_with_iobinding(io) + torch.cuda.synchronize() + best = min(best, (time.perf_counter() - t0) / repeat) + us = best * 1e6 + result = { + "case": name, + "m": m, + "k": k, + "n": n, + "block_size": block_size, + "bits": bits, + "dtype": dtype, + "avg_us": round(us, 2), + } + print(RESULT_PREFIX + json.dumps(result)) + return result + + +def main(): + p = argparse.ArgumentParser(description="Profile CUDA MatMulNBits") + p.add_argument("--k", type=int, help="in features (single-case mode)") + p.add_argument("--n", type=int, help="out features (single-case mode)") + p.add_argument("--m", type=int, help="rows (single-case mode)") + p.add_argument("--block-size", type=int, default=32) + p.add_argument("--bits", type=int, default=4) + p.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"]) + p.add_argument("--warmup", type=int, default=25) + p.add_argument("--repeat", type=int, default=200) + p.add_argument("--ms", type=int, nargs="+", default=DEFAULT_MS, help="row counts to sweep in table mode") + args = p.parse_args() + + if args.k and args.n and args.m: + run_case("custom", args.m, args.k, args.n, args.block_size, args.bits, args.dtype, args.warmup, args.repeat) + return + + rows = {} + for name, k, n in DEFAULT_CASES: + rows[name] = { + m: run_case(name, m, k, n, args.block_size, args.bits, args.dtype, args.warmup, args.repeat)["avg_us"] + for m in args.ms + } + + print(f"\n MatMulNBits {args.bits}-bit block{args.block_size} {args.dtype} (avg us, lower is better)") + print(" " + f"{'matrix':10s} {'K':>6} {'N':>7} " + " ".join(f"M={m:<5}" for m in args.ms)) + for name, k, n in DEFAULT_CASES: + print(" " + f"{name:10s} {k:>6} {n:>7} " + " ".join(f"{rows[name][m]:7.1f}" for m in args.ms)) + + +if __name__ == "__main__": + main()