[CUDA] Batched small-M GEMV for MatMulNBits (4-bit/8-bit)

[jambayk]

Description

Add a batched GEMV for the small-M range (M = 2..16) of CUDA MatMulNBits
(4-bit and 8-bit), using the standard [N, blocks, blob] weight layout with
no prepacking.

Previously, MatMulNBits had a fast single-row GEMV for M = 1 (decode), but
for M > 1 it fell back to full weight dequantization + cuBLAS, which
dequantizes the entire weight matrix regardless of M. That fallback is
therefore flat across small M and dominates latency at the row counts that
matter for multi-row decode.

The new half/bf16 path uses a CtaM x CtaN register-tiled kernel that streams
the quantized weight once per block row and reuses each activation load across
columns, so latency scales with M. M = 1 decode is unchanged. No prepacking is
used, so there is no extra resident weight memory and no GEMM tactic profiling
at session init.

Also adds profile_matmul_nbits.py (same style as profile_qmoe_gemv.py,
parseable with parse_nsys.py) and a docs experiment log.

Before (main: M > 1 dequant + cuBLAS) vs After (batched small-M GEMV) —
A100, block_size 32, fp16, average op latency in microseconds (lower is
better).

4-bit (M = 2..16)

Before:

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:

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):

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

8-bit (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;
the batched path covers M = 2..5 and M >= 6 keeps the fallback.

Before:

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:

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):

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

Motivation and Context

The small-M (M = 2..16) regime is exactly what speculative decoding hits when
verifying a block of draft tokens: each step runs the target model on a handful
of rows rather than a single token. With the previous dispatch, that step paid
the full-dequant + cuBLAS cost (flat ~172 us for an MLP projection, ~1.9 ms for
the lm_head even at M = 2), which erased much of the speculative-decoding
speedup over greedy decoding. Routing these row counts through a batched GEMV
that scales with M (2.6-5.7x faster at M = 2, at or above parity through M = 16)
restores the benefit, without the resident-memory and session-init costs of a
prepacked weight layout.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR improves CUDA MatMulNBits latency for small row counts by adding a batched GEMV path for 4-bit and 8-bit weight-only matmuls using the standard [N, blocks, blob] weight layout (no prepacking), and adds tooling/docs to profile and record baseline results.

Changes:

• Add/register-tiled batched small‑M CUDA kernels for 4-bit and 8-bit MatMulNBits and route small M through them instead of full dequantize + cuBLAS.
• Add a Python profiling script to benchmark representative decoder shapes with NVTX ranges for nsys parsing.
• Add a markdown experiment log capturing a point-in-time performance baseline.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu	Adds a batched small‑M GEMV implementation and dispatch routing for 4-bit.
onnxruntime/contrib_ops/cuda/quantization/matmul_8bits.cu	Adds a batched small‑M GEMV implementation and dispatch routing for 8-bit.
onnxruntime/test/python/transformers/profile_matmul_nbits.py	Adds a profiling utility for timing MatMulNBits across shapes/M values with NVTX ranges.
docs/contrib_ops/cuda/matmul_nbits_small_m_experiments.md	Adds a recorded baseline of profiling results and command templates for replication.

[tianleiwu]

Review: Batched small-M GEMV for MatMulNBits (4-bit/8-bit)

Overall this is a clean, well-scoped, and carefully-documented change. The dispatch is conservative (m==1 unchanged, larger m still falls back to dequant+cuBLAS), the register-tiled kernels avoid extra resident memory / session-init profiling, and the int4->half prmt repack being applied once per column weight rather than per activation row is a nice touch. The WPack/PackNatural ordering vs. the naturally-loaded uint4 activation checks out, the packed zero-point indexing matches the standard [N, blocks, blob] layout, and the uint4/uint64 load alignment holds given the k % 8 == 0 / k % block_size == 0 gates. No correctness blockers found.

A few non-blocking suggestions:

1. Test coverage for the new tile branches. The batched kernels introduce CtaM ∈ {2,4,8,16} and CtaN ∈ {1,2} tilings, but the CUDA cases in onnxruntime/test/contrib_ops/matmul_4bits_test.cc only exercise M ∈ {1, 2, 100}. That leaves the CtaM=4/8/16 tiles, the CtaN=2 column-reuse branch (requires N % 16 == 0), and the 8-bit CtaM=8 path with skipped rows (M=3, M=5) without direct coverage — a regression in those tiles would pass CI. Consider adding half/bf16 cases at M=4, 8, 16 (and 8-bit M=3, 5) with an N divisible by 16 so both the CtaN=2 and the row-skip (r < valid) paths are hit.

