[Yuto Ashida (y-chan)]

Description

Adds a fp16 fast path to TryMatMul8BitsM1 (matmul_8bits_m1_impl.cuh) for the single-row
(decode) GEMV of 8-bit MatMulNBits:

• one warp per block; each warp owns 1 output column (K >= 4096) or 2 columns (shorter K, so the
  A loads are shared between the two columns)
• every lane loads 16 bytes of B per step and up to four steps are issued before any is consumed
• scales and zero points are read directly from global memory (no shared-memory staging and no
  __syncthreads() before the first weight load)
• (q - zp) is formed exactly in fp16 and products are accumulated in fp32

The path is taken for fp16 when block_size is 16/32/64/128/256, K % 512 == 0 and (for the
2-column variant, K < 4096) N % 2 == 0. Everything else, and bf16/fp32, keeps the existing kernel.
ORT_MATMUL_NBITS_8BIT_M1_WIDE=0 selects the existing kernel (read once per process) for A/B
comparison; it can be removed before merge if preferred.

Motivation and Context

At M = 1 the 8-bit GEMV is bound by the weight read. The existing M=1 kernel reads 8 bytes of B per
lane per iteration and stages the scales in shared memory behind a barrier, so a warp keeps little
memory in flight: on A100 it streams a 2048x4096 int8 weight at about 0.6 TB/s from DRAM, while a
plain streaming read of the same bytes reaches about 1.1 TB/s.

Results

Kernel time per call (A100 80GB SXM4, CUDA 12.9, fp16, 8-bit weights with zero points, M = 1),
measured with the existing profile_matmul_nbits.py under nsys, existing kernel
(ORT_MATMUL_NBITS_8BIT_M1_WIDE=0) vs this PR:

Shape (K x N)	block 128: existing	block 128: this PR	block 32: existing	block 32: this PR
Qwen3-30B-A3B qkv (2048 x 5120)	11.97 us	9.30 us	12.60 us	9.40 us
Qwen3-30B-A3B o_proj (4096 x 2048)	11.79 us	8.30 us	12.91 us	8.39 us
Qwen3-30B-A3B lm_head (2048 x 151936)	258.1 us	203.8 us	280.8 us	209.8 us
Qwen3-8B-class qkv (4096 x 4096)	17.24 us	12.60 us	18.40 us	12.53 us
Qwen3-8B-class gate_up (4096 x 12288)	50.26 us	36.88 us	57.05 us	38.56 us
Qwen3-8B-class down (12288 x 4096)	51.02 us	40.87 us	58.97 us	43.50 us

The script runs one node repeatedly, so weights up to a few tens of MB stay in L2 (A100: 40 MB) and
the gap is smaller than in a model, where each layer's weights come from DRAM.

End to end, on a non-public decoder with the Qwen3-30B-A3B architecture (48 layers, 8-bit
MatMulNBits for the attention projections and lm_head, INT8 QMoE experts, CUDA graph, batch 1,
ORT_ENABLE_XQA=0 on both sides), measured on an earlier main (f2c39fe):

GPU	existing	this PR
A100 80GB PCIe	5.60 ms/token	5.20 ms/token
RTX PRO 6000 Blackwell Max-Q (sm120)	4.57 ms/token	4.24-4.41 ms/token (*)

(*) Shared machine with about 0.17 ms run-to-run variance; every run was faster than the baseline
of the same session. In the model, the 2048x4096 projection went from 14.2 us to 8.5 us (nsys).

Generated text is identical. Prefill is unchanged (M > 1 does not use this path).

Tests

• New MatMulNBits.Fp16_8b_M1_Cuda (matmul_8bits_test.cc): fp16, M = 1, with and without zero
  points, block sizes 16-256, K = 512 to 8192 so that each column/unroll instantiation runs, plus
  an odd N that falls back to the existing kernel. The existing CUDA fp16 8-bit tests only had
  M > 1 or K < 512.
• A100 (sm80, CUDA 12.9, Release): onnxruntime_provider_test --gtest_filter=*MatMulNBits*
  93 passed, onnxruntime_test_all --gtest_filter=*MatMulNBits* 54 passed.
• Tests and op-level numbers above are from main 7617fc8; the branch was then rebased onto
  c666fd4 without conflicts (the commits in between do not touch the M=1 8-bit path).

