[Saiyam Bharara (saiyambharara)]

Description

Add an opt-in direct-scale specialization for FP16, symmetric INT4, group-32 MatMulNBits with M=1. Read scales directly through the global-memory cache hierarchy, removing their cooperative shared-memory copy and block barrier. Preserve the eight-column block layout, accumulation order, warp reduction, eligibility checks and upstream tail handling.

Use four-way unrolling for K < 4096 and sixteen-way otherwise. Retain the staged kernel for K >= 4096 and N < 1024, where experiments found a regression.

Enable with ORT_MATMUL4BITS_M1_EXPERIMENT=2; mode 1 uses direct scales with sixteen-way unrolling. Unset/0 preserves default dispatch. The setting is cached per process.

Motivation and Context

Related to #32382. This reduces scale-staging and loop overhead while preserving the existing grid geometry. RTX 5070 Ti experiments also evaluated output regrouping, K-cooperating warps and weight-prefetch combinations; these did not provide consistent improvements across shapes.

Benchmarks

Real ORT models, M=1, FP16 inputs/scales, symmetric INT4, group size 32, CUDA 13.1.115. ORT_FPA_INTB_GEMM=0 isolates the original M1 route. CUDA events measure mean time per node in a 500-node CUDA graph with repeated resident weights. Twenty randomized ABBA/BAAB brackets per shape reject baseline drift above 15%; reductions use median paired speedups.

RTX 5070 Ti (12 GB) and H100 SXM 80 GB; the H100 repeat uses one ORT host thread per worker.

K	N	RTX 5070 Ti baseline -> candidate (us/node)	RTX 5070 Ti reduction	H100 initial reduction	H100 repeat reduction
3072	1024	3.1714 -> 2.5652	19.18%	19.61%	19.60%
3072	3072	9.7570 -> 8.7167	10.99%	17.16%	16.80%
3072	8192	23.9136 -> 22.7929	4.43%	7.92%	7.17%
8192	3072	21.1079 -> 20.9635	0.24%	5.12%	6.24%
3072	32000	211.2827 -> 207.6596	1.72%	1.48%	1.30%
4096	256	2.0029 -> 2.0104	-0.11%	-0.18%	-0.17%

Absolute timings are separate medians, so their ratio can differ from the median paired reduction. The RTX 5070 Ti K=8192/N=3072 result is inconclusive. K=4096/N=256 selects the original kernel in both modes. H100 K=3072/N=1024 measured 3.091136 -> 2.485032 us/node initially.

Validation

• RTX 5070 Ti, pre-rebase MatMulNBits.*: 73 passed, 8 existing skips, 0 failures. Final rebased source: 5 focused tests passed, zero skips, using selective compilation/linking against the cached runtime; the production FP16 CUDA translation unit also compiled.
• Both direct specializations execute with the experiment unset: 20 bitwise comparisons covering K tails, the 4096 boundary, multiple output blocks and cancellation inputs. A wrong-scale-indexing negative control failed both tests as expected.
• Historical H100 run: matching outputs across modes, 21 exact CUPTI route checks, and two focused Compute Sanitizer memcheck tests with zero errors.

Limitations

Opt-in only. Repeated weights and the strongest shapes fitting in L2 limit these results; integrated HBM-cold, fastest-default-route and full-model comparisons remain pending. Performance predates the latest upstream rebase, and H100 validation predates the final test layout. A full fresh-upstream build and Windows/CUDA 13.3 CI remain pending. H100 figures are retained run summaries; raw remote traces were not retained.

[azure-pipelines]

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

[Copilot]

🔵 Needs a closer look

The rebased CUDA kernel and test-target integration still need a fresh full-runtime build and outstanding Windows/CUDA 13.3 validation.

No open issues 🎉

What changed in this PR

