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

[jambayk]

Description

Add an opt-in session config option, ep.cuda.matmulnbits_tune_small_m (default off), that replaces
the built-in per-device cap on the small-M batched GEMV (added in #29451) with one measured on the
running GPU.

• 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 (process-global, keyed on dtype/K/N/block_size/bits/zero-point kind). Decode (M=1) and
  prefill (M above the kernel's structural max) never trigger it; the cached value is reused for every
  later call, so steady-state inference is unaffected.

• The batched kernel skips padded rows, so its cost is ~linear in M while dequant + cuBLAS is ~flat.
  "Batched wins" is therefore monotone in M with a single crossover, found by a binary search over the
  candidate range seeded at the built-in cap, using median-of-3 timing and a small margin for stability.

• A runtime max_batched_m is threaded through TryMatMulNBits / TryMatMul4Bits / TryMatMul8Bits,
  and the 8-bit kernel's structural max M (8) is separated from its default cap so tuning can probe the
  full supported range.

• profile_matmul_nbits.py is reworked into case / latency / sweep subcommands over the real
  Qwen3-4B, Qwen3-8B and Gemma4 decoder shapes. sweep reports the built-in vs tuned speedup across
  every bits x dtype x block_size x zero-point combination, so the payoff can be checked on any GPU
  with a single command:

  python profile_matmul_nbits.py sweep
  

Best built-in/tuned speedup measured on A100 (SM80), block_size 32, with zero points (>1.00 means the
tuned cap is faster; block_size and zero-point were swept too and change the numbers only within noise):

shape (K, N)	4-bit fp16	4-bit bf16	8-bit fp16	8-bit bf16
q8b:kv (4096, 1024)	1.08x	1.06x	1.00x	1.00x
q8b:down (12288, 4096)	1.03x	1.00x	1.00x	1.10x
q8b:qo (4096, 4096)	1.00x	1.00x	1.00x	1.11x
q8b:gate_up (4096,12288)	1.00x	1.00x	1.05x	1.23x
q4b:gate_up (2560, 9728)	1.00x	1.00x	1.13x	1.30x
gm:gate_up (1536, 12288)	1.00x	1.01x	1.04x	1.20x
gm:o (2048, 1536)	1.00x	1.00x	1.02x	1.09x

Dense 4-bit shapes are mostly 1.00x (the built-in cap of 16 is already right on A100). Tuning pays off
most for 8-bit bf16 wide-MoE projections, and modestly for 8-bit fp16 and narrow-N 4-bit.

Motivation and Context

The small-M cap (the largest M for which the batched GEMV beats dequant + cuBLAS) is a single constant
tuned on one reference GPU. The crossover depends on the weight shape and on the device (SM count, L2,
tensor-core vs CUDA-core balance), so a fixed cap can either leave speedup unused or run the batched
path past its crossover on other hardware. This option measures the crossover per shape on the running
GPU instead.

It is opt-in and off by default: the only cost is a one-time micro-benchmark on the first small-M call
for each distinct weight shape (roughly five shapes per decoder, deduped process-globally), and steady
state is unchanged. The sweep subcommand is meant to gather cross-GPU data to decide whether the
threshold varies enough on consumer GPUs to warrant enabling it more broadly, following the pattern of
the cuDNN conv algorithm search (cudnn_conv_algo_search).

[Copilot]

Pull request overview

Adds an opt-in CUDA Execution Provider session option to online-tune the small‑M MatMulNBits dispatch threshold (batched GEMV vs dequantize+cuBLAS) per shape/device, and updates the profiling tooling to measure the impact.

Changes:

• Add session config ep.cuda.matmulnbits_tune_small_m and plumb it through MatMulNBits so the small‑M cap can be resolved at runtime and cached.
• Extend the MatMulNBits fast-path dispatch (TryMatMulNBits / 4-bit / 8-bit) to accept a runtime max_batched_m override, including separating the 8-bit kernel’s structural max from its default cap.
• Rework profile_matmul_nbits.py into case / latency / sweep subcommands to compare built-in vs tuned behavior across representative decoder shapes.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/profile_matmul_nbits.py	Adds subcommands and tuning toggle to benchmark built-in vs tuned small‑M behavior.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h	Reads the new session config and adds per-node cached resolved cap storage.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh	Threads a runtime max_batched_m override through the TryMatMul* helpers.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc	Implements one-time micro-benchmarking + process-global caching for the tuned cap and uses it in dispatch.
onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu	Separates structural max (8) from default cap (5) and applies runtime cap override.
onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu	Applies runtime cap override to skip small‑M kernels above the tuned threshold.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Adds the public session config key constant and docs.