Address review feedback on fpA_intB tactic cache

- Reject cached CUDA-GEMV tactics when the runner has no GEMV kernel.
- Overwrite cache rows rejected on load with the freshly profiled tactic.
- Fail Flush when the advisory file lock cannot be acquired, instead of
  writing unlocked and risking lost rows.
- Add has_bias to GemmIdCore and the persistent key/TSV schema, and bump
  tactic_selection_version.
- Add the GPU name to GemmIdCore so the process-global profile map is not
  shared between devices that have the same packing SM.
- Renumber the tactic cache doc section to 6.3 and fix its references.

Author: tianleiwu

docs/contrib_ops/cuda/gemm_profiler_cache.md

@@ -99,7 +99,7 @@ punctuation) that make CSV quoting fragile.
 ```
 # ort_cuda_gemm_tactic_cache	v1
 # table	matmulnbits_fpa_intb
-# tactic_selection_version	1
+# tactic_selection_version	2
 # device_name	NVIDIA A100-SXM4-80GB
 # sm	80
 # multiprocessor_count	108
@@ -108,9 +108,9 @@ punctuation) that make CSV quoting fragile.
 # ort_version	1.23.0
 # ort_git_commit	<git-sha-or-unknown>
 # ort_build_config	Release
-n_16b	k	activation_dtype	weight_type	bits	block_size	has_zero_points	zero_point_dtype	gemv_enabled	packing_sm	m_bucket	valid_config	sm_version	tile80	tile90	tile100	tile120	split_k_style	split_k	stages	cluster	mainloop	epilogue	tma	enable_cuda_kernel
-12288	4096	half	uint4b_t	4	64	1	uint4b_t	1	80	1	1	80	23	0	0	0	0	1	1	3	0	0	0	0	1
-12288	4096	half	uint4b_t	4	64	1	uint4b_t	1	80	64	1	80	17	0	0	0	0	0	-1	4	0	0	0	0	0
+n_16b	k	activation_dtype	weight_type	bits	block_size	has_zero_points	zero_point_dtype	gemv_enabled	has_bias	packing_sm	m_bucket	valid_config	sm_version	tile80	tile90	tile100	tile120	split_k_style	split_k	stages	cluster	mainloop	epilogue	tma	enable_cuda_kernel
+12288	4096	half	uint4b_t	4	64	1	uint4b_t	1	0	80	1	1	80	23	0	0	0	0	1	1	3	0	0	0	0	1
+12288	4096	half	uint4b_t	4	64	1	uint4b_t	1	0	80	64	1	80	17	0	0	0	0	0	-1	4	0	0	0	0	0
 ...
 ```
 