Adds an opt-in direct-scale CUDA specialization for FP16, symmetric INT4, group-32 MatMulNBits with M=1. It targets scale-staging and loop overhead related to #32382 without changing default dispatch or addressing grid starvation.

Changes:

• Adds guarded direct-scale dispatch with four- or sixteen-way loop unrolling, preserving existing eligibility checks and fallbacks.
• Adds selection, fallback, and bitwise direct-kernel tests independent of the experiment setting.
• Integrates CUDA tests into the provider-test target and limits an MSVC warning option to C++ compilation.

No actionable code findings were identified.

File	Description
onnxruntime/​test/​contrib_ops/​matmul_4bits_test.cc	Adds selection, operator fallback, and direct-kernel test cases.
onnxruntime/​test/​contrib_ops/​cuda_kernels/​matmul_4bits_m1_test.cu	Compares both direct-scale variants bitwise against the staged kernel.
onnxruntime/​contrib_ops/​cuda/​quantization/​matmul_4bits_m1.cuh	Defines the guarded experiment-selection policy.
onnxruntime/​contrib_ops/​cuda/​quantization/​matmul_4bits_m1_impl.cuh	Implements direct-scale loading, configurable unrolling, and opt-in dispatch.
cmake/​onnxruntime_unittests.cmake	Wires CUDA tests and CCCL configuration into the test target.

_{Review effort: Balanced}

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) Akshay Sonawane (@apsonawane) Rui Ren (@rui-ren) could you review this FP16/group-32 MatMulNBits M=1 optimization related to #32382, or suggest the appropriate reviewer?

It adds opt-in direct-scale loads and guarded unrolling while preserving default dispatch. The PR includes RTX 5070 Ti/H100 benchmarks and direct correctness tests for both candidate kernels. Feedback on the selection policy would be helpful.

Could someone also approve the relevant CUDA/Windows CI runs?

[Tianlei Wu (tianleiwu)]

Thanks for the detailed experiment. Blocking: three Windows CUDA builds fail on the new .cu test source. Windows GPU CUDA CI (job), Windows CUDA Plugin EP Build (job) and Windows GPU TensorRT CI (job) all stop while nvcc + MSVC compile matmul_4bits_m1_test.cu: absl/container/internal/raw_hash_map.h C2794 / C2938, then MSB3721. It is the only error in those logs (the CUDA provider, including matmul_4bits.cu, compiles; Linux passes), so the Windows test jobs were skipped. The likely cause and a suggested fix are inline: this TU is built without the /permissive (and related) flags that the provider targets add for abseil headers under nvcc.

Otherwise I found no correctness problem in the kernel or dispatch: the direct path keeps the reduction order, the selector accepts only fp16 / group 32 / no zero point / N % 8 == 0 / K % 8 == 0, and the env var is read once into a function-local static (default unchanged, graph-consistent). One test gap (inline): default CI never sets ORT_MATMUL4BITS_M1_EXPERIMENT, so the op-level tests only run mode 0.

Validation: Linux CUDA, TensorRT and CUDA Plugin EP jobs pass (the new gtests run there); the three Windows build jobs fail. Locally on A100 (sm_80, CUDA 13.0, reduced build, gtests off): 176 fp16 M=1 group-32 cases vs a float64 reference, worst rel err 1.5e-3, modes 1/2 bit-identical to mode 0. Kernel time is mixed: mode 2 is -13% at N=32000/K=3072 and flat to -8% elsewhere, but +13..27% at N=1024/K=4096 (6.8 -> 7.7-8.6 us), worth a sweep before widening the dispatch. Not verified: Windows, sm_90, Blackwell.

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) I've pushed af5c14f addressing the two inline comments:

• Added the CUDA provider's /permissive, /Zc:__cplusplus, and /Zc:preprocessor options to CUDA compilation on the provider-test target, including the Windows target-name variation.
• Added separate-process runs of the two operator tests with ORT_MATMUL4BITS_M1_EXPERIMENT=1 and 2, both through CTest and in the CUDA plugin workflows that invoke the executable directly.

