diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 3efa6ae50faa7..d685831b891f5 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -518,6 +518,15 @@ static const char* const kOrtSessionOptionsQDQMatMulNBitsAccuracyLevel = "sessio // "-1": heuristic - largest power-of-2 <= min(K, 256) that minimizes padding. static const char* const kOrtSessionOptionsQDQMatMulNBitsBlockSize = "session.qdq_matmulnbits_block_size"; +// Enables online tuning of the CUDA MatMulNBits small-M batched GEMV cap (weight-only int4/int8). +// On the first call whose row count M falls in the batched-candidate range for a given weight shape, +// the batched GEMV is micro-benchmarked against the dequantize + cuBLAS fallback and the crossover M is +// cached, replacing the built-in per-device default with a measured one. The result is cached per shape. +// Option values: +// - "0": (default) use the built-in cap. +// - "1": enable online tuning. +static const char* const kOrtSessionOptionsConfigMatMulNBitsTuneSmallM = "ep.cuda.matmulnbits_tune_small_m"; + // Enable the DQ->MatMulNBits fusion graph transformer. // "0": disabled (default). "1": enabled. // This is typically set automatically by InferenceSession when the NvTensorRTRTX EP is registered. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index 1554cc2aecbea..a1155708cf034 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -449,7 +449,7 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt // 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). +// Upper bound on M is kSmallMMax for all dtypes (measured 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. @@ -1080,7 +1080,8 @@ bool TryMatMul4Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream) { + cudaStream_t stream, + int max_batched_m) { if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > SmallMCap()) { return false; } @@ -1105,13 +1106,21 @@ bool TryMatMul4Bits( // 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) { + // max_batched_m lets an online tuner cap the batched range below the compile-time cap for a given + // device/shape; above it, fall through to dequant + cuBLAS. + if (m >= 2 && m <= max_batched_m) { if (TryMatMulBatched4Bits(output, a_data, b_data_quant, scales_data, zero_points, m, n, k, block_size, 0, stream)) { return true; } } + // When the tuner selected the dequant path for this m, skip the shared-memory small-M fallback so the + // caller uses dequant + cuBLAS. (For the default cap this is unreachable since max_batched_m is INT_MAX.) + if (m >= 2 && m > max_batched_m) { + return false; + } + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, m); dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); int blocks_per_K = (k + block_size - 1) / block_size; @@ -1167,7 +1176,8 @@ template bool TryMatMul4Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); template bool TryMatMul4Bits( half* output, @@ -1181,7 +1191,8 @@ template bool TryMatMul4Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); template bool TryMatMul4Bits( nv_bfloat16* output, @@ -1195,7 +1206,8 @@ template bool TryMatMul4Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); template __global__ void MatMulNBitsBiasAddKernel(T* output, const T* bias_data, int n, int64_t total) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu index 495922185c939..39d8111ba44f4 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu @@ -5,6 +5,7 @@ #include #include #include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/quantization/matmul_nbits.cuh" @@ -182,11 +183,12 @@ __device__ __forceinline__ void AccumulateEightElements8b( // 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). +// +// kSmallMMax8 is the default (measured) cap used when online tuning is off. kBatchedStructuralMax8 +// is the largest M the kernel can actually run (CtaM in {2,4,8}); an online tuner may pick any cap up to +// this so a different device can use the batched path past the compile-time default. constexpr int kSmallMMax8 = 5; -template -__host__ __device__ constexpr int SmallMCap8() { - return kSmallMMax8; -} +constexpr int kBatchedStructuralMax8 = 8; template __device__ __forceinline__ float ToFloat8b(T v); @@ -530,15 +532,23 @@ bool TryMatMul8Bits( 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) { + cudaStream_t stream, + int max_batched_m) { // Constraints Check // 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. - // 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) { + // Reject m the batched kernel cannot run (> structural max); m==1 uses the single-row kernel below, + // and 2..cap use the batched kernel (cap resolved below). + if (m < 1 || m > kBatchedStructuralMax8 || n % kColsPerThreadBlock != 0 || k % kElementsPerThreadPerIteration != 0) { return false; } + // Effective batched cap: the compile-time default unless an online tuner overrides it (clamped to + // what the kernel structurally supports). INT_MAX is the sentinel for "no tuning override". + const int batched_cap = (max_batched_m == (std::numeric_limits::max)()) + ? kSmallMMax8 + : ((max_batched_m < kBatchedStructuralMax8) ? max_batched_m : kBatchedStructuralMax8); + // Ensure k_per_iter (kWarpSize * kElementsPerThreadPerIteration) is multiple of block_size. constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; if (k_per_iter % block_size != 0) { @@ -566,7 +576,9 @@ bool TryMatMul8Bits( // 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) { + // max_batched_m lets an online tuner cap the batched range below the compile-time cap for a given + // device/shape; above it, fall through to dequant + cuBLAS. + if (m >= 2 && m <= batched_cap) { const int cta_m = (m <= 2) ? 2 : (m <= 4) ? 4 : 8; const int cta_n = (n % (kColsPerThreadBlock * 2) == 0) ? 2 : 1; @@ -613,6 +625,11 @@ bool TryMatMul8Bits( return true; } + // Any remaining m >= 2 (above max_batched_m) uses dequant + cuBLAS; never the m==1 kernel below. + if (m >= 2) { + return false; + } + // --- 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)) * @@ -676,7 +693,8 @@ template bool TryMatMul8Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); template bool TryMatMul8Bits( half* output, @@ -689,7 +707,8 @@ template bool TryMatMul8Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); // Add template instantiation for nv_bfloat16 template bool TryMatMul8Bits( @@ -703,7 +722,8 @@ template bool TryMatMul8Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 896291ce01b5d..ed13c17a4a508 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -4,6 +4,12 @@ #include "contrib_ops/cuda/quantization/matmul_nbits.h" #include +#include +#include +#include +#include +#include +#include #include "core/common/status.h" #include "core/common/float16.h" @@ -252,6 +258,197 @@ Status MatMulNBits::PrePack_ZeroPoint([[maybe_unused]] const Tensor& tensor, } #endif +namespace { +// Process-global cache of tuned batched-GEMV caps, shared across all MatMulNBits nodes so that the +// handful of distinct weight shapes in a model are profiled once rather than per node. +std::mutex g_small_m_cap_mutex; +std::unordered_map g_small_m_cap_cache; +// Keys currently being profiled by some thread, so concurrent first calls for the same shape do not +// each run the (one-time) micro-benchmark. +std::unordered_set g_small_m_cap_in_progress; + +uint64_t MakeSmallMCapKey(int dtype_tag, int64_t k, int64_t n, int64_t block_size, int64_t bits, + int zp_kind, int sm, int sm_count, int64_t l2_cache_size) { + uint64_t h = 1469598103934665603ull; + auto mix = [&](uint64_t v) { + h ^= v; + h *= 1099511628211ull; + }; + mix(static_cast(dtype_tag)); + mix(static_cast(k)); + mix(static_cast(n)); + mix(static_cast(block_size)); + mix(static_cast(bits)); + mix(static_cast(zp_kind)); + // Device identity: compute capability alone does not distinguish two SM-equal GPUs with different + // SM counts / L2 (e.g. A100 vs A30), whose crossover differs, so mix those in too. + mix(static_cast(sm)); + mix(static_cast(sm_count)); + mix(static_cast(l2_cache_size)); + return h; +} +} // namespace + +template +int MatMulNBits::ResolveSmallMBatchedCap(OpKernelContext* ctx, + const uint8_t* blob_data, + const void* scales_data, + const void* zero_points_data, + bool zero_points_is_typed, + int n, int k, int64_t k_padded) const { + constexpr int kUseCompileTimeCap = std::numeric_limits::max(); + + if constexpr (!(std::is_same_v || std::is_same_v)) { + return kUseCompileTimeCap; // float has no batched GEMV to tune. + } else { + if (!tune_small_m_) { + return kUseCompileTimeCap; + } + + int cached = resolved_small_m_cap_.load(std::memory_order_relaxed); + if (cached >= 0) { + return cached; + } + + // The batched fast path (and hence tuning) only applies to the column-wise, non-typed-zero-point + // case that the dequant fallback below also serves. Typed zero points and the chunked large-N path + // are left on the compile-time cap. + const int64_t scratch_bytes = static_cast(N_) * k_padded * static_cast(sizeof(T)); + const bool will_use_chunked = column_wise_quant_blk_ && + (force_chunked_ || + ((scratch_bytes > 256 * 1024 * 1024) && (N_ > chunk_target_rows_ * 2))); + if (!column_wise_quant_blk_ || will_use_chunked || zero_points_is_typed) { + resolved_small_m_cap_.store(kUseCompileTimeCap, std::memory_order_relaxed); + return kUseCompileTimeCap; + } + + const int zp_kind = zero_points_data == nullptr ? 0 : 1; // 0: none, 1: uint8 + const int dtype_tag = std::is_same_v ? 1 : 2; + const auto& device_prop = this->GetDeviceProp(); + const uint64_t key = MakeSmallMCapKey(dtype_tag, K_, N_, block_size_, nbits_, zp_kind, sm_, + device_prop.multiProcessorCount, device_prop.l2CacheSize); + { + std::lock_guard lock(g_small_m_cap_mutex); + auto it = g_small_m_cap_cache.find(key); + if (it != g_small_m_cap_cache.end()) { + resolved_small_m_cap_.store(it->second, std::memory_order_relaxed); + return it->second; + } + if (g_small_m_cap_in_progress.count(key) != 0) { + // Another thread is profiling this shape. Use the built-in cap for this call (a safe default) + // and pick up the tuned value on a later call, rather than profiling redundantly. Do not store + // to resolved_small_m_cap_ so this node retries. + return kUseCompileTimeCap; + } + g_small_m_cap_in_progress.insert(key); + } + + typedef typename onnxruntime::cuda::OrtToCudaType::type CudaT; + cudaStream_t stream = this->Stream(ctx); + const int max_probe = (nbits_ == 8) ? 8 : 16; // structural max M of the batched kernels (CtaM range). + + auto a_buf = this->template GetScratchBuffer(static_cast(max_probe) * k, this->GetComputeStream(ctx)); + auto y_buf = this->template GetScratchBuffer(static_cast(max_probe) * n, this->GetComputeStream(ctx)); + auto b_buf = this->template GetScratchBuffer(static_cast(n) * k_padded, this->GetComputeStream(ctx)); + CudaT* a_data = a_buf.get(); + CudaT* y_data = y_buf.get(); + CudaT* b_dequant = b_buf.get(); + CUDA_CALL_THROW(cudaMemsetAsync(a_data, 0, static_cast(max_probe) * k * sizeof(CudaT), stream)); + + const CudaT* scales = reinterpret_cast(scales_data); + const uint8_t* zp = static_cast(zero_points_data); + const size_t shared_mem_per_block = this->GetDeviceProp().sharedMemPerBlock; + const CudaT alpha = onnxruntime::cuda::OrtToCudaType::FromFloat(1.f); + const CudaT zero = onnxruntime::cuda::OrtToCudaType::FromFloat(0.f); + + auto time_ms = [&](auto&& fn, int iters) -> float { + for (int i = 0; i < 5; i++) fn(); + cudaEvent_t start, end; + CUDA_CALL_THROW(cudaEventCreate(&start)); + CUDA_CALL_THROW(cudaEventCreate(&end)); + CUDA_CALL_THROW(cudaEventRecord(start, stream)); + for (int i = 0; i < iters; i++) fn(); + CUDA_CALL_THROW(cudaEventRecord(end, stream)); + CUDA_CALL_THROW(cudaEventSynchronize(end)); + float ms = 0.f; + CUDA_CALL_THROW(cudaEventElapsedTime(&ms, start, end)); + CUDA_CALL_THROW(cudaEventDestroy(start)); + CUDA_CALL_THROW(cudaEventDestroy(end)); + return ms / iters; + }; + + auto run_batched = [&](int m) -> bool { + return TryMatMulNBits(static_cast(nbits_), y_data, a_data, blob_data, scales, zp, + /*bias_data*/ static_cast(nullptr), m, n, k, + static_cast(block_size_), shared_mem_per_block, stream, /*max_batched_m*/ m); + }; + + auto run_dequant = [&](int m) { + if (nbits_ == 8) { + ORT_IGNORE_RETURN_VALUE((Dequantize8Bits(b_dequant, blob_data, scales, zp, nullptr, + SafeInt(k_padded), n, SafeInt(block_size_), stream))); + } else { + ORT_IGNORE_RETURN_VALUE((Dequantize4Bits(b_dequant, blob_data, scales, zp, nullptr, + SafeInt(k_padded), n, SafeInt(block_size_), stream))); + } + ORT_IGNORE_RETURN_VALUE(cublasGemmHelper(GetCublasHandle(ctx), CUBLAS_OP_T, CUBLAS_OP_N, n, m, k, + &alpha, b_dequant, SafeInt(k_padded), a_data, k, + &zero, y_data, n, GetDeviceProp(), UseTF32())); + }; + + constexpr int kIters = 30; + // Warm up the batched kernel's CtaM tiers (each distinct CtaM/CtaN is its own instantiation) and the + // dequant path so the timed measurements below exclude one-time module-load / cuBLAS-setup costs. + for (int wm : {2, 4, 8, max_probe}) { + if (wm >= 2 && wm <= max_probe && run_batched(wm)) run_dequant(wm); + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + + // Does the batched GEMV beat dequant + cuBLAS at M=m, by a margin? median-of-3 per path (robust to a + // single noisy sample) and a small margin give hysteresis so near-ties don't flip the cap run-to-run. + constexpr float kMargin = 0.02f; // batched must be >=2% faster to be preferred. + auto batched_wins = [&](int m) -> bool { + auto med3 = [&](auto&& fn) { + float a = time_ms(fn, kIters), b = time_ms(fn, kIters), c = time_ms(fn, kIters); + return a + b + c - std::min({a, b, c}) - std::max({a, b, c}); + }; + const float tb = med3([&] { run_batched(m); }); + const float td = med3([&] { run_dequant(m); }); + return tb < td * (1.0f - kMargin); + }; + + // The batched kernel's cost is ~linear in M: it skips padded rows (only the real M rows compute), so + // batched(M) rises with M while dequant + cuBLAS is ~flat. Hence "batched wins" is monotone in M with + // a single crossover, and the cap can land at any M (not only CtaM tier ends). Binary-search the + // boundary (largest M where batched wins), seeding the first probe at the built-in cap so a GPU that + // matches the reference converges in ~2 probes. + const int seed = std::min((nbits_ == 8) ? 5 : 16, max_probe); + int cap = 1; // 1 disables batched (m>=2 all use dequant). + if (run_batched(2) && batched_wins(2)) { + cap = 2; + int lo = 3, hi = max_probe; + int probe = std::max(lo, std::min(hi, seed)); + while (lo <= hi) { + if (run_batched(probe) && batched_wins(probe)) { + cap = probe; + lo = probe + 1; + } else { + hi = probe - 1; + } + probe = (lo + hi) / 2; + } + } + + { + std::lock_guard lock(g_small_m_cap_mutex); + g_small_m_cap_cache[key] = cap; + g_small_m_cap_in_progress.erase(key); + } + resolved_small_m_cap_.store(cap, std::memory_order_relaxed); + return cap; + } +} + template Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { if constexpr (std::is_same_v) { @@ -372,6 +569,20 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { #endif if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType())) { + // Resolve the batched-GEMV cap for this shape/device (INT_MAX unless online tuning is enabled). + // Only M in the batched-candidate range can use the batched path, so restrict tuning to it: M==1 + // (decode) uses the single-row kernel and M beyond the structural max (e.g. prefill) always uses + // dequant + cuBLAS, so neither should pay the one-time profiling cost. + int max_batched_m = std::numeric_limits::max(); + const int batched_structural_max = (nbits_ == 8) ? 8 : 16; + if (m >= 2 && m <= batched_structural_max) { + const int64_t k_padded_for_tune = (K_ + block_size_ - 1) / block_size_ * block_size_; + max_batched_m = ResolveSmallMBatchedCap( + ctx, blob_data, reinterpret_cast(scales_data), zero_points_data, + /*zero_points_is_typed*/ (zero_points != nullptr && zero_points->IsDataType()), + n, k, k_padded_for_tune); + } + // First, try the fused fast path. It handles bias only for the GPT-OSS router GEMV // specialization; for other shapes it fails when bias is present. if (TryMatMulNBits( @@ -387,7 +598,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { k, SafeInt(block_size_), GetDeviceProp().sharedMemPerBlock, - stream)) { + stream, + max_batched_m)) { return Status::OK(); } @@ -408,7 +620,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { k, SafeInt(block_size_), GetDeviceProp().sharedMemPerBlock, - stream)) { + stream, + max_batched_m)) { LaunchMatMulNBitsBiasAdd( reinterpret_cast(Y->MutableData()), reinterpret_cast(bias_data), diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index 61eb62fef4a2b..cd36ac29f3ac7 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -2,6 +2,7 @@ // Licensed under the MIT License. #pragma once +#include #include "core/providers/cuda/shared_inc/cuda_utils.h" namespace onnxruntime { @@ -21,7 +22,8 @@ bool TryMatMul4Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m = std::numeric_limits::max()); template bool TryMatMul8Bits( @@ -35,7 +37,8 @@ bool TryMatMul8Bits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream); + cudaStream_t stream, + int max_batched_m = std::numeric_limits::max()); template bool TryMatMulNBits( @@ -51,18 +54,19 @@ bool TryMatMulNBits( int k, int block_size, size_t shared_mem_per_block, - cudaStream_t stream) { + cudaStream_t stream, + int max_batched_m = std::numeric_limits::max()) { if (bits == 8) { if (bias_data != nullptr) { return false; } return TryMatMul8Bits(output, a_data, b_data_quant, scales_data, zero_points, - m, n, k, block_size, shared_mem_per_block, stream); + m, n, k, block_size, shared_mem_per_block, stream, max_batched_m); } if (bits == 4) { return TryMatMul4Bits(output, a_data, b_data_quant, scales_data, zero_points, bias_data, - m, n, k, block_size, shared_mem_per_block, stream); + m, n, k, block_size, shared_mem_per_block, stream, max_batched_m); } return false; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 6d3fae53b4abe..b05d2896790ed 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -7,11 +7,13 @@ // pre-packed and block-compacted into int4 // #pragma once +#include #include "core/common/safeint.h" #include "core/providers/cuda/cuda_kernel.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "core/platform/env_var_utils.h" +#include "core/session/onnxruntime_session_options_config_keys.h" namespace onnxruntime { namespace contrib { @@ -85,6 +87,8 @@ class MatMulNBits final : public CudaKernel { chunk_target_rows_ = kDefaultChunkTargetRows; } + tune_small_m_ = info.GetConfigOptions().GetConfigEntry(kOrtSessionOptionsConfigMatMulNBitsTuneSmallM) == "1"; + #if USE_FPA_INTB_GEMM if constexpr (std::is_same::value || std::is_same::value) { int option = ParseEnvironmentVariableWithDefault(kFpAIntBGemmOption, 0); @@ -134,6 +138,15 @@ class MatMulNBits final : public CudaKernel { #endif private: + // Returns the max M for which the batched GEMV should be used for this op's fixed weight shape. + // When tuning is disabled it returns INT_MAX (use the compile-time cap). When enabled it lazily + // micro-benchmarks the batched GEMV vs dequant + cuBLAS once per shape and caches the crossover. + int ResolveSmallMBatchedCap(OpKernelContext* ctx, + const uint8_t* blob_data, + const void* scales_data, + const void* zero_points_data, + bool zero_points_is_typed, + int n, int k, int64_t k_padded) const; #if USE_FPA_INTB_GEMM void InitGemmProfiler(int sm); void RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int max_m); @@ -157,6 +170,11 @@ class MatMulNBits final : public CudaKernel { bool force_chunked_{false}; int64_t chunk_target_rows_{kDefaultChunkTargetRows}; + bool tune_small_m_{false}; + // Resolved batched-GEMV cap for this op: -1 = unresolved, otherwise the max M to run batched + // (INT_MAX means "use the compile-time cap"). Resolved once, then reused. + mutable std::atomic resolved_small_m_cap_{-1}; + #if USE_FPA_INTB_GEMM bool has_fpA_intB_gemv_{false}; bool has_fpA_intB_gemm_{false}; diff --git a/onnxruntime/test/python/transformers/profile_matmul_nbits.py b/onnxruntime/test/python/transformers/profile_matmul_nbits.py index 75edc892ede1f..5e2b9ee1a867f 100644 --- a/onnxruntime/test/python/transformers/profile_matmul_nbits.py +++ b/onnxruntime/test/python/transformers/profile_matmul_nbits.py @@ -4,24 +4,32 @@ # -------------------------------------------------------------------------- """ -Profiling script for the CUDA MatMulNBits (4-bit / 8-bit weight-only) op. +Profiling script for the CUDA MatMulNBits (4-bit / 8-bit weight-only) op and its online small-M +tuning option (session config ep.cuda.matmulnbits_tune_small_m). -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. +Three subcommands: + + case Time one K x N weight at a single M. NVTX-wrapped so nsys can capture kernel-level timing. + latency Latency table (avg us) across the representative decoder shapes for one config. + sweep Does tuning help? For every bits x dtype x block_size x zero-point combination, time each + shape with the tuner off and on and report the best speedup (built-in / tuned). + +The representative shapes are taken from real quantized decoders: Qwen3-4B (q4b:*), Qwen3-8B (q8b:*), +and a Gemma4 target (gm:*). Usage: - # Host-timing table across all cases: - python profile_matmul_nbits.py --warmup 25 --repeat 200 + # Does the online tuning help on this GPU, across every combination (the main question): + python profile_matmul_nbits.py sweep + + # Narrow the sweep (e.g. only 8-bit, block 32/128): + python profile_matmul_nbits.py sweep --bits 8 --blocks 32 128 - # Single case: - python profile_matmul_nbits.py --k 4096 --n 4096 --m 8 --block-size 32 --bits 4 --dtype fp16 + # Latency table for one config, with the tuner on: + python profile_matmul_nbits.py latency --bits 4 --dtype bf16 --tune - # Kernel-level via nsys + the repo parser: + # Single case, 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 profile_matmul_nbits.py case --k 4096 --n 4096 --m 8 python parse_nsys.py mnb.sqlite --nvtx-range benchmark --pattern '%' """ @@ -57,14 +65,28 @@ _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. +# Representative decoder weight matrices (K = in features, N = out features), taken from real +# quantized decoders: Qwen3-4B (hidden 2560), Qwen3-8B (hidden 4096), and a Gemma4 target (hidden +# 1536, wide MoE). Labels are model:projection. lm_head shapes are large-N and take the chunked +# dequant path (not the small-M batched kernel), so they are omitted from the default small-M set. DEFAULT_CASES = [ - ("qkv", 4096, 4096), - ("o_proj", 4096, 4096), - ("gate_up", 4096, 12288), - ("down", 12288, 4096), - ("lm_head", 4096, 151936), + # Qwen3-4B + ("q4b:q", 2560, 4096), + ("q4b:kv", 2560, 1024), + ("q4b:o", 4096, 2560), + ("q4b:gate_up", 2560, 9728), + ("q4b:down", 9728, 2560), + # Qwen3-8B + ("q8b:qo", 4096, 4096), + ("q8b:kv", 4096, 1024), + ("q8b:gate_up", 4096, 12288), + ("q8b:down", 12288, 4096), + # Gemma4 target + ("gm:q", 1536, 2048), + ("gm:kv", 1536, 256), + ("gm:o", 2048, 1536), + ("gm:gate_up", 1536, 12288), + ("gm:down", 12288, 1536), ] DEFAULT_MS = [1, 2, 4, 8, 16] @@ -124,17 +146,22 @@ def build_model(k, n, block_size, bits, onnx_dtype, with_zero_point=True): return model.SerializeToString() -def make_session(model): +MATMULNBITS_TUNE_SMALL_M_CONFIG = "ep.cuda.matmulnbits_tune_small_m" + + +def make_session(model, tune=False): so = onnxruntime.SessionOptions() so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL so.log_severity_level = 3 + if tune: + so.add_session_config_entry(MATMULNBITS_TUNE_SMALL_M_CONFIG, "1") return onnxruntime.InferenceSession(model, so, providers=["CUDAExecutionProvider"]) -def run_case(name, m, k, n, block_size, bits, dtype, warmup, repeat): +def run_case(name, m, k, n, block_size, bits, dtype, warmup, repeat, tune=False, zp=True, quiet=False): onnx_dtype = _OT[dtype] torch_dtype = _TT[dtype] - sess = make_session(build_model(k, n, block_size, bits, onnx_dtype)) + sess = make_session(build_model(k, n, block_size, bits, onnx_dtype, with_zero_point=zp), tune=tune) 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") @@ -166,40 +193,120 @@ def run_case(name, m, k, n, block_size, bits, dtype, warmup, repeat): "block_size": block_size, "bits": bits, "dtype": dtype, + "zp": int(zp), + "tune": int(tune), "avg_us": round(us, 2), } - print(RESULT_PREFIX + json.dumps(result)) + if not quiet: + 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"] +def print_latency_table(args): + """Latency table (avg us) across the representative shapes for one config, tuner off or on.""" + zp = not args.no_zp + rows = { + name: { + m: run_case(name, m, k, n, args.block_size, args.bits, args.dtype, args.warmup, args.repeat, args.tune, zp)[ + "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 + } + state = "tuned" if args.tune else "built-in cap" + zp_tag = "zp" if zp else "no-zp" + print( + f"\n MatMulNBits {args.bits}-bit block{args.block_size} {args.dtype} {zp_tag} [{state}] (avg us, lower is better)" + ) + print(" " + f"{'shape':12s} {'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)) + print(" " + f"{name:12s} {k:>6} {n:>7} " + " ".join(f"{rows[name][m]:7.1f}" for m in args.ms)) + + +def run_sweep(args): + """Exhaustive check of whether the online-tuning session option helps. For every combination of + bits x dtype x block_size x zero-point, each representative shape is timed with the tuner off and on; + the cell is the best speedup (built-in / tuned) over the probed M values, so >1.00 means tuning is + faster and ~1.00 means it makes no difference. Lets a user on any GPU see where (if anywhere) the + per-device crossover differs from the built-in cap.""" + for bits in args.bits: + structural_max = 8 if bits == 8 else 16 + ms = [m for m in args.ms if 2 <= m <= structural_max] + for dtype in args.dtypes: + cols = [(b, zp) for b in args.blocks for zp in (True, False)] + print(f"\n=== {bits}-bit {dtype}: best built-in/tuned speedup over M={ms} (>1.00 = tuner helps) ===") + print( + " " + f"{'shape':12s} {'K':>6} {'N':>7} " + + " ".join((f"b{b}" + ("+zp" if zp else "")).rjust(8) for b, zp in cols) + ) + for name, k, n in DEFAULT_CASES: + cells = [] + for b, zp in cols: + best = 1.0 + for m in ms: + off = run_case(name, m, k, n, b, bits, dtype, args.warmup, args.repeat, False, zp, quiet=True) + on = run_case(name, m, k, n, b, bits, dtype, args.warmup, args.repeat, True, zp, quiet=True) + best = max(best, off["avg_us"] / on["avg_us"]) + cells.append(f"{best:.2f}x") + print(" " + f"{name:12s} {k:>6} {n:>7} " + " ".join(c.rjust(8) for c in cells)) + + +def main(): + p = argparse.ArgumentParser(description="Profile CUDA MatMulNBits and its online small-M tuning option.") + sub = p.add_subparsers(dest="cmd", required=True) + + def add_timing(sp): + sp.add_argument("--warmup", type=int, default=25) + sp.add_argument("--repeat", type=int, default=200) + + c = sub.add_parser("case", help="time one K x N weight at a single M (for nsys or a quick check)") + c.add_argument("--k", type=int, required=True, help="in features") + c.add_argument("--n", type=int, required=True, help="out features") + c.add_argument("--m", type=int, required=True, help="rows") + c.add_argument("--bits", type=int, default=4, choices=[4, 8]) + c.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"]) + c.add_argument("--block-size", type=int, default=32) + c.add_argument("--no-zp", action="store_true", help="build without zero points") + c.add_argument("--tune", action="store_true", help="enable the tuning session option") + add_timing(c) + + lat = sub.add_parser("latency", help="latency table across the representative shapes for one config") + lat.add_argument("--bits", type=int, default=4, choices=[4, 8]) + lat.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"]) + lat.add_argument("--block-size", type=int, default=32) + lat.add_argument("--no-zp", action="store_true", help="build without zero points") + lat.add_argument("--tune", action="store_true", help="enable the tuning session option") + lat.add_argument("--ms", type=int, nargs="+", default=DEFAULT_MS) + add_timing(lat) + + sw = sub.add_parser("sweep", help="does tuning help? speedup across bits x dtype x block x zero-point") + sw.add_argument("--bits", type=int, nargs="+", default=[4, 8], choices=[4, 8]) + sw.add_argument("--dtypes", nargs="+", default=["fp16", "bf16"], choices=["fp16", "bf16"]) + sw.add_argument("--blocks", type=int, nargs="+", default=[16, 32, 64, 128]) + sw.add_argument("--ms", type=int, nargs="+", default=[6, 16], help="row counts to probe for the best speedup") + add_timing(sw) + + args = p.parse_args() + if args.cmd == "case": + run_case( + "case", + args.m, + args.k, + args.n, + args.block_size, + args.bits, + args.dtype, + args.warmup, + args.repeat, + args.tune, + not args.no_zp, + ) + elif args.cmd == "latency": + print_latency_table(args) + elif args.cmd == "sweep": + run_sweep(args) if __name__ == "__main__":