onnxruntime
308f24c3 - Enable fpA_intB GEMM in CUDA builds and add configurable options (#29622)

Commit
67 days ago
Enable fpA_intB GEMM in CUDA builds and add configurable options (#29622) ## Enable fpA_intB GEMM in CUDA builds and add configurable options ### Summary This PR turns the CUDA fpA_intB (weight-only, FP activation × int weight) MatMulNBits path on by default in CUDA builds, replaces the ambiguous `ORT_FPA_INTB_GEMM` bitmask with a simple on/off flag, and adds session-config keys so the path and its autotuning buckets can be controlled per session. It also makes the CUTLASS tactic profiler CUDA-graph safe and configurable. ### Motivation The fpA_intB kernels were previously gated behind `onnxruntime_USE_FPA_INTB_GEMM=OFF` and an `ORT_FPA_INTB_GEMM` integer bitmask (`0x01=all`, `0x02=GEMV`, `0x04=int4`, `0x08=int8`). The bitmask was ambiguous (e.g. `6` could read as "int4 + GEMV" or "GEMV for int4 and int8") and the GEMM/GEMV kernels actually share one weight layout, so splitting them was never valid. Shipping the kernels by default and exposing a plain enable flag plus per-session config makes the feature usable and tunable without rebuilding. ### Key Changes #### Build enablement | File | Change | |------|--------| | [cmake/CMakeLists.txt](cmake/CMakeLists.txt) | `onnxruntime_USE_FPA_INTB_GEMM` becomes a `cmake_dependent_option`, defaulting **ON** when `onnxruntime_USE_CUDA` is enabled (OFF otherwise). | | [tools/ci_build/github/linux/build_cuda_c_api_package.sh](tools/ci_build/github/linux/build_cuda_c_api_package.sh), [build_linux_python_package.sh](tools/ci_build/github/linux/build_linux_python_package.sh), [build_tensorrt_c_api_package.sh](tools/ci_build/github/linux/build_tensorrt_c_api_package.sh) | Flip packaging builds from `onnxruntime_USE_FPA_INTB_GEMM=OFF` to `ON`. | #### Option simplification (bitmask → boolean) - Removed `kFpAIntBGemmOption_All/Gemv/Int4/Int8` and the old bitmask parsing. - Added `ParseFpAIntBEnabled`: `""` / `"0"` / `"off"` → disabled; any other value → enabled (numeric non-zero still works for back-compat). - The enable flag now only governs nodes **without** prepacked weights. A prepacked weight is already stored in the fpA_intB layout, so the choice was fixed at export time; the constructor forces the path on for prepacked nodes and only `ORT_ENFORCE`s that the shape/hardware actually support it. - GEMV is no longer independently toggleable: it is enabled whenever supported, since GEMM and GEMV share the same weight layout. #### New session-config keys (EP-agnostic, config wins over env) | Config key | Env fallback | Meaning | |------------|--------------|---------| | `ep.cuda.fpa_intb_gemm` | `ORT_FPA_INTB_GEMM` | Enable/disable the fpA_intB path (`0`/`off` vs `1`/`on`). | | `ep.cuda.fpa_intb_profile_m` | `ORT_FPA_INTB_PROFILE_M` | Comma-separated initial profile-M buckets (e.g. `"1,8,64,512"`); empty uses the default bucket set. | These are read by both the built-in CUDA EP and the CUDA plugin EP via `OpKernelInfo::GetConfigOptions()`. #### Profiler: CUDA-graph-safe, in-memory autotuning - Added `getBestConfigOrProfile()` for lazy single-bucket profiling outside CUDA-graph capture; during capture the kernel falls back to a pure lookup (`getBestConfig`) because profiling launches kernels, records/synchronizes events, and allocates scratch — all illegal during capture. - Added configurable profile-M buckets: `ParseProfileMList`, `setProfileMOverride`, `getProfileMBuckets`, plus `kEnvProfileM` and `kDefaultProfileMaxM` (default max M lowered to `2048`). - Clearer error when an M bucket was not profiled before capture ("run a warmup inference outside capture first"). #### Docs - [docs/contrib_ops/cuda/matmul_nbits.md](docs/contrib_ops/cuda/matmul_nbits.md): `ORT_FPA_INTB_GEMM` documented as int/string on/off; clarified prepacked-weight strictness. ### Testing Notes - New: [onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py](onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py) — exercises the prepacked fpA_intB path and the boolean/numeric back-compat values of the enable flag. - To verify locally (CUDA, SM ≥ 75): - Build with CUDA (fpA_intB now defaults ON): `./build.sh --use_cuda ...` - Run: `python -m pytest onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py` - Sanity-check config override: set `ep.cuda.fpa_intb_gemm=0` on a non-prepacked node and confirm the path is skipped; confirm a prepacked node still forces the path on.
Author
Parents
Loading