How to reproduce

cd onnxruntime/test/python/transformers
for w in 0 1; do
  ORT_MATMUL_NBITS_8BIT_M1_WIDE=$w nsys profile -t cuda,nvtx -o mnb_$w --export=sqlite \
    python profile_matmul_nbits.py --k 2048 --n 5120 --m 1 --bits 8 --block-size 128
  python parse_nsys.py mnb_$w.sqlite --nvtx-range benchmark --pattern '%'
done

🤖 Generated with Claude Code

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[Yuto Ashida (y-chan)]

@microsoft-github-policy-service agree company="InfoDeliver"

[Tianlei Wu (tianleiwu)]

Thanks for the thorough profiling. I read the wide kernel closely and found no correctness problem: the __byte_perm unpack gives an exact q - zp in fp16, the pairing with A is right, accumulation is fp32 (the scale is applied once per 16-element chunk), indexing is 64-bit, every (kCols, kUnroll) branch divides K, and the dispatch conditions match what the kernel supports.

Two small points are inline: three of the eight new test cases (N = 33 / 2) never reach the M1 path because TryMatMul8Bits requires N % 8 == 0 (so the "falls back to the original M1 kernel" comment is inaccurate), and the 16-byte vector loads of A have no alignment guard.

Validation: no CI has run on this PR yet (only license/cla); a maintainer needs to trigger /azp run. Locally on A100 (sm_80, CUDA 13.0, reduced build, gtests off): 920 fp16 M=1 shape x zero-point combinations (block 16-256, K 256-12288, N 8-4096) vs a float64 reference give worst relative error 5.7e-4 (6.9e-4 with the existing kernel); 769 of 770 wide-path outputs differ from the old kernel (max 9.5e-4 relative), so the new path is exercised. Chain benchmark with distinct weights (approximate): 4096x4096 19.9 -> 16.0 us (bs128) and 22.8 -> 16.6 us (bs32); 2048x2048 about 1.0-1.1x.

Recommend approving after verification on the impacted devices: Linux and Windows CUDA CI (including MatMulNBits.Fp16_8b_M1_Cuda), plus sm_90 (H100) and Blackwell (sm_100/sm_120). Verified so far: A100 / sm_80 locally. Not verified: any CI, Windows/MSVC (the new std::getenv call has precedent in group_query_attention_qdq.cuh, which builds on Windows CI), sm_86/89/90/100/120.

[Tianlei Wu (tianleiwu)]

Re-reviewed the vector-load alignment, block/scale boundaries, K-divisibility and unroll dispatch, and one-/two-column grid coverage. The earlier alignment guard and small-N coverage fixes are present, with tighter tolerances. No additional actionable findings.

The only current-head reported check is CLA. I did not repeat local CUDA validation; please obtain current-head CUDA build/tests covering the wide path and fallback before merging. No approval submitted.

[Yuto Ashida (y-chan)]

Thanks for the re-review. It looks like the workflow runs for the current head (ed7985f) have not started because they are waiting for approval, as this is my first contribution: 32 runs are in action_required, including Linux CUDA CI and ONNX Runtime CUDA Builds. Could a maintainer approve them? Locally this head passes onnxruntime_provider_test --gtest_filter=MatMulNBits.* on both the wide path and with ORT_MATMUL_NBITS_8BIT_M1_WIDE=0.

[Yuto Ashida (y-chan)]

While testing this on an RTX 3090 with an INT4 (block size 32) Qwen3-Omni-30B-A3B export, I confirmed that the same approach also helps 4-bit weights. The 4-bit decode GEMVs still go through MatMulFloat4BitsKernelM1. Routing them through this PR's wide kernel (16-byte B loads per lane, multi-column warps, scales read directly instead of staged behind a barrier) reduced MatMulNBits time from 818 µs to 755 µs per token (48 layers, packed QKV + o_proj). End-to-end decode went from 186.4 to 188.5 tok/s.

The 4-bit and 8-bit cases differ only in how a 16-byte chunk is unpacked and how many K elements it covers. So rather than keeping the kernel 8-bit only, I'd like to change this PR to a more general implementation that takes the bit width as a template parameter, with unchanged 8-bit behavior. Would that be OK?