2. Register-tiled path gated behind a shared-memory budget it doesn't use. In TryMatMul4Bits, the shared_mem_size > shared_mem_per_block early-return runs before the if (m >= 2) branch. TryMatMulBatched4Bits (half/bf16) launches with 0 shared memory, so for large K where the staged scale/zp buffer would exceed the per-block limit, the fast register-tiled path is skipped and it falls back to cuBLAS even though it needs no shared memory. Rare in practice, but the gate is stricter than necessary for the batched path.

3. bf16 batched path on pre-SM80 (question). DequantizeEight / AccumulateRow / PackNatural / DotAccum / HorizontalAdd for nv_bfloat16 are all guarded by __CUDA_ARCH__ >= 800 with empty/zero fallbacks, and TryMatMulBatched4Bits has no runtime SM check, so on SM < 800 with bf16 the kernel would launch and write zeros rather than returning false. This mirrors the existing single-row bf16 pattern (bf16 MatMulNBits is effectively SM80+ only), so it's likely a non-issue — just confirming bf16 can't reach this on pre-Ampere.

Nitpicks: WarpUniform(threadIdx.y) in the 4-bit batched kernel vs. plain threadIdx.y in the 8-bit one — threadIdx.y is already warp-uniform, so the wrapper is redundant; worth making the two consistent. SmallMCap<T>() / SmallMCap8<T>() are templated on T but always return the constant, which reads as if per-dtype caps exist when they don't.

[jambayk]

Thanks for the thorough review @tianleiwu. Addressed in c43d376:

1. Test coverage for the new tile branches. Added CUDA half/bf16 tests that exercise the tiles the previous M ∈ {1,2,100} cases missed:

• Fp16_Int4_SmallMBatchedTiles / BFloat16_Int4_SmallMBatchedTiles: M ∈ {3,4,5,8,12,16} (CtaM 4/8/16, with M values not a multiple of CtaM to hit the r < valid row-skip path), N ∈ {256, 24} (N%16==0 → CtaN=2 column-reuse, N%16!=0 → CtaN=1), block_size {32,128}, with and without zero points.
• Fp16_Int8_SmallMBatchedTiles / BFloat16_Int8_SmallMBatchedTiles: same, M ∈ {2,3,4,5} (8-bit caps the batched path at M=5; M=3/5 hit the CtaM=8 row-skip).

All four pass on A100 (SM80); the full *MatMulNBits* suite is green (58 passed, 1 pipeline-skip).

2. Register-tiled path gated behind a shared-memory budget it does not use. Fixed — the batched attempt now runs before the shared-memory gate in both TryMatMul4Bits and TryMatMul8Bits, so large-K shapes no longer fall back to cuBLAS. Details in the inline thread.

3. bf16 batched path on pre-SM80. Confirmed non-issue — it mirrors the existing single-row bf16 pattern exactly (same __CUDA_ARCH__ >= 800 empty-fallback helpers, no host SM check; bf16 MatMulNBits is effectively SM80+). Details in the inline thread.

Nitpicks:

• WarpUniform(threadIdx.y): kept for within-file consistency — all four kernels in matmul_4bits.cu (including the pre-existing single-row ones) wrap threadIdx.y, while matmul_8bits.cu has never defined the helper, so each file is internally consistent. Happy to drop it across the 4-bit file if you would prefer, but that touches the pre-existing kernels.
• SmallMCap<T>() / SmallMCap8<T>(): agreed the template parameter is unused today; kept as a seam for a future per-dtype cap. The comment now states the cap is the same for all dtypes to avoid implying otherwise.

[tianleiwu]

LGTM.

The threshold tuned in A100 machine might not fit consumer GPUs. It is better to use online profiler to choose best kernel for given input shape. That could be a future work, and need not be in this PR.

BTW, the gemv kernel of fpA_intB is faster and that support M <= 16 too. It is not enabled in production, and you can test it with a build flag for now.