Enable fpA_intB GEMM in CUDA builds and add configurable options#29622
Open
tianleiwu wants to merge 11 commits into
Open
Enable fpA_intB GEMM in CUDA builds and add configurable options#29622tianleiwu wants to merge 11 commits into
tianleiwu wants to merge 11 commits into
Conversation
Reduce the construction-time tactic profiling sweep and lazily profile any
runtime M bucket on demand, keeping all tuned tactics in the existing
process-global in-memory profile map. No disk/file interaction.
- gemm_profiler.h: fix reader/writer lock naming (shared read, exclusive
write); add getProfileMBuckets() (virtual, default = historical dense
sweep) so subclasses can profile a smaller bucket set; add
getBestConfigOrProfile() that serves an already-profiled bucket under a
shared lock and otherwise profiles a single bucket under an exclusive lock
and caches it in the in-process map; remember mDims/mHasWeightOnlyCudaKernel
from the initial sweep so lazy profiling can rebuild the tactic context.
- fpA_intB_gemm_profiler.{h,cc}: override getProfileMBuckets() with a small
default bucket set, configurable via ORT_FPA_INTB_PROFILE_M; add
ProfileMaxM()/ParseProfileMOverride() to bound the initial sweep top-M.
- matmul_nbits.{h,cc}: size the initial sweep with ProfileMaxM(); during
inference call getBestConfigOrProfile() to tune unseen M on the fly, but
fall back to getBestConfig() while the compute stream is being captured
into a CUDA graph (lazy profiling launches kernels and is illegal during
capture).
This is the in-memory half of the fpA_intB autotune work; the persistent
(on-disk) tactic cache is a separate change.
Make the fpA_intB MatMulNBits path (and its in-memory tactic autotuning) part of the default CUDA build so it is actually exercised: - cmake: option(... OFF) -> cmake_dependent_option(... ON "onnxruntime_USE_CUDA" OFF), matching the FLASH_ATTENTION / MEMORY_EFFICIENT_ATTENTION pattern. DISABLE_CONTRIB_OPS still forces it OFF. - Linux CUDA / CUDA-C-API / TensorRT-C-API packaging scripts: flip the explicit onnxruntime_USE_FPA_INTB_GEMM=OFF override to =ON so shipped packages include the kernels.
Provide session-option config keys that work identically for the built-in CUDA EP and the CUDA plugin EP (both create the MatMulNBits kernel via KernelRegistryManager::CreateKernel, which injects the session-level ConfigOptions; the plugin EP reuses the same kernel). Each key overrides its existing ORT_* environment variable, with config taking precedence: - ep.cuda.fpa_intb_gemm <-> ORT_FPA_INTB_GEMM (enable/mode; accepts off/on/all or a numeric kFpAIntBGemmOption_* bitmask) - ep.cuda.fpa_intb_profile_m<-> ORT_FPA_INTB_PROFILE_M (initial profile M buckets) The kernel resolves both settings (ResolveFpAIntBConfigOrEnv) and threads the profile-M list into the profiler instance (setProfileMOverride), so the bucket override is now per-session instead of process-global. ParseProfileMOverride/ ProfileMaxM (env-reading statics) are replaced by ParseProfileMList(string); the profiler no longer reads the environment itself. Uses the provider-bridge ConfigOptions already visible via cuda_kernel.h (no extra include), so it also compiles in the plugin build.
Add TestFpAIntBConfigKeys (CUDA, needs only quantize_matmul_4bits, not the offline packer) covering: the config key enabling the fpA_intB path (parity vs the standard dequant baseline for off/on/all/0x4 forms), session config winning over ORT_FPA_INTB_GEMM env, the profile-M key being accepted, env-var backward compat, and the prepacked-requires-enable error referencing ep.cuda.fpa_intb_gemm.
The ORT_FPA_INTB_GEMM / ep.cuda.fpa_intb_gemm bitmask (gemv/int4/int8/all bits) was ambiguous (e.g. "6" could read as "int4+gemv" or "gemv for int4 and int8"). Replace it with a plain on/off flag: - "" / 0 / off / false / none -> disabled; 1 / on / true / all -> enabled. Any other non-zero integer is still treated as enabled for backward compat. - "on" enables the full fpA_intB path: the CUTLASS GEMM plus the GEMV decode kernel where supported. GEMM and GEMV share one weight layout, so GEMV can no longer be toggled independently; it is enabled whenever supported. The 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 and cannot be turned off from configuration: the constructor forces the path on for prepacked nodes and only ORT_ENFORCEs that the shape/hardware actually supports it. Drop kFpAIntBGemmOption_All/Gemv/Int4/Int8 and ParseFpAIntBOption; add ParseFpAIntBEnabled. Update tests (boolean values, numeric back-compat) and the runtime comment.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR enables building the CUDA fpA_intB (weight-only FP activation × int weight) MatMulNBits GEMM path by default in CUDA builds, simplifies runtime enablement to a boolean-style flag, and adds per-session configuration keys for toggling the path and controlling initial autotuning buckets. It also updates the CUTLASS tactic profiler to avoid illegal work during CUDA graph capture via a capture-aware lazy-profiling path.
Changes:
- Turn on
onnxruntime_USE_FPA_INTB_GEMMby default for CUDA builds and flip Linux packaging builds to compile it in. - Add session-config keys (
ep.cuda.fpa_intb_gemm,ep.cuda.fpa_intb_profile_m) with env-var fallback and update fpA_intB enablement/autotuning behavior. - Make profiler selection CUDA-graph safe by avoiding lazy profiling during capture; add configurable/reduced M-bucket profiling.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/ci_build/github/linux/build_tensorrt_c_api_package.sh | Build package with fpA_intB kernels compiled in (onnxruntime_USE_FPA_INTB_GEMM=ON). |
| tools/ci_build/github/linux/build_linux_python_package.sh | Enable fpA_intB kernels in CUDA Python packaging build. |
| tools/ci_build/github/linux/build_cuda_c_api_package.sh | Enable fpA_intB kernels in CUDA C API packaging build. |
| onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py | Add tests for new session-config keys and env-var back-compat; adjust skips. |
| onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h | Add session-config key plumbing and simplify fpA_intB enable flag parsing. |
| onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc | Use capture-aware tactic selection (lookup vs lazy profile) and improve error messaging. |
| onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h | Add lazy single-bucket profiling API and customizable M-bucket profiling. |
| onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h | Add profile-M env key/defaults and override bucket support. |
| onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc | Implement profile-M parsing and reduced/configurable bucket selection. |
| docs/contrib_ops/cuda/matmul_nbits.md | Update fpA_intB enable flag documentation and clarify prepacked-weight behavior. |
| cmake/CMakeLists.txt | Make onnxruntime_USE_FPA_INTB_GEMM default ON when CUDA is enabled. |
- Fix Windows C4267 (size_t->int) warnings-treated-as-errors in fpA_intB_gemm_profiler.cc (getWorkspaceSize/computeTmpSize casts). These surfaced now that onnxruntime_USE_FPA_INTB_GEMM defaults ON. - Update Fp16_Int4_PrepackedWeightRequiresFpAIntBGemm test: a prepacked weight now forces the fpA_intB path on regardless of the enable flag, so the run no longer fails when the flag is off. Retarget it to verify rejection on an unsupported node (block_size=256) via the prepacked support ORT_ENFORCE. - Make ParseFpAIntBEnabled case-insensitive for "off" and fix a comment typo. - Query CUDA graph capture status unconditionally (a null stream is the default stream, a valid capture target under per-thread default streams). - Clarify prepacked-weight enable-flag wording in matmul_nbits.md.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_GEMMbitmask 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=OFFand anORT_FPA_INTB_GEMMinteger bitmask (0x01=all,0x02=GEMV,0x04=int4,0x08=int8). The bitmask was ambiguous (e.g.6could 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
onnxruntime_USE_FPA_INTB_GEMMbecomes acmake_dependent_option, defaulting ON whenonnxruntime_USE_CUDAis enabled (OFF otherwise).onnxruntime_USE_FPA_INTB_GEMM=OFFtoON.Option simplification (bitmask → boolean)
kFpAIntBGemmOption_All/Gemv/Int4/Int8and the old bitmask parsing.ParseFpAIntBEnabled:""/"0"/"off"→ disabled; any other value → enabled (numeric non-zero still works for back-compat).ORT_ENFORCEs that the shape/hardware actually support it.New session-config keys (EP-agnostic, config wins over env)
ep.cuda.fpa_intb_gemmORT_FPA_INTB_GEMM0/offvs1/on).ep.cuda.fpa_intb_profile_mORT_FPA_INTB_PROFILE_M"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
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.ParseProfileMList,setProfileMOverride,getProfileMBuckets, pluskEnvProfileMandkDefaultProfileMaxM(default max M lowered to2048).Docs
ORT_FPA_INTB_GEMMdocumented as int/string on/off; clarified prepacked-weight strictness.Testing Notes
./build.sh --use_cuda ...python -m pytest onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.pyep.cuda.fpa_intb_gemm=0on a non-prepacked node and confirm the path is skipped; confirm a prepacked node still forces the path on.