The focused mode-1/2 operator tests pass locally on an RTX 5070 Ti, and the CUDA test helper compiles locally with Windows CUDA 13.4/MSVC 14.51. Confirmation on the failing CI toolchain is still pending.

Your A100 regression at N=1024/K=4096 falls just outside our current long-K/narrow-N fallback (N < 1024). Could you share your existing timing results and benchmark harness/command? If convenient, a small mode-0 versus mode-2 sweep around N={512,1024,2048}, K={3072,4096,5120} would help establish the boundary. Otherwise, I can reproduce it on a rented A100.

[Tianlei Wu (tianleiwu)]

Thanks for the follow-up in af5c14fd68 — both previous review items (the MSVC NVCC host-compiler flags for onnxruntime_provider_test_target and running Fp16_Int4_DirectScaleExperiment / Cuda_DirectScaleExperimentFallbacks with ORT_MATMUL4BITS_M1_EXPERIMENT=1 and 2 in separate processes) are addressed.

resolved

[Copilot]

🟡 Changes recommended

The tests can pass without exercising the intended dispatch routes, and the non-plugin CUDA test linkage violates C++’s one-definition rule.

1 open issue

• Assert the expected kernel dispatch in mode tests · New

1 previously missed issue

In code that hasn't changed since last review

Avoid duplicate framework type definitions in CUDA tests

cmake/​onnxruntime_unittests.cmake:1638

In a non-plugin CUDA build, this source includes matmul_4bits_m1_impl.cuh, which pulls in provider_api.h and its bridge definitions of onnxruntime::Tensor and other framework types. Linking that translation unit into onnxruntime_provider_test alongside the native framework definitions violates C++'s one-definition rule. Move the kernel comparison into onnxruntime_providers_cuda_ut, or isolate the tested kernel in a provider-neutral header before compiling it into this executable.

_{Review effort: Max}

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) CI is green and your previous review items are resolved. I plan to address Copilot’s ODR concern by moving the kernel comparisons into the CUDA test module. Would you also require explicit dispatch-path assertions before approval?

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) I addressed Copilot?s findings in ffdd3fd663. The kernel comparisons now run in the CUDA test module, and the operator tests verify dispatch for modes 0, 1 and 2.

Local 5070 Ti tests and Windows compilation passed. Forcing a staged fallback correctly failed the new assertions.

The receipts follow the existing MatMulNBitsUnaligned4 stderr pattern. Would you prefer ORT logging here? Could you review the updated PR?

[Tianlei Wu (tianleiwu)]

Re-reviewed the direct/staged specialization, selector boundaries, native/plugin test compilation, and fresh-process modes 0/1/2. The earlier Windows configuration and cached-mode coverage concerns are addressed. This head also adds actual launch receipts to the operator tests, addressing the remaining thread about numerically equivalent fallback masking dispatch. No additional actionable findings.

Only CLA is reported on this head. Please obtain the current Windows/Linux CUDA checks, especially the new kernel-test compilation and mode-specific runs, before merging. I did not rebuild the CUDA provider locally. No approval submitted.

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) The current CUDA checks have run, and the new kernel comparisons and dispatch tests for modes 0, 1 and 2 passed.

The four remaining failed checks hit the same six T5 Python tests. CI installs SentencePiece 0.2.3 with Protobuf code generated by 7.36.1, while the installed runtime is 6.33.6. This breaks tokenizer initialization before inference.

These failures appear independent of the CUDA changes. Could you advise how we should unblock CI?

[Saiyam Bharara (saiyambharara)]

Tianlei Wu (@tianleiwu) I applied the upstream SentencePiece pin to address the T5 dependency failures and retained the adjacent upstream Pillow minimum to resolve the merge conflict. The CUDA implementation remains as reviewed. Could you approve the new CI runs and renew approval once they pass?