From d66449b826988d992c62d5db4545c663355cf2dc Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 09:06:00 +0000 Subject: [PATCH 01/12] [CUDA] In-memory tactic autotuning for fpA_intB MatMulNBits 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. --- .../cuda/llm/fpA_intB_gemm_profiler.cc | 72 ++++++++++ .../cuda/llm/fpA_intB_gemm_profiler.h | 17 +++ .../contrib_ops/cuda/llm/gemm_profiler.h | 134 ++++++++++++++---- .../cuda/quantization/matmul_nbits.cc | 20 ++- .../cuda/quantization/matmul_nbits.h | 2 +- 5 files changed, 216 insertions(+), 29 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc index e5b15856a6c05..150cb2aacc2b3 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -17,6 +17,11 @@ #if USE_FPA_INTB_GEMM #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "contrib_ops/cuda/llm/common/workspace.h" +#include "core/platform/env_var_utils.h" + +#include +#include +#include using namespace onnxruntime::llm::common; using namespace onnxruntime::llm::kernels::cutlass_kernels; @@ -97,5 +102,72 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int /*n*/, i return true; } +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMOverride() { + const std::string value = onnxruntime::ParseEnvironmentVariableWithDefault(kEnvProfileM, ""); + std::vector result; + if (value.empty()) { + return result; + } + std::stringstream ss(value); + std::string token; + std::set unique; + while (std::getline(ss, token, ',')) { + // Trim surrounding whitespace. + size_t start = token.find_first_not_of(" \t"); + size_t end = token.find_last_not_of(" \t"); + if (start == std::string::npos) { + continue; + } + token = token.substr(start, end - start + 1); + try { + int m = std::stoi(token); + if (m > 0) { + unique.insert(m); + } + } catch (const std::exception&) { + // Ignore malformed entries. + } + } + result.assign(unique.begin(), unique.end()); + return result; +} + +int WeightOnlyGroupwiseQuantGemmPluginProfiler::ProfileMaxM() { + auto override_ms = ParseProfileMOverride(); + if (!override_ms.empty()) { + return override_ms.back(); // sorted ascending + } + return kDefaultProfileMaxM; +} + +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getProfileMBuckets( + int minM, int maxM, bool /*hasWeightOnlyCudaKernel*/) const { + int const lo = std::max(1, minM); + int const hi = std::max(lo, maxM); + + std::set buckets; + + auto override_ms = ParseProfileMOverride(); + if (!override_ms.empty()) { + for (int m : override_ms) { + buckets.insert(std::min(std::max(lo, m), hi)); + } + } else { + // Small default bucket set clamped to [lo, hi]. + static const int kDefault[] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048}; + for (int m : kDefault) { + if (m >= lo && m <= hi) { + buckets.insert(m); + } + } + } + + // Always include the decode bucket (M=1) and the top bucket so both extremes are tuned. + buckets.insert(lo); + buckets.insert(hi); + + return std::vector(buckets.begin(), buckets.end()); +} + } // namespace onnxruntime::llm::kernels::weight_only #endif diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h index b1336f45cab27..c7c7ec661545f 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h @@ -43,12 +43,27 @@ constexpr int32_t INT4_BITS = 4; constexpr int32_t FP16_INT4_RATIO = FP16_BITS / INT4_BITS; constexpr int32_t FP16_INT8_RATIO = FP16_BITS / INT8_BITS; +// Comma-separated list of M buckets to profile for MatMulNBits/fpA_intB. Overrides the default +// reduced bucket set. Example: ORT_FPA_INTB_PROFILE_M="1,8,64,512". +constexpr const char* kEnvProfileM = "ORT_FPA_INTB_PROFILE_M"; + +// Default top M that bounds the initial profile sweep when no override is given. Larger runtime +// M values are handled by lazy single-bucket profiling. +constexpr int kDefaultProfileMaxM = 2048; + class WeightOnlyGroupwiseQuantGemmPluginProfiler : public GemmPluginProfiler { public: using Config = onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; + // Parses kEnvProfileM into a sorted, de-duplicated, positive M list (empty when unset). + static std::vector ParseProfileMOverride(); + + // Returns the top M used to size the initial profile range. Honors kEnvProfileM if set, + // otherwise kDefaultProfileMaxM. + static int ProfileMaxM(); + void setQuant(int bits, bool has_bias, bool has_zeros) { mQuantBits = bits; mHasBiases = has_bias; @@ -74,6 +89,8 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler bool checkTactic(int m, int n, int k, Config const& tactic) const override; + std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const override; + private: bool mHasBiases; bool mHasZeros; diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h index d9c19cb734e91..291d8ac87e791 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h @@ -110,10 +110,10 @@ class GemmPluginProfiler { using MProfileMap = std::unordered_map>; using MProfileMapPtr = std::shared_ptr; + // requires shared ownership to read from *this + using reader_lock = std::shared_lock; // requires exclusive ownership to write to *this - using reader_lock = std::unique_lock; - // requires shared ownership to read from other - using writer_lock = std::shared_lock; + using writer_lock = std::unique_lock; // Struct of continuing map if GEMMs to the best profiles for different Ms struct MNKProfileMap { @@ -168,6 +168,13 @@ class GemmPluginProfiler { std::optional getBestConfig(int m, GemmIdType const& gemmId) const; + // Like getBestConfig, but if the requested M bucket has not been profiled yet, profiles it + // lazily (single bucket) and inserts it into the in-process map. This briefly blocks the caller + // but guarantees a tuned tactic for any runtime M, which is what makes the reduced first-time M + // sweep safe. Must not be called while the compute stream is being captured into a CUDA graph + // (the caller is responsible for using getBestConfig instead during capture). + std::optional getBestConfigOrProfile(int m, GemmIdType const& gemmId); + virtual int getMaxProfileM() const; protected: @@ -183,6 +190,11 @@ class GemmPluginProfiler { virtual void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream); + // Returns the ordered set of M buckets to profile during the initial sweep, given the + // (rounded) profile range [minM, maxM]. The default reproduces the historical dense sweep. + // Subclasses may override to profile a smaller, configurable bucket set. + virtual std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const; + private: std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics); @@ -214,6 +226,10 @@ class GemmPluginProfiler { bool mSkip{false}; + // Remembered from the initial profileTactics call so lazy single-bucket profiling can + // reproduce the same tactic candidate set. + bool mHasWeightOnlyCudaKernel{false}; + onnxruntime::AllocatorPtr mAllocator; }; @@ -277,8 +293,10 @@ void GemmPluginProfiler::profileT mRunner = runner; mType = type; + mDims = dims; + mHasWeightOnlyCudaKernel = hasWeightOnlyCudaKernel; - int const maxM = std::min(nextPowerOfTwo(dims.maxM), getMaxProfileM()); + int const maxM = std::min(nextPowerOfTwo(static_cast(dims.maxM)), getMaxProfileM()); size_t workspace_bytes = computeTmpSize(maxM, dims.n, dims.k); @@ -317,28 +335,12 @@ void GemmPluginProfiler::profileT CUDA_CALL_THROW(cudaStreamCreate(&mStream)); - int const startMinMRounded = nextPowerOfTwo(dims.minM); - - if (hasWeightOnlyCudaKernel) { - // Profile tactics for finer granularity of M, - // if CUDA kernel is enabled for weight-only plugins - int minM = dims.minM; - for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) { - profileTactics(m, dims.n, dims.k); - } - - for (int m = 16; m < maxM; m *= 2) { - profileTactics(m, dims.n, dims.k); - } - } else { - // Profile tactics for CUTLASS kernel only - for (int m = std::max(1, startMinMRounded); m < maxM; m *= 2) { - profileTactics(m, dims.n, dims.k); - } + // Profile the (possibly reduced) set of M buckets. Any unprofiled runtime M is handled + // later by lazy single-bucket profiling in getBestConfigOrProfile. + for (int m : getProfileMBuckets(static_cast(dims.minM), maxM, hasWeightOnlyCudaKernel)) { + profileTactics(m, static_cast(dims.n), static_cast(dims.k)); } - profileTactics(maxM, dims.n, dims.k); - if (isAllocated) { // Free tmp data mWorkspaceTmp.reset(); @@ -372,6 +374,90 @@ std::optional GemmPluginProfiler +std::vector GemmPluginProfiler::getProfileMBuckets( + int minM, int maxM, bool hasWeightOnlyCudaKernel) const { + // Default: reproduce the historical dense sweep so any other users of this template keep + // their behavior. Subclasses may override this to profile a smaller bucket set. + std::vector buckets; + if (hasWeightOnlyCudaKernel) { + for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) { + buckets.push_back(m); + } + for (int m = 16; m < maxM; m *= 2) { + buckets.push_back(m); + } + } else { + for (int m = std::max(1, nextPowerOfTwo(minM)); m < maxM; m *= 2) { + buckets.push_back(m); + } + } + buckets.push_back(maxM); + return buckets; +} + +template +std::optional GemmPluginProfiler::getBestConfigOrProfile( + int m, GemmIdType const& gemmId) { + if (mSkip) { + return std::nullopt; + } + + int const target = std::min(std::max(1, nextPowerOfTwo(m)), getMaxProfileM()); + + // Fast path: an already-profiled (exact or rounded) bucket under a shared read lock. + { + reader_lock lock(mMNKProfileMap->mutex); + if (mMNKProfileMap->existsMProfileMap(gemmId)) { + auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); + if (mProfileMap->count(m) > 0) { + return mProfileMap->at(m); + } + if (mProfileMap->count(target) > 0) { + return mProfileMap->at(target); + } + } + } + + writer_lock lock(mMNKProfileMap->mutex); + + if (!mMNKProfileMap->existsMProfileMap(gemmId)) { + mMNKProfileMap->createMProfileMap(gemmId); + } + auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); + + // Re-check under the writer lock: another thread may have profiled it meanwhile. + if (mProfileMap->count(m) > 0) { + return mProfileMap->at(m); + } + if (mProfileMap->count(target) > 0) { + return mProfileMap->at(target); + } + + // We can only profile lazily if the profiling context from construction is available. + if (mRunner == nullptr || mAllocator == nullptr || !mDims.isInitialized()) { + ORT_LLM_LOG_WARNING("Cannot lazily profile an unprofiled M bucket: profiler context is unavailable."); + return std::nullopt; + } + + int const n = static_cast(mDims.n); + int const k = static_cast(mDims.k); + size_t const workspace_bytes = computeTmpSize(target, n, k); + + CUDA_CALL_THROW(cudaStreamCreate(&mStream)); + mWorkspaceTmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); + initTmpData(target, n, k, mWorkspaceTmp.get(), workspace_bytes, mStream); + + auto tactics = this->getTactics(target, n, k); + auto best = this->profileTacticsForProblem(target, n, k, tactics); + mProfileMap->insert({target, best}); + + mWorkspaceTmp.reset(); + CUDA_CALL_THROW(cudaStreamDestroy(mStream)); + + return best; +} + template std::optional GemmPluginProfiler::profileTacticsForProblem( int m, int n, int k, std::vector const& tactics) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index abd43dc6f1f6c..08d23e5c00302 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -17,6 +17,7 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #endif #include "contrib_ops/cuda/llm/common/logger.h" #include "contrib_ops/cpu/quantization/matmul_nbits_helper.h" @@ -378,11 +379,22 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const void* fpA_intB_weight = is_prepacked_weight_ ? fpA_intB_weight_buffer_.get() : static_cast(blob_data); - auto const bestTactic = gemmProfiler_->getBestConfig(m, gemmId_); + // During CUDA graph capture we must not lazily profile: profiling launches kernels, records + // and synchronizes events, and allocates/frees scratch, all of which are illegal while the + // compute stream is being captured. Fall back to a lookup of an already-profiled bucket + // (warmup runs before capture populate these); only outside capture do we allow lazy + // single-bucket profiling. + const bool stream_is_capturing = + stream != nullptr && onnxruntime::llm::common::isCapturing(stream); + auto const bestTactic = stream_is_capturing ? gemmProfiler_->getBestConfig(m, gemmId_) + : gemmProfiler_->getBestConfigOrProfile(m, gemmId_); if (!bestTactic.has_value()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "No valid fpA_intB MatMulNBits tactic for M=", m, - ", N=", n, ", K=", k); + return ORT_MAKE_STATUS( + ONNXRUNTIME, FAIL, + "No valid fpA_intB MatMulNBits tactic for M=", m, ", N=", n, ", K=", k, + stream_is_capturing + ? ". The M bucket was not profiled before CUDA graph capture; run a warmup inference outside capture first." + : ""); } // Env-gated diagnostics (ORT_FPA_INTB_DEBUG=1): dump the selected tactic, the kernel path diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 9a7bfe2895582..03fffa7c88e19 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -135,7 +135,7 @@ class MatMulNBits final : public CudaKernel { InitGemmProfiler(FpAIntBPackingSmForKernel()); - constexpr int max_m = 8291; + int max_m = WeightOnlyGroupwiseQuantGemmPluginProfiler::ProfileMaxM(); RunGemmProfile(has_fpA_intB_gemv_, 1, max_m); has_fpA_intB_gemm_ = true; } From 864535398e14f22933807227bc5ec392d78d3915 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 09:19:48 +0000 Subject: [PATCH 02/12] Enable onnxruntime_USE_FPA_INTB_GEMM by default when CUDA is enabled 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. --- cmake/CMakeLists.txt | 2 +- tools/ci_build/github/linux/build_cuda_c_api_package.sh | 2 +- tools/ci_build/github/linux/build_linux_python_package.sh | 2 +- tools/ci_build/github/linux/build_tensorrt_c_api_package.sh | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 09e307e124316..16e3eb172017e 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -107,7 +107,7 @@ option(onnxruntime_USE_LEAN_ATTENTION "Build lean attention kernel for scaled do cmake_dependent_option(onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION "Build memory efficient attention kernel for scaled dot product attention" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_FP4_QMOE "Build CUDA QMoE FP4 kernels" OFF) option(onnxruntime_USE_FP8_QMOE "Build CUDA QMoE FP8 kernels" OFF) -option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" OFF) +cmake_dependent_option(onnxruntime_USE_FPA_INTB_GEMM "Build FpA IntB gemm cuda kernels" ON "onnxruntime_USE_CUDA" OFF) option(onnxruntime_USE_INT4_KV_CACHE "Build cuda kernels for int4 kv cache" OFF) option(onnxruntime_USE_FP8_KV_CACHE "Build cuda kernels for fp8 kv cache" ON) option(onnxruntime_QUICK_BUILD "Speed up build by skipping some kernels for faster development" OFF) diff --git a/tools/ci_build/github/linux/build_cuda_c_api_package.sh b/tools/ci_build/github/linux/build_cuda_c_api_package.sh index b06fbc70c8ca9..ffae4b2ff2cae 100755 --- a/tools/ci_build/github/linux/build_cuda_c_api_package.sh +++ b/tools/ci_build/github/linux/build_cuda_c_api_package.sh @@ -23,5 +23,5 @@ docker run -e SYSTEM_COLLECTIONURI --rm \ --parallel --nvcc_threads 1 --flash_nvcc_threads 1 \ --use_cuda --cuda_version=$CUDA_VERSION --cuda_home=/usr/local/cuda-$CUDA_VERSION --cudnn_home=/usr/local/cuda-$CUDA_VERSION \ --skip_tests --use_vcpkg --use_vcpkg_ms_internal_asset_cache \ ---cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=OFF' \ +--cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=ON' \ && cd /build/Release && make install DESTDIR=/build/installed" diff --git a/tools/ci_build/github/linux/build_linux_python_package.sh b/tools/ci_build/github/linux/build_linux_python_package.sh index 2db052a15f22e..e7695c393fe04 100755 --- a/tools/ci_build/github/linux/build_linux_python_package.sh +++ b/tools/ci_build/github/linux/build_linux_python_package.sh @@ -88,7 +88,7 @@ if [ "$BUILD_DEVICE" == "GPU" ]; then #Enable CUDA EP. BUILD_ARGS+=("--use_cuda" "--cuda_version=$SHORT_CUDA_VERSION" "--cuda_home=$CUDA_HOME" "--cudnn_home=$CUDA_HOME") BUILD_ARGS+=("--nvcc_threads=1" "--flash_nvcc_threads=1") - BUILD_ARGS+=("--cmake_extra_defines" "CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}" "onnxruntime_USE_FPA_INTB_GEMM=OFF") + BUILD_ARGS+=("--cmake_extra_defines" "CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}" "onnxruntime_USE_FPA_INTB_GEMM=ON") # Enable TRT EP only if TensorRT is installed. if [ -f /usr/include/NvInfer.h ]; then BUILD_ARGS+=("--use_tensorrt" "--tensorrt_home=/usr") diff --git a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh index 8f03577bb546e..675d48c5bf7f9 100755 --- a/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh +++ b/tools/ci_build/github/linux/build_tensorrt_c_api_package.sh @@ -23,5 +23,5 @@ docker run -e SYSTEM_COLLECTIONURI --rm --volume /data/onnx:/data/onnx:ro --volu --skip_submodule_sync --use_binskim_compliant_compile_flags --build_shared_lib --build_java --build_nodejs \ --parallel --nvcc_threads 1 --flash_nvcc_threads 1 \ --use_tensorrt --cuda_version=$CUDA_VERSION --cuda_home=/usr/local/cuda-$CUDA_VERSION --cudnn_home=/usr --tensorrt_home=/usr \ - --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=OFF' \ + --cmake_extra_defines 'CMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}' 'onnxruntime_USE_FPA_INTB_GEMM=ON' \ --use_vcpkg --use_vcpkg_ms_internal_asset_cache && cd /build/Release && make install DESTDIR=/build/installed" From 89a9ca46548de68ad6d01a00a3b1c07c97ed280b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 09:52:58 +0000 Subject: [PATCH 03/12] Add ep.cuda.fpa_intb_gemm / fpa_intb_profile_m session-config keys 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. --- .../cuda/llm/fpA_intB_gemm_profiler.cc | 17 +---- .../cuda/llm/fpA_intB_gemm_profiler.h | 17 +++-- .../cuda/quantization/matmul_nbits.h | 65 ++++++++++++++++++- 3 files changed, 76 insertions(+), 23 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc index 150cb2aacc2b3..f102587a00197 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -17,7 +17,6 @@ #if USE_FPA_INTB_GEMM #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "contrib_ops/cuda/llm/common/workspace.h" -#include "core/platform/env_var_utils.h" #include #include @@ -102,8 +101,7 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int /*n*/, i return true; } -std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMOverride() { - const std::string value = onnxruntime::ParseEnvironmentVariableWithDefault(kEnvProfileM, ""); +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(const std::string& value) { std::vector result; if (value.empty()) { return result; @@ -132,14 +130,6 @@ std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMOverri return result; } -int WeightOnlyGroupwiseQuantGemmPluginProfiler::ProfileMaxM() { - auto override_ms = ParseProfileMOverride(); - if (!override_ms.empty()) { - return override_ms.back(); // sorted ascending - } - return kDefaultProfileMaxM; -} - std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getProfileMBuckets( int minM, int maxM, bool /*hasWeightOnlyCudaKernel*/) const { int const lo = std::max(1, minM); @@ -147,9 +137,8 @@ std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getProfileMBuckets( std::set buckets; - auto override_ms = ParseProfileMOverride(); - if (!override_ms.empty()) { - for (int m : override_ms) { + if (!mProfileMOverride.empty()) { + for (int m : mProfileMOverride) { buckets.insert(std::min(std::max(lo, m), hi)); } } else { diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h index c7c7ec661545f..be6f1a2d82d1c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h @@ -57,12 +57,16 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler public: using Config = onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; - // Parses kEnvProfileM into a sorted, de-duplicated, positive M list (empty when unset). - static std::vector ParseProfileMOverride(); - - // Returns the top M used to size the initial profile range. Honors kEnvProfileM if set, - // otherwise kDefaultProfileMaxM. - static int ProfileMaxM(); + // Parses a comma-separated list of M buckets (e.g. "1,8,64,512") into a sorted, de-duplicated, + // positive list (empty when the string is empty/blank). Used for the ep.cuda.fpa_intb_profile_m + // session-config key and the ORT_FPA_INTB_PROFILE_M env var, both resolved by the kernel. + static std::vector ParseProfileMList(const std::string& value); + + // Overrides the initial profile M-bucket set for this profiler instance (per session). An empty + // list keeps the built-in default bucket set. Resolved by the kernel from session config / env. + void setProfileMOverride(std::vector ms) { + mProfileMOverride = std::move(ms); + } void setQuant(int bits, bool has_bias, bool has_zeros) { mQuantBits = bits; @@ -98,6 +102,7 @@ class WeightOnlyGroupwiseQuantGemmPluginProfiler int mGroupSize; KernelType mCudaKernelType; int mArch; + std::vector mProfileMOverride; }; } // namespace onnxruntime::llm::kernels::weight_only diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 03fffa7c88e19..3fe687dcc5216 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -13,6 +13,10 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "core/platform/env_var_utils.h" +#include +#include +#include + namespace onnxruntime { namespace contrib { namespace cuda { @@ -43,6 +47,51 @@ constexpr int kFpAIntBGemmOption_Int8 = 0x08; constexpr int64_t kMatMulNBitsWeightNotPrepacked = 0; constexpr int64_t kMatMulNBitsWeightPrepackedSm80 = 1; constexpr int64_t kMatMulNBitsWeightPrepackedSm90 = 2; + +// Session-option config keys. These are readable by BOTH the built-in CUDA EP and the CUDA plugin +// EP: every kernel is created via KernelRegistryManager::CreateKernel, which injects the +// session-level ConfigOptions, and the plugin CUDA EP wraps a CUDAExecutionProvider that reuses this +// same kernel. Each key overrides its ORT_* environment-variable equivalent (config wins). +// ep.cuda.fpa_intb_gemm <-> ORT_FPA_INTB_GEMM (enable/mode) +// ep.cuda.fpa_intb_profile_m <-> ORT_FPA_INTB_PROFILE_M (initial profile M buckets) +constexpr const char* kConfigFpAIntBGemm = "ep.cuda.fpa_intb_gemm"; +constexpr const char* kConfigFpAIntBProfileM = "ep.cuda.fpa_intb_profile_m"; + +// Resolves a setting from the session config first (per-session, EP-agnostic), then the environment +// variable, else empty. Session config wins so a model/session can override a process-wide env var. +inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const char* config_key, + const char* env_key) { + const std::string from_config = info.GetConfigOptions().GetConfigOrDefault(config_key, ""); + if (!from_config.empty()) { + return from_config; + } + return ParseEnvironmentVariableWithDefault(env_key, ""); +} + +// Parses the fpA_intB enable/mode setting. Accepts (case-insensitive): "" / "0" / "off" / "false" / +// "none" -> disabled; "on" / "true" / "all" -> all; otherwise a numeric bitmask of +// kFpAIntBGemmOption_* (decimal or 0x-prefixed hex), e.g. "4" = int4 only, "6" = int4 + GEMV. +inline int ParseFpAIntBOption(const std::string& value) { + if (value.empty()) { + return 0; + } + std::string v; + v.reserve(value.size()); + for (char c : value) { + v.push_back(static_cast(std::tolower(static_cast(c)))); + } + if (v == "0" || v == "off" || v == "false" || v == "none") { + return 0; + } + if (v == "on" || v == "true" || v == "all") { + return kFpAIntBGemmOption_All; + } + try { + return static_cast(std::stol(v, nullptr, /*base*/ 0)); + } catch (const std::exception&) { + return 0; + } +} #endif template @@ -106,9 +155,11 @@ class MatMulNBits final : public CudaKernel { } if constexpr (std::is_same::value || std::is_same::value) { - int option = ParseEnvironmentVariableWithDefault(kFpAIntBGemmOption, 0); + // Enable/mode from session config (ep.cuda.fpa_intb_gemm) with ORT_FPA_INTB_GEMM env fallback. + int option = ParseFpAIntBOption(ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBGemm, kFpAIntBGemmOption)); ORT_ENFORCE(!(weight_prepacked_ != kMatMulNBitsWeightNotPrepacked && option == 0), - "weight_prepacked requires the fpA_intB path, but ORT_FPA_INTB_GEMM is off for this node"); + "weight_prepacked requires the fpA_intB path, but it is disabled for this node " + "(set session config ep.cuda.fpa_intb_gemm=1 or env ORT_FPA_INTB_GEMM=1)"); // Note: a fused bias (input[5]) is fully supported by the fpA_intB GEMV, CUTLASS SM80/SM90 // GEMM (EpilogueOpBias), and the tactic profiler, so bias-bearing nodes (e.g. gpt-oss // qkv_proj/o_proj) are eligible. Only g_idx/reorder remains unsupported by this path. @@ -135,7 +186,15 @@ class MatMulNBits final : public CudaKernel { InitGemmProfiler(FpAIntBPackingSmForKernel()); - int max_m = WeightOnlyGroupwiseQuantGemmPluginProfiler::ProfileMaxM(); + // Initial profile M buckets from session config (ep.cuda.fpa_intb_profile_m) with + // ORT_FPA_INTB_PROFILE_M env fallback; empty -> profiler uses its default bucket set. + std::vector profile_m = WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList( + ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBProfileM, + onnxruntime::llm::kernels::weight_only::kEnvProfileM)); + gemmProfiler_->setProfileMOverride(profile_m); + + int max_m = profile_m.empty() ? onnxruntime::llm::kernels::weight_only::kDefaultProfileMaxM + : profile_m.back(); RunGemmProfile(has_fpA_intB_gemv_, 1, max_m); has_fpA_intB_gemm_ = true; } From 5ec9dc5f416dae7ce185bf67ba87bbbc0a923d8d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 17:17:27 +0000 Subject: [PATCH 04/12] Test ep.cuda.fpa_intb_gemm / fpa_intb_profile_m session-config keys 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. --- .../test_op_matmulnbits_prepacked_cuda.py | 111 ++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 406eee21c059f..6647165cacbc7 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -171,5 +171,116 @@ def test_int8_sm90_prepacked_weight_matches_runtime_prepack(self): self._check_sm90_parity(bits=8, block_size=128, m=32) +@unittest.skipIf("CUDAExecutionProvider" not in ort.get_available_providers(), "CUDA is not available") +@unittest.skipUnless(hasattr(_pybind, "quantize_matmul_4bits"), "MatMulNBits 4-bit quantizer is unavailable") +class TestFpAIntBConfigKeys(unittest.TestCase): + """Session-config keys ep.cuda.fpa_intb_gemm / ep.cuda.fpa_intb_profile_m. + + These do not need the offline weight packer (pack_weights_for_cuda_mixed_gemm), so they run in + more build configurations than TestMatMulNBitsPrepackedCuda. They cover: the config key enabling + the fpA_intB path, session config overriding the ORT_FPA_INTB_GEMM env var, the profile-M key + being accepted, and the prepacked-requires-enable error message referencing the config key. + """ + + def setUp(self): + # Make sure no env override leaks in from the process / other tests. + for name in ("ORT_FPA_INTB_GEMM", "ORT_FPA_INTB_PROFILE_M"): + os.environ.pop(name, None) + + def _quantize_weight(self, weight: np.ndarray, bits: int, block_size: int): + k, n = weight.shape + k_blocks = (k + block_size - 1) // block_size + blob_size = block_size * bits // 8 + q_weight = np.zeros((n, k_blocks, blob_size), dtype=np.uint8) + scales = np.zeros((n, k_blocks), dtype=np.float16) + zero_points = np.zeros((n, (k_blocks + 1) // 2), dtype=np.uint8) + _pybind.quantize_matmul_4bits(q_weight, weight, scales, zero_points, block_size, n, k, True) + return q_weight, np.abs(scales) + + def _make_model(self, m, k, n, q_weight, scales, bits, block_size, weight_prepacked=0) -> ModelProto: + node = helper.make_node( + "MatMulNBits", + ["A", "B", "scales"], + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=bits, + block_size=block_size, + weight_prepacked=weight_prepacked, + ) + graph = helper.make_graph( + [node], + "fpa_intb_config_keys_test", + [helper.make_tensor_value_info("A", TensorProto.FLOAT16, [m, k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [m, n])], + initializer=[ + numpy_helper.from_array(q_weight, name="B"), + numpy_helper.from_array(scales, name="scales"), + ], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ) + model.ir_version = 10 + return model + + def _run(self, model: ModelProto, a: np.ndarray, config: dict[str, str] | None = None) -> np.ndarray: + so = ort.SessionOptions() + for key, value in (config or {}).items(): + so.add_session_config_entry(key, value) + sess = ort.InferenceSession(model.SerializeToString(), so, providers=["CUDAExecutionProvider"]) + return sess.run(None, {"A": a})[0] + + def _make_int4_case(self, m=32, k=256, n=512, block_size=64): + rng = np.random.default_rng(2024) + a = rng.normal(0.0, 0.25, size=(m, k)).astype(np.float16) + weight = rng.normal(0.0, 0.25, size=(k, n)).astype(np.float16) + q_weight, scales = self._quantize_weight(weight, 4, block_size) + model = self._make_model(m, k, n, q_weight, scales, 4, block_size) + return model, a, q_weight, scales + + def test_config_key_enables_fpa_intb(self): + # Baseline runs the standard dequant path (fpA_intB disabled); the config key must switch to + # the fpA_intB path and stay numerically equivalent. + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + for value in ("1", "all", "on", "0x4"): + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": value}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2, err_msg=f"value={value}") + + def test_profile_m_config_key_accepted(self): + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": "1", "ep.cuda.fpa_intb_profile_m": "1,8,32"}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) + + def test_session_config_overrides_env(self): + # env var says off, session config says on -> the session config must win. + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + with set_env("ORT_FPA_INTB_GEMM", "0"): + out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": "1"}) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) + + def test_env_var_backward_compatible(self): + model, a, _, _ = self._make_int4_case() + ref = self._run(model, a) + with set_env("ORT_FPA_INTB_GEMM", "1"): + out = self._run(model, a) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) + + def test_prepacked_requires_enable_reports_config_key(self): + # A prepacked node requires the fpA_intB path; when it is disabled the construction-time + # error must point at the session config key (and the env var). + model, a, q_weight, scales = self._make_int4_case() + prepacked = self._make_model(a.shape[0], a.shape[1], q_weight.shape[0], q_weight, scales, 4, 64, + weight_prepacked=1) + with self.assertRaises(Exception) as ctx: + self._run(prepacked, a) # nothing enables the path + self.assertIn("ep.cuda.fpa_intb_gemm", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() From c1bc617ff73ea8d5eac48515e5d5d6e2afab4c05 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 8 Jul 2026 19:11:23 +0000 Subject: [PATCH 05/12] Simplify fpA_intB enable option to boolean; prepacked forces the path 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. --- docs/contrib_ops/cuda/matmul_nbits.md | 5 +- .../cuda/quantization/matmul_nbits.cc | 3 +- .../cuda/quantization/matmul_nbits.h | 81 ++++++++----------- .../test_op_matmulnbits_prepacked_cuda.py | 26 +++--- 4 files changed, 46 insertions(+), 69 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index a09a80a4732ed..637b4a685737b 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -263,8 +263,7 @@ Prepacked weights are intentionally strict: - If ORT was built without `onnxruntime_USE_FPA_INTB_GEMM=ON`, any nonzero `weight_prepacked` value throws during kernel construction. -- If `ORT_FPA_INTB_GEMM` is unset or `0`, any nonzero `weight_prepacked` value - throws instead of silently falling back to a raw-layout path. +- If any nonzero `weight_prepacked` value will cause `ORT_FPA_INTB_GEMM` be ignored. - Nonzero `weight_prepacked` requires FP16 or BF16 input `A`, because only the CUDA fpA_intB path consumes this layout. - `weight_prepacked` must match the layout the selected kernel expects: `1` is @@ -293,7 +292,7 @@ present. `ComputeInternal` then: | Variable | Type / default | Effect | |----------|----------------|--------| | `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION` | bool, `0` | Disable the router GEMV specialization (§4.2); shapes fall back to the generic GEMV / dequant path. Useful for A/B benchmarking. | -| `ORT_FPA_INTB_GEMM` | int bitmask, `0` | Enable the CUTLASS weight-only path (§6). `0x01` = all, `0x02` = CUDA GEMV, `0x04` = int4, `0x08` = int8. `0` disables it. | +| `ORT_FPA_INTB_GEMM` | int/string, `0` | Enable the CUTLASS weight-only path (§6). `0` or `off` disables it, otherwise enables it. | | `ORT_MATMULNBITS_FORCE_CHUNKED` | int, `0` | Force the chunked dequant+GEMM fallback (§5) regardless of the size heuristic. | | `ORT_MATMULNBITS_CHUNK_SIZE` | int64, `32768` | Target rows per chunk in the chunked fallback. Values `< 1` reset to the default. | diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 08d23e5c00302..3e52aeee2afb8 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -372,7 +372,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { if constexpr (std::is_same::value || std::is_same::value) { if (has_fpA_intB_gemm_) { // We expect weight/scale/zero_point(optional) inputs are initializers and have been prepacked. - // User could disable it by setting ORT_FPA_INTB_GEMM=0 if those tensors cannot be prepacked (It is rare). + // A non-prepacked node can opt out by setting session config ep.cuda.fpa_intb_gemm=0 (or env + // ORT_FPA_INTB_GEMM=0) if those tensors cannot be prepacked (it is rare). const bool has_fpA_intB_weight = is_prepacked_weight_ || weight_prepacked_ != kMatMulNBitsWeightNotPrepacked; ORT_ENFORCE(has_fpA_intB_weight && is_prepacked_scale_ && (is_prepacked_zero_point_ || !has_zero_points_), "To use fpA_intB_gemm, prepacking must be done on weight, scale and zero point."); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 3fe687dcc5216..49840cb0dae26 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -38,12 +38,10 @@ using onnxruntime::llm::kernels::weight_only::WeightOnlyGroupwiseQuantGemmPlugin using GemmProfilerPtr = std::shared_ptr; using WeightOnlyGemmRunnerPtr = std::shared_ptr; -// Environment variable to configure fpA_intB_gemm for experiments. Set it to 0 to disable, 1 to eanble all. +// Environment variable to enable/disable the fpA_intB path: unset/0/off to disable, other value to enable. +// This only affects nodes whose weights are NOT prepacked (see the constructor). constexpr const char* kFpAIntBGemmOption = "ORT_FPA_INTB_GEMM"; -constexpr int kFpAIntBGemmOption_All = 0x01; -constexpr int kFpAIntBGemmOption_Gemv = 0x02; -constexpr int kFpAIntBGemmOption_Int4 = 0x04; -constexpr int kFpAIntBGemmOption_Int8 = 0x08; + constexpr int64_t kMatMulNBitsWeightNotPrepacked = 0; constexpr int64_t kMatMulNBitsWeightPrepackedSm80 = 1; constexpr int64_t kMatMulNBitsWeightPrepackedSm90 = 2; @@ -52,7 +50,7 @@ constexpr int64_t kMatMulNBitsWeightPrepackedSm90 = 2; // EP: every kernel is created via KernelRegistryManager::CreateKernel, which injects the // session-level ConfigOptions, and the plugin CUDA EP wraps a CUDAExecutionProvider that reuses this // same kernel. Each key overrides its ORT_* environment-variable equivalent (config wins). -// ep.cuda.fpa_intb_gemm <-> ORT_FPA_INTB_GEMM (enable/mode) +// ep.cuda.fpa_intb_gemm <-> ORT_FPA_INTB_GEMM (0/off, 1/on) // ep.cuda.fpa_intb_profile_m <-> ORT_FPA_INTB_PROFILE_M (initial profile M buckets) constexpr const char* kConfigFpAIntBGemm = "ep.cuda.fpa_intb_gemm"; constexpr const char* kConfigFpAIntBProfileM = "ep.cuda.fpa_intb_profile_m"; @@ -68,29 +66,14 @@ inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const cha return ParseEnvironmentVariableWithDefault(env_key, ""); } -// Parses the fpA_intB enable/mode setting. Accepts (case-insensitive): "" / "0" / "off" / "false" / -// "none" -> disabled; "on" / "true" / "all" -> all; otherwise a numeric bitmask of -// kFpAIntBGemmOption_* (decimal or 0x-prefixed hex), e.g. "4" = int4 only, "6" = int4 + GEMV. -inline int ParseFpAIntBOption(const std::string& value) { - if (value.empty()) { - return 0; - } - std::string v; - v.reserve(value.size()); - for (char c : value) { - v.push_back(static_cast(std::tolower(static_cast(c)))); - } - if (v == "0" || v == "off" || v == "false" || v == "none") { - return 0; - } - if (v == "on" || v == "true" || v == "all") { - return kFpAIntBGemmOption_All; - } - try { - return static_cast(std::stol(v, nullptr, /*base*/ 0)); - } catch (const std::exception&) { - return 0; +// Parses the fpA_intB enable flag. "on" enables the full fpA_intB path (the CUTLASS GEMM and, where +// supported, the GEMV decode kernel; they share one weight layout and cannot be split). Accepts +// (case-insensitive): "" / "0" / "off" -> disabled; otherise, enabled. +inline bool ParseFpAIntBEnabled(const std::string& value) { + if (value.empty() || value == "0" || value == "off") { + return false; } + return true; } #endif @@ -155,33 +138,34 @@ class MatMulNBits final : public CudaKernel { } if constexpr (std::is_same::value || std::is_same::value) { - // Enable/mode from session config (ep.cuda.fpa_intb_gemm) with ORT_FPA_INTB_GEMM env fallback. - int option = ParseFpAIntBOption(ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBGemm, kFpAIntBGemmOption)); - ORT_ENFORCE(!(weight_prepacked_ != kMatMulNBitsWeightNotPrepacked && option == 0), - "weight_prepacked requires the fpA_intB path, but it is disabled for this node " - "(set session config ep.cuda.fpa_intb_gemm=1 or env ORT_FPA_INTB_GEMM=1)"); + const bool prepacked = weight_prepacked_ != kMatMulNBitsWeightNotPrepacked; + // The enable flag (session config ep.cuda.fpa_intb_gemm, else ORT_FPA_INTB_GEMM env) only + // chooses the path for weights that are NOT prepacked. A prepacked weight is already stored in + // the fpA_intB layout, so the choice was made at export time and cannot be turned off here. + const bool enable_fpa_intb = + prepacked || + ParseFpAIntBEnabled(ResolveFpAIntBConfigOrEnv(info, kConfigFpAIntBGemm, kFpAIntBGemmOption)); // Note: a fused bias (input[5]) is fully supported by the fpA_intB GEMV, CUTLASS SM80/SM90 // GEMM (EpilogueOpBias), and the tactic profiler, so bias-bearing nodes (e.g. gpt-oss // qkv_proj/o_proj) are eligible. Only g_idx/reorder remains unsupported by this path. - if ((option & (static_cast(nbits_) | kFpAIntBGemmOption_All)) != 0 && + if (enable_fpa_intb && (block_size_ == 32 || block_size_ == 64 || block_size_ == 128) && (nbits_ == 4 || nbits_ == 8) && !has_g_idx_ && N_ % (nbits_ == 8 ? 32 : 64) == 0 && K_ % block_size_ == 0 && sm_ >= 75) { - if ((option & (kFpAIntBGemmOption_Gemv | kFpAIntBGemmOption_All)) != 0) { - using onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; - KernelType cuda_kernel_type; - if constexpr (std::is_same::value) { - cuda_kernel_type = (nbits_ == 8) ? KernelType::FP16Int8Groupwise : KernelType::FP16Int4Groupwise; - } else if constexpr (std::is_same::value) { - cuda_kernel_type = (nbits_ == 8) ? KernelType::BF16Int8Groupwise : KernelType::BF16Int4Groupwise; - } - - if (onnxruntime::llm::kernels::fpA_intB_gemv::is_supported(sm_, cuda_kernel_type)) { - has_fpA_intB_gemv_ = true; - } + // The CUTLASS GEMM and the GEMV decode kernel consume the same fpA_intB weight layout, so + // enable GEMV whenever it is supported; a node cannot mix fpA_intB and legacy layouts. + using onnxruntime::llm::kernels::fpA_intB_gemv::KernelType; + KernelType cuda_kernel_type; + if constexpr (std::is_same::value) { + cuda_kernel_type = (nbits_ == 8) ? KernelType::FP16Int8Groupwise : KernelType::FP16Int4Groupwise; + } else if constexpr (std::is_same::value) { + cuda_kernel_type = (nbits_ == 8) ? KernelType::BF16Int8Groupwise : KernelType::BF16Int4Groupwise; + } + if (onnxruntime::llm::kernels::fpA_intB_gemv::is_supported(sm_, cuda_kernel_type)) { + has_fpA_intB_gemv_ = true; } InitGemmProfiler(FpAIntBPackingSmForKernel()); @@ -199,9 +183,10 @@ class MatMulNBits final : public CudaKernel { has_fpA_intB_gemm_ = true; } - if (weight_prepacked_ != kMatMulNBitsWeightNotPrepacked) { + if (prepacked) { ORT_ENFORCE(has_fpA_intB_gemm_, - "weight_prepacked requires the fpA_intB path, but it is disabled or unsupported for this node"); + "weight_prepacked requires the fpA_intB path, but it is unsupported for this node " + "(check bits, block_size, N/K alignment, g_idx, and compute capability >= 7.5)"); ORT_ENFORCE(weight_prepacked_ == RequiredWeightPrepackedFormat(), "weight_prepacked=", weight_prepacked_, " does not match the format required by the selected fpA_intB kernel: ", RequiredWeightPrepackedFormat()); diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 6647165cacbc7..a039501e19550 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -178,8 +178,8 @@ class TestFpAIntBConfigKeys(unittest.TestCase): These do not need the offline weight packer (pack_weights_for_cuda_mixed_gemm), so they run in more build configurations than TestMatMulNBitsPrepackedCuda. They cover: the config key enabling - the fpA_intB path, session config overriding the ORT_FPA_INTB_GEMM env var, the profile-M key - being accepted, and the prepacked-requires-enable error message referencing the config key. + the fpA_intB path (on/off only), session config overriding the ORT_FPA_INTB_GEMM env var, the + profile-M key being accepted, and env-var backward compatibility. """ def setUp(self): @@ -243,10 +243,10 @@ def _make_int4_case(self, m=32, k=256, n=512, block_size=64): def test_config_key_enables_fpa_intb(self): # Baseline runs the standard dequant path (fpA_intB disabled); the config key must switch to - # the fpA_intB path and stay numerically equivalent. + # the fpA_intB path and stay numerically equivalent. Only on/off is accepted. model, a, _, _ = self._make_int4_case() ref = self._run(model, a) - for value in ("1", "all", "on", "0x4"): + for value in ("1", "on", "all", "true"): out = self._run(model, a, {"ep.cuda.fpa_intb_gemm": value}) np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2, err_msg=f"value={value}") @@ -267,19 +267,11 @@ def test_session_config_overrides_env(self): def test_env_var_backward_compatible(self): model, a, _, _ = self._make_int4_case() ref = self._run(model, a) - with set_env("ORT_FPA_INTB_GEMM", "1"): - out = self._run(model, a) - np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2) - - def test_prepacked_requires_enable_reports_config_key(self): - # A prepacked node requires the fpA_intB path; when it is disabled the construction-time - # error must point at the session config key (and the env var). - model, a, q_weight, scales = self._make_int4_case() - prepacked = self._make_model(a.shape[0], a.shape[1], q_weight.shape[0], q_weight, scales, 4, 64, - weight_prepacked=1) - with self.assertRaises(Exception) as ctx: - self._run(prepacked, a) # nothing enables the path - self.assertIn("ep.cuda.fpa_intb_gemm", str(ctx.exception)) + # "1" plus a legacy non-zero numeric value (previously a bitmask) both mean "enabled" now. + for value in ("1", "4"): + with set_env("ORT_FPA_INTB_GEMM", value): + out = self._run(model, a) + np.testing.assert_allclose(out, ref, rtol=2e-2, atol=2e-2, err_msg=f"env={value}") if __name__ == "__main__": From 0a6c50188a4cfc8a77786abd6656ffa99a734122 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 00:36:01 +0000 Subject: [PATCH 06/12] Fix fpA_intB MatMulNBits CI failures and address review feedback - 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. --- docs/contrib_ops/cuda/matmul_nbits.md | 5 ++++- .../cuda/llm/fpA_intB_gemm_profiler.cc | 20 +++++++++---------- .../cuda/quantization/matmul_nbits.cc | 6 +++--- .../cuda/quantization/matmul_nbits.h | 10 ++++++++-- .../test/contrib_ops/matmul_4bits_test.cc | 11 +++++++--- 5 files changed, 33 insertions(+), 19 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index d12d8400b6e71..6e382826de37d 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -263,7 +263,10 @@ Prepacked weights are intentionally strict: - If ORT was built without `onnxruntime_USE_FPA_INTB_GEMM=ON`, any nonzero `weight_prepacked` value throws during kernel construction. -- If any nonzero `weight_prepacked` value will cause `ORT_FPA_INTB_GEMM` be ignored. +- Any nonzero `weight_prepacked` value forces the fpA_intB path on, so the enable + flag (`ep.cuda.fpa_intb_gemm` session config, or the `ORT_FPA_INTB_GEMM` env + var) is ignored for prepacked weights — the layout choice was fixed at export + time and cannot be turned off at run time. - Nonzero `weight_prepacked` requires FP16 or BF16 input `A`, because only the CUDA fpA_intB path consumes this layout. - `weight_prepacked` must match the layout the selected kernel expects: `1` is diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc index f102587a00197..a56b8ea632007 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -62,7 +62,7 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic( onnxruntime::llm::kernels::fpA_intB_gemv::kernel_launcher(mArch, params, stream); } else { // run CUTLASS kernel - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); + int const wsSize = static_cast(mRunner->getWorkspaceSize(m, originalN, k)); if (mQuantBits == 8) { mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, outputPtr, m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); @@ -75,17 +75,17 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic( size_t WeightOnlyGroupwiseQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) { // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantBits == 8 ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; + int const originalN = static_cast(mQuantBits == 8 ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO); std::vector workspaces = { - maxM * k * sizeof(half), // A - k * n * sizeof(half), // B - k * originalN * sizeof(half) / mGroupSize, // scales - k * originalN * sizeof(half) / mGroupSize, // zeros - originalN * sizeof(half), // biases - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace + maxM * k * sizeof(half), // A + k * n * sizeof(half), // B + k * originalN * sizeof(half) / mGroupSize, // scales + k * originalN * sizeof(half) / mGroupSize, // zeros + originalN * sizeof(half), // biases + maxM * originalN * sizeof(half), // C + mRunner->getWorkspaceSize(static_cast(maxM), originalN, static_cast(k)) // workspace }; - return calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); + return calculateTotalWorkspaceSize(workspaces.data(), static_cast(workspaces.size())); } std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics( diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 3e52aeee2afb8..34a0f668c1f88 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -384,9 +384,9 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { // and synchronizes events, and allocates/frees scratch, all of which are illegal while the // compute stream is being captured. Fall back to a lookup of an already-profiled bucket // (warmup runs before capture populate these); only outside capture do we allow lazy - // single-bucket profiling. - const bool stream_is_capturing = - stream != nullptr && onnxruntime::llm::common::isCapturing(stream); + // single-bucket profiling. Note: a null cudaStream_t is the default stream (a valid capture + // target under per-thread default streams), so query the capture status unconditionally. + const bool stream_is_capturing = onnxruntime::llm::common::isCapturing(stream); auto const bestTactic = stream_is_capturing ? gemmProfiler_->getBestConfig(m, gemmId_) : gemmProfiler_->getBestConfigOrProfile(m, gemmId_); if (!bestTactic.has_value()) { diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 49840cb0dae26..b0eca2466426b 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -68,9 +68,15 @@ inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const cha // Parses the fpA_intB enable flag. "on" enables the full fpA_intB path (the CUTLASS GEMM and, where // supported, the GEMV decode kernel; they share one weight layout and cannot be split). Accepts -// (case-insensitive): "" / "0" / "off" -> disabled; otherise, enabled. +// (case-insensitive): "" / "0" / "off" -> disabled; otherwise, enabled (a non-zero numeric value +// still enables, for backward compatibility). inline bool ParseFpAIntBEnabled(const std::string& value) { - if (value.empty() || value == "0" || value == "off") { + std::string lowered; + lowered.reserve(value.size()); + for (char c : value) { + lowered.push_back(static_cast(std::tolower(static_cast(c)))); + } + if (lowered.empty() || lowered == "0" || lowered == "off") { return false; } return true; diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 51895938ce3aa..2927477a66244 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -984,8 +984,13 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint_Bias_Prepacked) { RunTest(opts, std::move(eps)); } -TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRequiresFpAIntBGemm) { - ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_FPA_INTB_GEMM", "0"}}}; +// A prepacked weight (weight_prepacked!=0) forces the fpA_intB path on regardless of the enable +// flag, and the constructor ORT_ENFORCEs that the path is actually supported for the node. Here the +// block_size (256) is outside the fpA_intB-supported set {32, 64, 128}, so kernel construction is +// rejected up front with "weight_prepacked requires the fpA_intB path, but it is unsupported ...", +// even though ORT_FPA_INTB_GEMM is enabled. +TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRejectedWhenFpAIntBUnsupported) { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_FPA_INTB_GEMM", "1"}}}; auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) { @@ -994,7 +999,7 @@ TEST(MatMulNBits, Fp16_Int4_PrepackedWeightRequiresFpAIntBGemm) { TestOptions opts{}; opts.M = 1, opts.N = 256, opts.K = 1024; - opts.block_size = 64; + opts.block_size = 256; opts.disable_cpu_ep_fallback = true; opts.weight_prepacked = 1; opts.expected_failure = "weight_prepacked requires"; From 06eb3ae878862491f98ff8186d066be6a289a285 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 00:38:47 +0000 Subject: [PATCH 07/12] Clarify fpA_intB config-toggle test comment (numeric-equivalence scope) --- .../quantization/test_op_matmulnbits_prepacked_cuda.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py index 28f26a3fd3ad8..0fcc6eacd9dc4 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -281,8 +281,12 @@ def _make_int4_case(self, m=32, k=256, n=512, block_size=64): return model, a, q_weight, scales def test_config_key_enables_fpa_intb(self): - # Baseline runs the standard dequant path (fpA_intB disabled); the config key must switch to - # the fpA_intB path and stay numerically equivalent. Only on/off is accepted. + # On fpA_intB-capable hardware (compute capability >= 7.5) the baseline (no config) runs the + # standard dequant path -- for a non-prepacked node the enable flag defaults to disabled -- + # while the config key selects the fpA_intB path; the two paths must stay numerically + # equivalent. On sm < 75 both fall back to the dequant path, so this asserts equivalence + # rather than the switch itself (the prepacked tests force and exercise the fpA_intB kernel). + # Only on/off is accepted. model, a, _, _ = self._make_int4_case() ref = self._run(model, a) for value in ("1", "on", "all", "true"): From 7e0065fb544a8536e1343afda3a84988cc8758f9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 06:38:34 +0000 Subject: [PATCH 08/12] refactoring --- .../cuda/llm/fpA_intB_gemm_profiler.cc | 37 ++++----- .../contrib_ops/cuda/llm/gemm_profiler.h | 77 +++++++++---------- .../cuda/quantization/matmul_nbits.h | 14 ++-- 3 files changed, 57 insertions(+), 71 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc index a56b8ea632007..2b3d63c0585d3 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -20,7 +20,9 @@ #include #include -#include + +#include "core/common/parse_string.h" +#include "core/common/string_utils.h" using namespace onnxruntime::llm::common; using namespace onnxruntime::llm::kernels::cutlass_kernels; @@ -77,12 +79,12 @@ size_t WeightOnlyGroupwiseQuantGemmPluginProfiler::computeTmpSize(size_t maxM, s // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) int const originalN = static_cast(mQuantBits == 8 ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO); std::vector workspaces = { - maxM * k * sizeof(half), // A - k * n * sizeof(half), // B - k * originalN * sizeof(half) / mGroupSize, // scales - k * originalN * sizeof(half) / mGroupSize, // zeros - originalN * sizeof(half), // biases - maxM * originalN * sizeof(half), // C + /* A */ maxM * k * sizeof(half), + /* B */ k * n * sizeof(half), + /* scales */ k * originalN * sizeof(half) / mGroupSize, + /* zeros */ k * originalN * sizeof(half) / mGroupSize, + /* biases */ originalN * sizeof(half), + /* C */ maxM * originalN * sizeof(half), mRunner->getWorkspaceSize(static_cast(maxM), originalN, static_cast(k)) // workspace }; return calculateTotalWorkspaceSize(workspaces.data(), static_cast(workspaces.size())); @@ -106,24 +108,15 @@ std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(c if (value.empty()) { return result; } - std::stringstream ss(value); - std::string token; std::set unique; - while (std::getline(ss, token, ',')) { - // Trim surrounding whitespace. - size_t start = token.find_first_not_of(" \t"); - size_t end = token.find_last_not_of(" \t"); - if (start == std::string::npos) { + for (const auto token : onnxruntime::utils::SplitString(value, ",", true)) { + const std::string trimmed_token = onnxruntime::utils::TrimString(token); + if (trimmed_token.empty()) { continue; } - token = token.substr(start, end - start + 1); - try { - int m = std::stoi(token); - if (m > 0) { - unique.insert(m); - } - } catch (const std::exception&) { - // Ignore malformed entries. + int m = 0; + if (TryParseStringWithClassicLocale(trimmed_token, m) && m > 0) { + unique.insert(m); } } result.assign(unique.begin(), unique.end()); diff --git a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h index 291d8ac87e791..e695ae1af4914 100644 --- a/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h +++ b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h @@ -196,9 +196,10 @@ class GemmPluginProfiler { virtual std::vector getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const; private: - std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics); + std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics, + char* workspace, cudaStream_t stream); - float profileTacticForProblem(int m, int n, int k, Config const& tactic); + float profileTacticForProblem(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t stream); int nextPowerOfTwo(int v) const { --v; @@ -218,10 +219,6 @@ class GemmPluginProfiler { private: MNKProfileMapPtr mMNKProfileMap{}; - onnxruntime::IAllocatorUniquePtr mWorkspaceTmp{nullptr}; - - cudaStream_t mStream; - GemmDims mDims{}; bool mSkip{false}; @@ -311,11 +308,13 @@ void GemmPluginProfiler::profileT auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); bool isAllocated{false}; + onnxruntime::IAllocatorUniquePtr workspace_tmp{nullptr}; + cudaStream_t stream; auto profileTactics = [&](int m, int n, int k) { if (mProfileMap->count(m) == 0) { if (!isAllocated) { - this->mWorkspaceTmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); + workspace_tmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); #if ORT_LLM_VERBOSE AllocatorStats stats; this->mAllocator->GetStats(&stats); @@ -325,15 +324,15 @@ void GemmPluginProfiler::profileT isAllocated = true; } - initTmpData(m, n, k, this->mWorkspaceTmp.get(), workspace_bytes, this->mStream); + initTmpData(m, n, k, workspace_tmp.get(), workspace_bytes, stream); auto tactics = this->getTactics(m, n, k); // Profile different tactics for particular m and insert best config to the map - mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics)}); + mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics, workspace_tmp.get(), stream)}); } }; - CUDA_CALL_THROW(cudaStreamCreate(&mStream)); + CUDA_CALL_THROW(cudaStreamCreate(&stream)); // Profile the (possibly reduced) set of M buckets. Any unprofiled runtime M is handled // later by lazy single-bucket profiling in getBestConfigOrProfile. @@ -343,9 +342,9 @@ void GemmPluginProfiler::profileT if (isAllocated) { // Free tmp data - mWorkspaceTmp.reset(); + workspace_tmp.reset(); } - CUDA_CALL_THROW(cudaStreamDestroy(mStream)); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); } template @@ -419,14 +418,33 @@ std::optional GemmPluginProfilermutex); + // We can only profile lazily if the profiling context from construction is available. + if (mRunner == nullptr || mAllocator == nullptr || !mDims.isInitialized()) { + ORT_LLM_LOG_WARNING("Cannot lazily profile an unprofiled M bucket: profiler context is unavailable."); + return std::nullopt; + } + + int const n = static_cast(mDims.n); + int const k = static_cast(mDims.k); + size_t const workspace_bytes = computeTmpSize(target, n, k); + + cudaStream_t stream; + CUDA_CALL_THROW(cudaStreamCreate(&stream)); + auto workspace_tmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); + initTmpData(target, n, k, workspace_tmp.get(), workspace_bytes, stream); + + auto tactics = this->getTactics(target, n, k); + auto best = this->profileTacticsForProblem(target, n, k, tactics, workspace_tmp.get(), stream); + + workspace_tmp.reset(); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); + writer_lock lock(mMNKProfileMap->mutex); if (!mMNKProfileMap->existsMProfileMap(gemmId)) { mMNKProfileMap->createMProfileMap(gemmId); } auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); - // Re-check under the writer lock: another thread may have profiled it meanwhile. if (mProfileMap->count(m) > 0) { return mProfileMap->at(m); } @@ -434,33 +452,14 @@ std::optional GemmPluginProfilerat(target); } - // We can only profile lazily if the profiling context from construction is available. - if (mRunner == nullptr || mAllocator == nullptr || !mDims.isInitialized()) { - ORT_LLM_LOG_WARNING("Cannot lazily profile an unprofiled M bucket: profiler context is unavailable."); - return std::nullopt; - } - - int const n = static_cast(mDims.n); - int const k = static_cast(mDims.k); - size_t const workspace_bytes = computeTmpSize(target, n, k); - - CUDA_CALL_THROW(cudaStreamCreate(&mStream)); - mWorkspaceTmp = onnxruntime::IAllocator::MakeUniquePtr(mAllocator, workspace_bytes, true); - initTmpData(target, n, k, mWorkspaceTmp.get(), workspace_bytes, mStream); - - auto tactics = this->getTactics(target, n, k); - auto best = this->profileTacticsForProblem(target, n, k, tactics); mProfileMap->insert({target, best}); - mWorkspaceTmp.reset(); - CUDA_CALL_THROW(cudaStreamDestroy(mStream)); - return best; } template std::optional GemmPluginProfiler::profileTacticsForProblem( - int m, int n, int k, std::vector const& tactics) { + int m, int n, int k, std::vector const& tactics, char* workspace, cudaStream_t stream) { ORT_LLM_LOG_ENTRY(); float bestTime = std::numeric_limits::max(); @@ -480,7 +479,7 @@ std::optional GemmPluginProfiler 1 if constexpr (std::is_same_v) { @@ -527,15 +526,13 @@ std::optional GemmPluginProfiler float GemmPluginProfiler::profileTacticForProblem( - int m, int n, int k, Config const& tactic) { + int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t stream) { constexpr int warmup = 5; constexpr int runs = 10; - cudaStream_t stream = mStream; - // Warmup the execution for (int i = 0; i < warmup; ++i) { - runTactic(m, n, k, tactic, mWorkspaceTmp.get(), stream); + runTactic(m, n, k, tactic, workspace, stream); } cudaEvent_t start; @@ -547,7 +544,7 @@ float GemmPluginProfiler::profile // Profile GEMM for (int i = 0; i < runs; ++i) { - runTactic(m, n, k, tactic, mWorkspaceTmp.get(), stream); + runTactic(m, n, k, tactic, workspace, stream); } CUDA_CALL_THROW(cudaEventRecord(stop, stream)); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index b0eca2466426b..30c5edaa82c15 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -8,12 +8,12 @@ // #pragma once #include "core/common/safeint.h" +#include "core/common/string_utils.h" #include "core/providers/cuda/cuda_kernel.h" #include "core/providers/cuda/shared_inc/fpgeneric.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "core/platform/env_var_utils.h" -#include #include #include @@ -59,9 +59,9 @@ constexpr const char* kConfigFpAIntBProfileM = "ep.cuda.fpa_intb_profile_m"; // variable, else empty. Session config wins so a model/session can override a process-wide env var. inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const char* config_key, const char* env_key) { - const std::string from_config = info.GetConfigOptions().GetConfigOrDefault(config_key, ""); - if (!from_config.empty()) { - return from_config; + const auto from_config = info.GetConfigOptions().GetConfigEntry(config_key); + if (from_config.has_value()) { + return *from_config; } return ParseEnvironmentVariableWithDefault(env_key, ""); } @@ -71,11 +71,7 @@ inline std::string ResolveFpAIntBConfigOrEnv(const OpKernelInfo& info, const cha // (case-insensitive): "" / "0" / "off" -> disabled; otherwise, enabled (a non-zero numeric value // still enables, for backward compatibility). inline bool ParseFpAIntBEnabled(const std::string& value) { - std::string lowered; - lowered.reserve(value.size()); - for (char c : value) { - lowered.push_back(static_cast(std::tolower(static_cast(c)))); - } + const std::string lowered = onnxruntime::utils::GetLowercaseString(onnxruntime::utils::TrimString(value)); if (lowered.empty() || lowered == "0" || lowered == "off") { return false; } From 6d2d77640b139f33707c9db7335dc97466e94232 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 18:19:48 +0000 Subject: [PATCH 09/12] fix(cuda): use ORT env helper instead of std::getenv in cuda_runtime_utils std::getenv triggers MSVC C4996 (treated as error) now that fpA_intB GEMM is compiled in Windows CUDA builds. Replace with the cross-platform ParseEnvironmentVariableWithDefault helper, matching sibling env_utils.h. --- .../contrib_ops/cuda/llm/common/cuda_runtime_utils.h | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index 8b9e35adc106e..718e31d150173 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -21,6 +21,7 @@ #ifdef ENABLE_FP8 #include #endif +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/shared_inc/cuda_call.h" namespace onnxruntime::llm::common { @@ -61,12 +62,7 @@ inline std::optional isCudaLaunchBlocking() { thread_local bool firstCall = true; thread_local std::optional result = std::nullopt; if (!firstCall) { - char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); - if (env != nullptr && std::string(env) == "1") { - result = true; - } else { - result = false; - } + result = ParseEnvironmentVariableWithDefault("CUDA_LAUNCH_BLOCKING", 0) == 1; firstCall = false; } return result; From 050b9a95b079ce2ffa41004aeaab12e2bda2bf4c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 21:13:26 +0000 Subject: [PATCH 10/12] fix Windows CI --- .../cuda/llm/common/cuda_runtime_utils.h | 16 ++++++++++++++-- .../cuda/quantization/matmul_nbits.cc | 4 ++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index 718e31d150173..5926939375a53 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -16,12 +16,13 @@ */ #pragma once +#include #include +#include #include #ifdef ENABLE_FP8 #include #endif -#include "core/platform/env_var_utils.h" #include "core/providers/cuda/shared_inc/cuda_call.h" namespace onnxruntime::llm::common { @@ -62,7 +63,18 @@ inline std::optional isCudaLaunchBlocking() { thread_local bool firstCall = true; thread_local std::optional result = std::nullopt; if (!firstCall) { - result = ParseEnvironmentVariableWithDefault("CUDA_LAUNCH_BLOCKING", 0) == 1; + // Read the env var directly here instead of via core/platform/env_var_utils.h. + // This is a leaf CUDA header pulled in by kernel headers (e.g. compute_occupancy.h) + // BEFORE the SHARED_PROVIDER bridge (provider_api.h) is established. env_var_utils.h + // unconditionally includes core/common/logging/logging.h while SHARED_PROVIDER is + // undefined, which then clashes with the logging stubs provider_api.h defines later in + // the same translation unit. Using std::getenv keeps this header self-contained. + char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); + if (env != nullptr && std::string(env) == "1") { + result = true; + } else { + result = false; + } firstCall = false; } return result; diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 34a0f668c1f88..9bd6f49b5757a 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -640,7 +640,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { // Chunked dequant+GEMM: scratch buffer is already sized for one chunk. const int64_t chunk_n = chunk_target_rows; - auto* out_data = reinterpret_cast(Y->MutableData()); + auto* chunk_out_data = reinterpret_cast(Y->MutableData()); for (int64_t n_start = 0; n_start < N_; n_start += chunk_n) { const int64_t n_end = std::min(n_start + chunk_n, N_); @@ -686,7 +686,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { reinterpret_cast(a_data), // A [M, K] helper.Lda(transa), &zero, - out_data + n_start, // C[:, n_start] — strided output + chunk_out_data + n_start, // C[:, n_start] — strided output SafeInt(helper.Ldc()), // ldc = N (full output stride) GetDeviceProp(), UseTF32())); From 0360f4e40898e74e0b2531446aa5da5f0707b3ee Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 21:20:46 +0000 Subject: [PATCH 11/12] fix(cuda): use _dupenv_s on Windows to avoid C4996 in cuda_runtime_utils --- .../cuda/llm/common/cuda_runtime_utils.h | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index 5926939375a53..7ae61835a8540 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -68,13 +68,22 @@ inline std::optional isCudaLaunchBlocking() { // BEFORE the SHARED_PROVIDER bridge (provider_api.h) is established. env_var_utils.h // unconditionally includes core/common/logging/logging.h while SHARED_PROVIDER is // undefined, which then clashes with the logging stubs provider_api.h defines later in - // the same translation unit. Using std::getenv keeps this header self-contained. - char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); - if (env != nullptr && std::string(env) == "1") { - result = true; + // the same translation unit. Read the variable directly to keep this header self-contained. + // std::getenv is avoided on MSVC because it raises C4996 (treated as an error); _dupenv_s + // is the supported Windows-safe replacement. +#if defined(_WIN32) + char* env = nullptr; + size_t env_len = 0; + if (_dupenv_s(&env, &env_len, "CUDA_LAUNCH_BLOCKING") == 0 && env != nullptr) { + result = std::string(env) == "1"; } else { result = false; } + std::free(env); +#else + char const* env = std::getenv("CUDA_LAUNCH_BLOCKING"); + result = (env != nullptr && std::string(env) == "1"); +#endif firstCall = false; } return result; From 9079f41d87312831c5407c7d6e68c42a6bc12e16 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 9 Jul 2026 23:17:30 +0000 Subject: [PATCH 12/12] fix build --- .../cuda/quantization/matmul_nbits.cc | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index 9bd6f49b5757a..2630d76b6e4dc 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -110,8 +110,8 @@ void MatMulNBits::InitGemmProfiler(int sm) { } gemmProfiler_->setCudaKernelType(cuda_kernel_type, sm); - gemmProfiler_->setQuant(nbits_, has_bias_, has_zero_points_); - gemmProfiler_->setGroupSize(block_size_); + gemmProfiler_->setQuant(static_cast(nbits_), has_bias_, has_zero_points_); + gemmProfiler_->setGroupSize(static_cast(block_size_)); auto allocator = this->Info().GetAllocator(OrtMemType::OrtMemTypeDefault); gemmProfiler_->setAllocator(allocator); @@ -120,15 +120,15 @@ void MatMulNBits::InitGemmProfiler(int sm) { template void MatMulNBits::RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int max_m) { // Number of 16-bit elements after casting int8/int4 to fp16. - int n_16b = N_ / (nbits_ == 8 ? 2 : 4); + int n_16b = static_cast(N_ / (nbits_ == 8 ? 2 : 4)); // Include the packing/kernel SM in the GEMM id so the SM80-compatibility and native SM90 kernels // (which need different tactics) do not share profiled configs for the same (N, K, dtype). const int kernel_sm = FpAIntBPackingSmForKernel(); if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, K_, onnxruntime::llm::nvinfer::DataType::kHALF, kernel_sm); + gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kHALF, kernel_sm); } else if constexpr (std::is_same_v) { - gemmId_ = GemmIdCore(n_16b, K_, onnxruntime::llm::nvinfer::DataType::kBF16, kernel_sm); + gemmId_ = GemmIdCore(n_16b, static_cast(K_), onnxruntime::llm::nvinfer::DataType::kBF16, kernel_sm); } GemmDims dims = {min_m, max_m, n_16b, K_}; @@ -206,10 +206,10 @@ Status MatMulNBits::PrePack_B([[maybe_unused]] const Tensor& tensor, if (nbits_ == 4) { // Transpose the weight and add default zero point. onnxruntime::llm::kernels::fpA_intB_gemv::unpack_uint4_transposed_to_int8_direct_cuda( - stream, packed_transposed_weight, blob_data, n, k); + stream, packed_transposed_weight, blob_data, static_cast(n), static_cast(k)); } else { onnxruntime::llm::kernels::fpA_intB_gemv::transpose_uint8_matrix_and_convert_to_int8( - stream, packed_transposed_weight, blob_data, n, k); + stream, packed_transposed_weight, blob_data, static_cast(n), static_cast(k)); } using onnxruntime::llm::kernels::weight_only::QuantType; @@ -252,7 +252,7 @@ Status MatMulNBits::PrePack_Scale([[maybe_unused]] const Tensor& tensor, CudaT* transposed_scales = reinterpret_cast(fpA_intB_scale_buffer_.get()); onnxruntime::llm::kernels::fpA_intB_gemv::launch_transpose_scale_kernel( - stream, reinterpret_cast(tensor.Data()), transposed_scales, n, k_blocks); + stream, reinterpret_cast(tensor.Data()), transposed_scales, static_cast(n), static_cast(k_blocks)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); DUMP_TENSOR_INIT(); @@ -288,16 +288,16 @@ Status MatMulNBits::PrePack_ZeroPoint([[maybe_unused]] const Tensor& tensor, if (nbits_ == 4) { onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } else { onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } } else { // zero point is not uint8_t type onnxruntime::llm::kernels::fpA_intB_gemv::launch_scaled_zero_point_kernel( stream, reinterpret_cast(zero_points_data), - transposed_scales, scaled_zero_points, n, k_blocks, default_zero_point); + transposed_scales, scaled_zero_points, static_cast(n), static_cast(k_blocks), default_zero_point); } CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream)); @@ -441,7 +441,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { a_data, pre_quant_scale_ptr, fpA_intB_weight, fpA_intB_scale_buffer_.get(), has_zero_points_ ? fpA_intB_zero_buffer_.get() : nullptr, bias_data, out_data, - alpha, m, n, k, block_size_, cuda_kernel_type, apply_alpha_in_advance); + alpha, m, n, k, static_cast(block_size_), cuda_kernel_type, apply_alpha_in_advance); // Launch the GEMV with the arch the weights were PACKED for (FpAIntBPackingSmForKernel), // not the raw device SM. The GEMV interleave layout is arch-dependent: arch in [90,100) @@ -463,7 +463,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { 1.f, out_data, m, n, k, - block_size_, + static_cast(block_size_), *bestTactic, reinterpret_cast(workspace_buffer.get()), workspace_size,