@@ -119,7 +119,7 @@ n_16b	k	activation_dtype	weight_type	bits	block_size	has_zero_points	zero_point_
 ```
 # ort_cuda_gemm_tactic_cache	v1
 # table	qmoe_gemm
-# tactic_selection_version	1
+# tactic_selection_version	2
 # device_name	NVIDIA A100-SXM4-80GB
 # sm	80
 # multiprocessor_count	108
@@ -179,7 +179,7 @@ Use op-specific keys rather than a single sparse superset.
 ```
 matmulnbits_key = {
   n_16b, k, activation_dtype, weight_type, bits, block_size,
-  has_zero_points, zero_point_dtype, gemv_enabled, packing_sm
+  has_zero_points, zero_point_dtype, gemv_enabled, has_bias, packing_sm
 }
 
 qmoe_key = {
@@ -236,7 +236,7 @@ unit-testable.
 - *Strategy*: Wire the `gemm_tactic_cache` utility independently into `GemmPluginProfiler` and `MoeGemmProfiler` rather than attempting a large refactor. The shared cache utility provides a simple `GetConfig(...)` / `PutConfig(...)` API.
 - `gemm_profiler.h`: before the M sweep in `profileTactics`, load matching entries and skip if hit;
   after profiling, merge into the disk cache. Re-enable the serialize stubs.
-- `matmul_nbits.{h,cc}`: pass `bits / block_size / zero-point mode / gemv_enabled / packing_sm`
+- `matmul_nbits.{h,cc}`: pass `bits / block_size / zero-point mode / gemv_enabled / has_bias / packing_sm`
   into the MatMulNBits key.
 - `moe_quantization.cc` + `moe_gemm_profiler.{h,cc}`: use the same disk cache + key; QMoE gains
   cross-node dedup via the shared cache. Add a process-global MoE profiler manager

docs/contrib_ops/cuda/matmul_nbits.md

@@ -531,7 +531,7 @@ Prepacked weights are intentionally strict:
 
 ---
 
-## 6.1 fpA_intB Tactic Autotune Cache
+### 6.3 fpA_intB Tactic Autotune Cache
 
 The fpA_intB profiler (§6) tunes CUTLASS/CUDA tactics per problem shape. By
 default this is in-process only and re-run every session. When a cache location
@@ -658,9 +658,9 @@ present. `ComputeInternal` then:
 |----------|----------------|--------|
 | `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION` | bool, `0` | Disable the router GEMV specialization (§4.3); shapes fall back to the generic GEMV / dequant path. Useful for A/B benchmarking. |
 | `ORT_FPA_INTB_GEMM` | int/string, `0` | Enable the CUTLASS weight-only path (§6). `0` or `off` disables it, otherwise enables it. |
-| `ORT_FPA_INTB_PROFILE_M` | comma list, unset | Override the M buckets profiled for the fpA_intB tactic cache (§6.1). The maximum value also bounds the initial profile range. Session-config equivalent: `ep.cuda.fpa_intb_profile_m`. |
-| `ORT_CUDA_GEMM_TACTIC_CACHE_DIR` | path, unset | Directory for the persistent fpA_intB tactic cache (§6.1). Unset means the cache is in-process only. Session-config equivalent: `ep.cuda.gemm_tactic_cache_dir`. |
-| `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` | path prefix, unset | Explicit cache file prefix (§6.1); writes `<prefix>.matmulnbits_fpa_intb.tsv`. Session-config equivalent: `ep.cuda.gemm_tactic_cache_prefix`. |
+| `ORT_FPA_INTB_PROFILE_M` | comma list, unset | Override the M buckets profiled for the fpA_intB tactic cache (§6.3). The maximum value also bounds the initial profile range. Session-config equivalent: `ep.cuda.fpa_intb_profile_m`. |
+| `ORT_CUDA_GEMM_TACTIC_CACHE_DIR` | path, unset | Directory for the persistent fpA_intB tactic cache (§6.3). Unset means the cache is in-process only. Session-config equivalent: `ep.cuda.gemm_tactic_cache_dir`. |
+| `ORT_CUDA_GEMM_TACTIC_CACHE_PREFIX` | path prefix, unset | Explicit cache file prefix (§6.3); writes `<prefix>.matmulnbits_fpa_intb.tsv`. Session-config equivalent: `ep.cuda.gemm_tactic_cache_prefix`. |
 | `ORT_MATMULNBITS_FORCE_CHUNKED` | int, `0` | Force the chunked dequant+GEMM fallback (§5) regardless of the size heuristic, and bypass the fpA_intB M-chunking size condition (§6.2). |
 | `ORT_MATMULNBITS_CHUNK_SIZE` | int64, `32768` | Target rows per chunk in the chunked fallback. Values `< 1` reset to the default. |
 | `ORT_MATMULNBITS_M_CHUNK_SIZE` | int, `0` | Max rows of `A` per fpA_intB launch (§6.2). `0` disables M chunking. Overridden by the `ep.cuda.matmul_nbits_m_chunk_size` session config entry. Also applies to the CUDA plugin EP. |
@@ -730,7 +730,7 @@ present. `ComputeInternal` then:
   [onnxruntime/test/contrib_ops/matmul_2bits_test.cc](../../../onnxruntime/test/contrib_ops/matmul_2bits_test.cc).
   Run it from `onnxruntime_provider_test` with
   `--gtest_filter=MatMul2BitsCuda.*`.
-- fpA_intB tactic cache (§6.1) unit tests:
+- fpA_intB tactic cache (§6.3) unit tests:
   [onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc](../../../onnxruntime/test/contrib_ops/cuda_kernels/gemm_tactic_cache_test.cc)
   (config serialize/parse round-trip, signature-mismatch rejection, appended-column
   tolerance, store/load and merge). Build with `onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON`

onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc

@@ -30,6 +30,18 @@
 using namespace onnxruntime::llm::common;
 using namespace onnxruntime::llm::kernels::cutlass_kernels;
 
+namespace {
+// Compares tactics by their persisted representation, which is exactly what a cache hit stores.
+bool SameConfigColumns(const std::optional<onnxruntime::llm::cutlass_extensions::CutlassGemmConfig>& a,
+                       const std::optional<onnxruntime::llm::cutlass_extensions::CutlassGemmConfig>& b) {
+  std::vector<std::string> columns_a;
+  std::vector<std::string> columns_b;
+  onnxruntime::llm::gemm_cache::AppendConfigColumns(columns_a, a);
+  onnxruntime::llm::gemm_cache::AppendConfigColumns(columns_b, b);
+  return columns_a == columns_b;
+}
+}  // namespace
+
 namespace onnxruntime::llm::kernels::weight_only {
 
 std::optional<size_t> ComputeWeightOnlyGemmProfilerScratchSize(
@@ -227,6 +239,7 @@ onnxruntime::llm::gemm_cache::MatMulNBitsKey WeightOnlyGroupwiseQuantGemmPluginP
   key.has_zero_points = mHasZeros;
   key.zero_point_dtype = mHasZeros ? key.weight_type : "none";
   key.gemv_enabled = hasWeightOnlyCudaKernel;
+  key.has_bias = mHasBiases;
   key.packing_sm = mArch;
   return key;
 }
@@ -242,12 +255,12 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::loadPersistentCache(
     return;
   }
 
-  // Validate CUTLASS tactics loaded from disk against the tactics this runner can actually
-  // dispatch. A parseable-but-incompatible cache row (e.g. hand-edited, or written by a build
-  // whose signature happens to match but whose tactic set differs) would otherwise be handed
-  // straight to the kernel. Non-matching CUTLASS tactics are dropped so the bucket is re-profiled.
-  // The synthetic CUDA-GEMV tactic (enableCudaKernel) is not part of getConfigs(); its validity
-  // is already keyed by gemv_enabled in the cache key, so it is accepted as-is.
+  // Validate tactics loaded from disk against the tactics this runner can actually dispatch. A
+  // parseable-but-incompatible cache row (e.g. hand-edited, or written by a build whose signature
+  // happens to match but whose tactic set differs) would otherwise be handed straight to the kernel.
+  // Rejected rows are dropped so the bucket is re-profiled (and the stale row overwritten on staging).
+  // The CUDA-GEMV tactic is only dispatchable when this runner has a GEMV kernel, so it is rejected
+  // otherwise, regardless of what the row's key claims.
   auto const valid_configs = getTactics(0, gemmId.n, gemmId.k);
   auto is_valid_cutlass = [&valid_configs](Config const& c) {
     for (auto const& v : valid_configs) {
@@ -269,6 +282,10 @@ void WeightOnlyGroupwiseQuantGemmPluginProfiler::loadPersistentCache(
         ORT_LLM_LOG_WARNING("Dropping unsupported cached fpA_intB tactic from the tactic cache; re-profiling.");
         continue;
       }
+      if (config->enableCudaKernel && !hasWeightOnlyCudaKernel) {
+        ORT_LLM_LOG_WARNING("Dropping cached fpA_intB CUDA-GEMV tactic: no GEMV kernel is available; re-profiling.");
+        continue;
+      }
       if (!config->enableCudaKernel && !is_valid_cutlass(*config)) {
         ORT_LLM_LOG_WARNING("Dropping incompatible cached fpA_intB tactic from the tactic cache; re-profiling.");
         continue;
@@ -287,11 +304,14 @@ bool WeightOnlyGroupwiseQuantGemmPluginProfiler::stageProfiledTactics(
   auto key = makeCacheKey(gemmId, hasWeightOnlyCudaKernel);
   bool added = false;
   for (auto const& [m, config] : map) {
-    // Only stage buckets that are not already recorded (skips re-staging cache hits).
-    if (!mCache->Get(key, m).has_value()) {
-      mCache->Put(key, m, config);
-      added = true;
+    // Skip buckets already recorded with the same tactic (cache hits). A row rejected on load is still
+    // in the cache, so it must be overwritten by the freshly profiled tactic.
+    auto const cached = mCache->Get(key, m);
+    if (cached.has_value() && SameConfigColumns(*cached, config)) {
+      continue;
     }
+    mCache->Put(key, m, config);
+    added = true;
   }
   return added;
 }

onnxruntime/contrib_ops/cuda/llm/gemm_profiler.h

@@ -26,7 +26,9 @@
 #include <optional>
 #include <shared_mutex>
 #include <sstream>
+#include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "contrib_ops/cuda/llm/nv_infer_datatype.h"
@@ -92,7 +94,9 @@ struct GemmDims {
 // In our case GEMM is uniquely identified by N and K, plus the target SM architecture (so the
 // SM80-compatibility and native SM90 kernels for the same shape do not share profiled configs).
 // The quantization fields keep the process-global profile map aligned with the persistent cache key:
-// kernels with different bits/group size/zero points/GEMV support must not share tactics.
+// kernels with different bits/group size/zero points/bias/GEMV support must not share tactics.
+// device_name separates GPUs that share a packing SM (e.g. RTX 4090 and RTX 4060) so a tactic profiled
+// on one device is never reused, or written to a disk cache, for another.
 class GemmIdCore {
  public:
   int n;
@@ -103,10 +107,13 @@ class GemmIdCore {
   int group_size;
   bool has_zeros;
   bool gemv_enabled;
+  bool has_bias;
+  std::string device_name;
 
   GemmIdCore(int n_, int k_, nvinfer::DataType const& dtype_, int sm_ = 0, int bits_ = 0, int group_size_ = 0,
-             bool has_zeros_ = false, bool gemv_enabled_ = false)
-      : n(n_), k(k_), dtype(dtype_), sm(sm_), bits(bits_), group_size(group_size_), has_zeros(has_zeros_), gemv_enabled(gemv_enabled_) {
+             bool has_zeros_ = false, bool gemv_enabled_ = false, bool has_bias_ = false,
+             std::string device_name_ = std::string())
+      : n(n_), k(k_), dtype(dtype_), sm(sm_), bits(bits_), group_size(group_size_), has_zeros(has_zeros_), gemv_enabled(gemv_enabled_), has_bias(has_bias_), device_name(std::move(device_name_)) {
   }
 
   GemmIdCore()
@@ -115,7 +122,8 @@ class GemmIdCore {
         bits(0),
         group_size(0),
         has_zeros(false),
-        gemv_enabled(false) {
+        gemv_enabled(false),
+        has_bias(false) {
   }
 
   bool operator==(GemmIdCore const& id) const {
@@ -127,14 +135,16 @@ class GemmIdCore {
     out << " type=" << static_cast<int>(id.dtype);
     out << " sm=" << id.sm;
     out << " bits=" << id.bits << " group_size=" << id.group_size;
-    out << " has_zeros=" << id.has_zeros << " gemv=" << id.gemv_enabled;
+    out << " has_zeros=" << id.has_zeros << " gemv=" << id.gemv_enabled << " has_bias=" << id.has_bias;
+    out << " device=" << id.device_name;
     return out;
   }
 
  protected:
   bool isEqual(GemmIdCore const& id) const {
     return n == id.n && k == id.k && dtype == id.dtype && sm == id.sm && bits == id.bits &&
-           group_size == id.group_size && has_zeros == id.has_zeros && gemv_enabled == id.gemv_enabled;
+           group_size == id.group_size && has_zeros == id.has_zeros && gemv_enabled == id.gemv_enabled &&
+           has_bias == id.has_bias && device_name == id.device_name;
   }
 };
 
@@ -145,9 +155,10 @@ struct GemmIdCoreHash {
     auto h2 = std::hash<int>{}(id.k);
     auto h3 = std::hash<int>{}(static_cast<int>(id.dtype));
     auto h4 = std::hash<int>{}(id.sm);
-    auto h5 = std::hash<int>{}((id.bits << 16) ^ (id.group_size << 2) ^ (id.has_zeros ? 2 : 0) ^
-                               (id.gemv_enabled ? 1 : 0));
-    return h1 ^ h2 ^ h3 ^ h4 ^ h5;
+    auto h5 = std::hash<int>{}((id.bits << 16) ^ (id.group_size << 3) ^ (id.has_bias ? 4 : 0) ^
+                               (id.has_zeros ? 2 : 0) ^ (id.gemv_enabled ? 1 : 0));
+    auto h6 = std::hash<std::string>{}(id.device_name);
+    return h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6;
   }
 };
 

onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.cc

@@ -356,7 +356,7 @@ bool MatMulNBitsKey::operator==(const MatMulNBitsKey& o) const {
          activation_dtype == o.activation_dtype && weight_type == o.weight_type &&
          bits == o.bits && block_size == o.block_size &&
          has_zero_points == o.has_zero_points && zero_point_dtype == o.zero_point_dtype &&
-         gemv_enabled == o.gemv_enabled && packing_sm == o.packing_sm;
+         gemv_enabled == o.gemv_enabled && has_bias == o.has_bias && packing_sm == o.packing_sm;
 }
 
 std::size_t MatMulNBitsKeyHash::operator()(const MatMulNBitsKey& k) const {
@@ -372,6 +372,7 @@ std::size_t MatMulNBitsKeyHash::operator()(const MatMulNBitsKey& k) const {
   mix(std::hash<bool>{}(k.has_zero_points));
   mix(std::hash<std::string>{}(k.zero_point_dtype));
   mix(std::hash<bool>{}(k.gemv_enabled));
+  mix(std::hash<bool>{}(k.has_bias));
   mix(std::hash<int>{}(k.packing_sm));
   return h;
 }
@@ -382,7 +383,7 @@ namespace {
 const std::vector<std::string>& MatMulNBitsColumnNames() {
   static const std::vector<std::string> names = {
       "n_16b", "k", "activation_dtype", "weight_type", "bits",
-      "block_size", "has_zero_points", "zero_point_dtype", "gemv_enabled", "packing_sm",
+      "block_size", "has_zero_points", "zero_point_dtype", "gemv_enabled", "has_bias", "packing_sm",
       "m_bucket", "valid_config", "sm_version", "tile80", "tile90", "tile100", "tile120",
       "split_k_style", "split_k", "stages", "cluster", "mainloop", "epilogue", "tma",
       "enable_cuda_kernel"};
@@ -551,6 +552,9 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Load() {
     int gemv = 0;
     ParseInt(field("gemv_enabled"), gemv);
     key.gemv_enabled = gemv != 0;
+    int has_bias = 0;
+    ParseInt(field("has_bias"), has_bias);
+    key.has_bias = has_bias != 0;
 
     int m_bucket = 0;
     if (!ParseInt(field("m_bucket"), m_bucket)) {
@@ -590,7 +594,7 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Load() {
 
 onnxruntime::common::Status MatMulNBitsTacticCache::WriteAllLocked(
     const std::unordered_map<MatMulNBitsKey, BucketMap, MatMulNBitsKeyHash>& table) const {
-  // Unique per writer so concurrent writers never share a temp file, even if the file lock failed.
+  // Unique per writer so concurrent writers never share a temp file.
   static std::atomic<uint64_t> tmp_counter{0};
   const std::string tmp_path = file_path_ + ".tmp." + std::to_string(CurrentProcessId()) + "." +
                                std::to_string(tmp_counter.fetch_add(1));
@@ -633,6 +637,7 @@ onnxruntime::common::Status MatMulNBitsTacticCache::WriteAllLocked(
         row.emplace_back(key.has_zero_points ? "1" : "0");
         row.emplace_back(TsvEncode(key.zero_point_dtype));
         row.emplace_back(key.gemv_enabled ? "1" : "0");
+        row.emplace_back(key.has_bias ? "1" : "0");
         row.emplace_back(std::to_string(key.packing_sm));
         row.emplace_back(std::to_string(m_bucket));
         AppendConfigColumns(row, cfg);
@@ -684,9 +689,12 @@ onnxruntime::common::Status MatMulNBitsTacticCache::Flush() {
     }
   }
 
+  // Without the lock two writers could load the same snapshot and the last rename would drop the
+  // other's rows, so fail the flush (the rows stay dirty and are retried at the next flush point).
   ScopedFileLock lock(file_path_);
-  // Even if the OS lock could not be acquired, proceed best-effort; atomic rename
-  // still guarantees a consistent file, only lost updates become possible.
+  if (!lock.locked()) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to acquire the gemm tactic cache file lock: ", file_path_, ".lock");
+  }
 
   // Reload the current on-disk file so concurrently-written rows are not dropped,
   // then overlay the in-memory table (in-memory wins on conflict).

onnxruntime/contrib_ops/cuda/llm/gemm_tactic_cache.h

@@ -43,7 +43,7 @@ constexpr const char* kCacheMagic = "ort_cuda_gemm_tactic_cache";
 constexpr const char* kCacheFormatVersion = "v1";
 
 // Bump whenever profiling or tactic selection changes which tactic wins, so stale caches are rejected.
-constexpr const char* kTacticSelectionVersion = "1";
+constexpr const char* kTacticSelectionVersion = "2";
 
 constexpr const char* kTableMatMulNBits = "matmulnbits_fpa_intb";
 
@@ -119,6 +119,7 @@ struct MatMulNBitsKey {
   bool has_zero_points = false;
   std::string zero_point_dtype;  // e.g. "uint4b_t" / activation dtype / "none"
   bool gemv_enabled = false;
+  bool has_bias = false;
   int packing_sm = 0;
 
   bool operator==(const MatMulNBitsKey& o) const;

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc

@@ -716,7 +716,8 @@ void MatMulNBits<T>::RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int
   const auto dtype = std::is_same_v<T, BFloat16> ? onnxruntime::llm::nvinfer::DataType::kBF16
                                                  : onnxruntime::llm::nvinfer::DataType::kHALF;
   gemmId_ = GemmIdCore(n_16b, static_cast<int>(K_), dtype, kernel_sm, static_cast<int>(nbits_),
-                       static_cast<int>(block_size_), has_zero_points_, hasWeightOnlyCudaKernel);
+                       static_cast<int>(block_size_), has_zero_points_, hasWeightOnlyCudaKernel, has_bias_,
+                       this->GetDeviceProp().name);
 
   GemmDims dims = {min_m, max_m, n_16b, K_};
   gemmProfiler_->profileTactics(weightOnlyGemmRunner_, gemmId_.dtype, dims, gemmId_, hasWeightOnlyCudaKernel);

onnxruntime/python/tools/fpa_intb_tune.py

@@ -144,6 +144,7 @@ def _summarize_cache(cache_path: str) -> None:
                 "block_size",
                 "has_zero_points",
                 "gemv_enabled",
+                "has_bias",
                 "packing_sm",
             ]
             key = tuple(fields[n_key_col[c]] for c in key_cols if c in n_key_col)