onnxruntime
cd50f9d0 - [CUDA] Add opt-in wave-aware fpA_intB GEMV dispatch (#33186)

Commit
4 days ago
[CUDA] Add opt-in wave-aware fpA_intB GEMV dispatch (#33186) ## Description Add an opt-in wave-aware column-tile heuristic for the CUDA MatMulNBits fpA_intB GEMV. Decode projections can leave the final block wave underfilled; selecting among the existing CtaN=2/4/8 kernels improves wave utilization on the measured sm_120 configuration. ## Summary of Changes - Enable with `SessionOptions.add_session_config_entry("ep.cuda.fpa_intb_gemv_wave_aware", "1")`. Unset or `"0"` retains the existing dispatch; other explicit values are rejected. - Keep the existing sm_12x, fp16/int4 SM80-interleaved, no-zero-point, M=8 guards. Other devices, types, layouts, and M values retain their existing dispatch even with the option enabled. - Carry the per-session flag into both tactic profiling and inference. Include dispatch mode in the shared tactic-cache key so opt-in and opt-out sessions cannot reuse one another's profiling results. - Add off/on/off session-option parity tests for M=7/8/9 and N=512/10240, plus invalid-value coverage. - Squash the original implementation and its follow-up dispatch fix into a single commit, rebased onto current `main`. ## Testing - Repository lintrunner: RUFF, RUFF-FORMAT, and CLANGFORMAT passed. - Compile-only checks of the current MatMulNBits and fpA_intB profiler translation units passed with the compact CUDA plugin defines. - Standalone CUDA graph probe against this branch on RTX 5090: 18 launch/output checks passed. Default/off uses CtaN=4; opt-in M=8 uses CtaN=2 for N=512/10240; M=7/9 remain unchanged. Exact controlled-data outputs and distinct cache entries passed. - Added Python session tests have not run against a rebuilt full plugin yet. Run `python -m pytest onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py -k wave_aware` with a build containing this change. - Earlier controlled model A/B for the wave-aware heuristic measured approximately 2.3% decode improvement on RTX 5090. That measurement predates this default-off session-option wrapper; no performance claim is made for other GPUs. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading