Enable fpA_intB GEMM in CUDA builds and add configurable options

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)

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. |
 

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc

@@ -18,6 +18,10 @@
 #include "contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h"
 #include "contrib_ops/cuda/llm/common/workspace.h"
 
+#include <algorithm>
+#include <set>
+#include <sstream>
+
 using namespace onnxruntime::llm::common;
 using namespace onnxruntime::llm::kernels::cutlass_kernels;
 
@@ -58,7 +62,7 @@
     onnxruntime::llm::kernels::fpA_intB_gemv::kernel_launcher(mArch, params, stream);
   } else {
     // run CUTLASS kernel
     int const wsSize = mRunner->getWorkspaceSize(m, originalN, k);
     if (mQuantBits == 8) {
       mRunner->gemm(actPtr, reinterpret_cast<int8_t*>(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, outputPtr,
                     m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream);
@@ -71,7 +75,7 @@
 
 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;
   std::vector<size_t> workspaces = {
       maxM * k * sizeof(half),                       // A
       k * n * sizeof(half),                          // B
@@ -97,5 +101,62 @@
   return true;
 }
 
+std::vector<int> WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(const std::string& value) {
+  std::vector<int> result;
+  if (value.empty()) {
+    return result;
+  }
+  std::stringstream ss(value);
+  std::string token;
+  std::set<int> 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;
+}
+
+std::vector<int> 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<int> 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<int>(buckets.begin(), buckets.end());
+}
+
 }  // namespace onnxruntime::llm::kernels::weight_only
 #endif

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h

@@ -43,12 +43,31 @@
 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<onnxruntime::llm::cutlass_extensions::CutlassGemmConfig, WeightOnlyGemmRunnerPtr,
                                 GemmIdCore, GemmIdCoreHash> {
  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<int> 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<int> ms) {
+    mProfileMOverride = std::move(ms);
+  }
+
   void setQuant(int bits, bool has_bias, bool has_zeros) {
     mQuantBits = bits;
     mHasBiases = has_bias;
@@ -74,13 +93,16 @@
 
   bool checkTactic(int m, int n, int k, Config const& tactic) const override;
 
+  std::vector<int> getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const override;
+
  private:
   bool mHasBiases;
   bool mHasZeros;
   int mQuantBits;
   int mGroupSize;
   KernelType mCudaKernelType;
   int mArch;
+  std::vector<int> mProfileMOverride;
 };
 
 }  // namespace onnxruntime::llm::kernels::weight_only

onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h

@@ -110,10 +110,10 @@
   using MProfileMap = std::unordered_map<int, std::optional<Config>>;
   using MProfileMapPtr = std::shared_ptr<MProfileMap>;
 
+  // requires shared ownership to read from *this
+  using reader_lock = std::shared_lock<std::shared_timed_mutex>;
   // requires exclusive ownership to write to *this
-  using reader_lock = std::unique_lock<std::shared_timed_mutex>;
-  // requires shared ownership to read from other
-  using writer_lock = std::shared_lock<std::shared_timed_mutex>;
+  using writer_lock = std::unique_lock<std::shared_timed_mutex>;
 
   // Struct of continuing map if GEMMs to the best profiles for different Ms
   struct MNKProfileMap {
@@ -168,6 +168,13 @@
 
   std::optional<Config> 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<Config> getBestConfigOrProfile(int m, GemmIdType const& gemmId);
+
   virtual int getMaxProfileM() const;
 
  protected:
@@ -183,6 +190,11 @@
 
   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<int> getProfileMBuckets(int minM, int maxM, bool hasWeightOnlyCudaKernel) const;
+
  private:
   std::optional<Config> profileTacticsForProblem(int m, int n, int k, std::vector<Config> const& tactics);
 
@@ -214,6 +226,10 @@
 
   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 @@
 
   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<int>(dims.maxM)), getMaxProfileM());
 
   size_t workspace_bytes = computeTmpSize(maxM, dims.n, dims.k);
 
@@ -317,28 +335,12 @@
 
   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<int>(dims.minM), maxM, hasWeightOnlyCudaKernel)) {
+    profileTactics(m, static_cast<int>(dims.n), static_cast<int>(dims.k));
   }
 
-  profileTactics(maxM, dims.n, dims.k);
-
   if (isAllocated) {
     // Free tmp data
     mWorkspaceTmp.reset();
@@ -372,6 +374,90 @@
   }
 }
 
+template <typename Config, typename RunnerPtr, typename GemmIdType, typename GemmIdHashType>
+std::vector<int> GemmPluginProfiler<Config, RunnerPtr, GemmIdType, GemmIdHashType>::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<int> 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 <typename Config, typename RunnerPtr, typename GemmIdType, typename GemmIdHashType>
+std::optional<Config> GemmPluginProfiler<Config, RunnerPtr, GemmIdType, GemmIdHashType>::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<int>(mDims.n);
+  int const k = static_cast<int>(mDims.k);
+  size_t const workspace_bytes = computeTmpSize(target, n, k);
+
+  CUDA_CALL_THROW(cudaStreamCreate(&mStream));
+  mWorkspaceTmp = onnxruntime::IAllocator::MakeUniquePtr<char>(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 <typename Config, typename RunnerPtr, typename GemmIdType, typename GemmIdHashType>
 std::optional<Config> GemmPluginProfiler<Config, RunnerPtr, GemmIdType, GemmIdHashType>::profileTacticsForProblem(
     int m, int n, int k, std::vector<Config> const& tactics) {

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"
@@ -371,18 +372,30 @@ Status MatMulNBits<T>::ComputeInternal(OpKernelContext* ctx) const {
   if constexpr (std::is_same<T, MLFloat16>::value || std::is_same<T, BFloat16>::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<const void*>(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