[CUDA] Online tuning of the MatMulNBits small-M batched GEMV cap

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.

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<T>()) {
     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<T>(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<float>(
     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>(
     half* output,
@@ -1181,7 +1191,8 @@ template bool TryMatMul4Bits<half>(
     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>(
     nv_bfloat16* output,
@@ -1195,7 +1206,8 @@ template bool TryMatMul4Bits<nv_bfloat16>(
     int k,
     int block_size,
     size_t shared_mem_per_block,
-    cudaStream_t stream);
+    cudaStream_t stream,
+    int max_batched_m);
 
 template <typename T>
 __global__ void MatMulNBitsBiasAddKernel(T* output, const T* bias_data, int n, int64_t total) {

onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu

@@ -5,6 +5,7 @@
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
 #include <math_constants.h>
+#include <limits>
 #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 <class T>
-__host__ __device__ constexpr int SmallMCap8() {
-  return kSmallMMax8;
-}
+constexpr int kBatchedStructuralMax8 = 8;
 
 template <class T>
 __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<T>() || 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<int>::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<float>(
     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>(
     half* output,
@@ -689,7 +707,8 @@ template bool TryMatMul8Bits<half>(
     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<nv_bfloat16>(
@@ -703,7 +722,8 @@ template bool TryMatMul8Bits<nv_bfloat16>(
     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