diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index ad446214cfc8f..6b8bdd951d5ad 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -108,7 +108,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/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index 3f8e1f1d8f616..6e382826de37d 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -263,8 +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 `ORT_FPA_INTB_GEMM` is unset or `0`, any nonzero `weight_prepacked` value - throws instead of silently falling back to a raw-layout path. +- 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 @@ -293,7 +295,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/llm/common/cuda_runtime_utils.h b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h index 8b9e35adc106e..7ae61835a8540 100644 --- a/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h +++ b/onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h @@ -16,7 +16,9 @@ */ #pragma once +#include #include +#include #include #ifdef ENABLE_FP8 #include @@ -61,12 +63,27 @@ 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; + // 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. 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; 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..2b3d63c0585d3 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc @@ -18,6 +18,12 @@ #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h" #include "contrib_ops/cuda/llm/common/workspace.h" +#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; @@ -58,7 +64,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); @@ -71,17 +77,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 + /* 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(), workspaces.size()); + return calculateTotalWorkspaceSize(workspaces.data(), static_cast(workspaces.size())); } std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics( @@ -97,5 +103,53 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int /*n*/, i return true; } +std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(const std::string& value) { + std::vector result; + if (value.empty()) { + return result; + } + std::set unique; + for (const auto token : onnxruntime::utils::SplitString(value, ",", true)) { + const std::string trimmed_token = onnxruntime::utils::TrimString(token); + if (trimmed_token.empty()) { + continue; + } + int m = 0; + if (TryParseStringWithClassicLocale(trimmed_token, m) && m > 0) { + unique.insert(m); + } + } + result.assign(unique.begin(), unique.end()); + return result; +} + +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; + + if (!mProfileMOverride.empty()) { + for (int m : mProfileMOverride) { + 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..be6f1a2d82d1c 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,31 @@ 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 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; mHasBiases = has_bias; @@ -74,6 +93,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; @@ -81,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/llm/gemm_profiler.h b/onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h index d9c19cb734e91..e695ae1af4914 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,10 +190,16 @@ 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); + 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; @@ -206,14 +219,14 @@ class GemmPluginProfiler { private: MNKProfileMapPtr mMNKProfileMap{}; - onnxruntime::IAllocatorUniquePtr mWorkspaceTmp{nullptr}; - - cudaStream_t mStream; - GemmDims mDims{}; 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 +290,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); @@ -293,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); @@ -307,43 +324,27 @@ 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)); - - int const startMinMRounded = nextPowerOfTwo(dims.minM); + CUDA_CALL_THROW(cudaStreamCreate(&stream)); - 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(); + workspace_tmp.reset(); } - CUDA_CALL_THROW(cudaStreamDestroy(mStream)); + CUDA_CALL_THROW(cudaStreamDestroy(stream)); } template @@ -372,9 +373,93 @@ 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); + } + } + } + + // 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); + + if (mProfileMap->count(m) > 0) { + return mProfileMap->at(m); + } + if (mProfileMap->count(target) > 0) { + return mProfileMap->at(target); + } + + mProfileMap->insert({target, best}); + + 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(); @@ -394,7 +479,7 @@ std::optional GemmPluginProfiler 1 if constexpr (std::is_same_v) { @@ -441,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; @@ -461,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.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index abd43dc6f1f6c..2630d76b6e4dc 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" @@ -109,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); @@ -119,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_}; @@ -205,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; @@ -251,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(); @@ -287,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)); @@ -371,18 +372,30 @@ 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."); 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. 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()) { - 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 @@ -428,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) @@ -450,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, @@ -627,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_); @@ -673,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())); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h index 9a7bfe2895582..30c5edaa82c15 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h @@ -8,11 +8,15 @@ // #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 + namespace onnxruntime { namespace contrib { namespace cuda { @@ -34,15 +38,45 @@ 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; + +// 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 (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"; + +// 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 auto from_config = info.GetConfigOptions().GetConfigEntry(config_key); + if (from_config.has_value()) { + return *from_config; + } + return ParseEnvironmentVariableWithDefault(env_key, ""); +} + +// 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; otherwise, enabled (a non-zero numeric value +// still enables, for backward compatibility). +inline bool ParseFpAIntBEnabled(const std::string& value) { + const std::string lowered = onnxruntime::utils::GetLowercaseString(onnxruntime::utils::TrimString(value)); + if (lowered.empty() || lowered == "0" || lowered == "off") { + return false; + } + return true; +} #endif template @@ -106,43 +140,55 @@ class MatMulNBits final : public CudaKernel { } if constexpr (std::is_same::value || std::is_same::value) { - int option = ParseEnvironmentVariableWithDefault(kFpAIntBGemmOption, 0); - ORT_ENFORCE(!(weight_prepacked_ != kMatMulNBitsWeightNotPrepacked && option == 0), - "weight_prepacked requires the fpA_intB path, but ORT_FPA_INTB_GEMM is off for this node"); + 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()); - constexpr int max_m = 8291; + // 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; } - 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/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"; 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 9d3a378f6db19..0fcc6eacd9dc4 100644 --- a/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py +++ b/onnxruntime/test/python/quantization/test_op_matmulnbits_prepacked_cuda.py @@ -177,7 +177,6 @@ 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(_cuda_quant is not None, "standalone CUDA weight packer (parity oracle) is unavailable") class TestCudaQuantizerTorchPackerParity(unittest.TestCase): """Validate the PyTorch mixed-GEMM packer in cuda_quantizer.py against the CUDA oracle. @@ -211,5 +210,112 @@ def test_torch_packer_matches_cuda_oracle(self): self._check(bits, force_arch, n, k) +@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 (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): + # 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): + # 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"): + 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) + # "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__": unittest.main() 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"