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

onnxruntime/contrib_ops/cuda/llm/common/cuda_runtime_utils.h

@@ -16,7 +16,9 @@
  */
 #pragma once
 
+#include <cstdlib>
 #include <optional>
+#include <string>
 #include <cuda_runtime_api.h>
 #ifdef ENABLE_FP8
 #include <cuda_fp8.h>
@@ -61,12 +63,27 @@ inline std::optional<bool> isCudaLaunchBlocking() {
   thread_local bool firstCall = true;
   thread_local std::optional<bool> 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;

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 <algorithm>
+#include <set>
+
+#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 @@
     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<int>(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,17 +77,17 @@
 
 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<int>(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
-      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<int>(maxM), originalN, static_cast<int>(k))  // workspace
   };
-  return calculateTotalWorkspaceSize(workspaces.data(), workspaces.size());
+  return calculateTotalWorkspaceSize(workspaces.data(), static_cast<int>(workspaces.size()));
 }
 
 std::vector<WeightOnlyGroupwiseQuantGemmPluginProfiler::Config> WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics(
@@ -97,5 +103,53 @@
   return true;
 }
 
+std::vector<int> WeightOnlyGroupwiseQuantGemmPluginProfiler::ParseProfileMList(const std::string& value) {
+  std::vector<int> result;
+  if (value.empty()) {
+    return result;
+  }
+  std::set<int> 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<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