update(rocm): skip the dequant LRU for GEMMs routed above the WMMA ceiling

docs/benchmark_results/rocm-bf16-qmm-route-gfx1151-2026-09-30.md

@@ -46,3 +46,31 @@ Decode is unchanged in every cell (M is 1, which never took the WMMA kernel). ML
 ## The dense prefill path on ROCm
 
 `MLXCEL_PREFILL_DEQUANT_MIN_M=1024` against unset, same binary, pp2048, three ABBA runs: gemma-3-4b-it-4bit 2871 against 2852 tok/s, Qwen3-0.6B-4bit 4174 against 4185. With the route change the dense path is eligible on ROCm only where `quantized_matmul` already runs the same GEMM, so turning it on changes neither bytes nor speed there.
+
+## 2026-10-07: the dequantized-weight cache above the WMMA ceiling (issue #2151)
+
+The route above keeps every dequantized weight in an LRU of 8 matrices or 256 MB (`MLX_ROCM_QMM_DEQUANT_CACHE_SIZE`, `MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES`). A forward pass runs more distinct projections than that, so the cache never hits, and its last entries stay alive into decode. GEMMs routed above the ceiling (`QmmRoute::DequantGemmAboveWmmaCeiling`) now dequantize into a temporary and leave the cache alone; f16 and the other shapes on the same GEMM keep it. `MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` prints the cache counters at exit.
+
+Same host. Before is main at `87538835`; after is that commit plus the change. The before and after medians were taken on 2026-10-07, except the third Gemma before run, which was taken on 2026-10-08 (the first failed in the harness, not the bench). `scripts/bench_decode.sh <model> --prompt-tokens 2048 --max-tokens 128` (20-token same-process warmup, then the measured pass), each run in its own `scripts/rocm_gpu_guard.sh --idle-secs 60` with before and after alternated, every run CLEAN. Medians of three runs (range). "Active after prefill" is the `[Memory] after warmup (prefill + warmup tokens)` line, read once the warmup prefill and its decode tokens are done; "active at end" is the `after measured pass` line.
+
+| Model | Scales | Arm | Prefill tok/s | Decode tok/s | MLX peak | Active after prefill | Active at end |
+|---|---|---|---:|---:|---:|---:|---:|
+| gemma-3-4b-it-4bit | bf16 | before | 2872 (2819-2910) | 62.75 (62.68-62.76) | 4.52 GB | 2.80 GB | 2.97 GB |
+| gemma-3-4b-it-4bit | bf16 | after | 2868 (2839-2945) | 62.71 (62.55-62.73) | 4.52 GB | 2.56 GB | 2.73 GB |
+| Meta-Llama-3.1-8B-Instruct-4bit | f16 | before | 1136 (1106-1136) | 36.63 (32.36-36.84) | 6.92 GB | 4.75 GB | 4.75 GB |
+| Meta-Llama-3.1-8B-Instruct-4bit | f16 | after | 1132 (1048-1143) | 36.67 (36.58-36.92) | 6.92 GB | 4.75 GB | 4.75 GB |
+
+Cache counters for the same run (warmup and measured prefill together). The before counts come from a build of the counters alone, which routes and caches exactly as `87538835` does; one run each, since the counts are deterministic and matched across the three after runs.
+
+| Model | Arm | Hits | Misses | Inserts | Evictions | Bypassed | Entries at exit | Bytes at exit |
+|---|---|---:|---:|---:|---:|---:|---:|---:|
+| gemma-3-4b-it-4bit | before (counters build) | 0 | 340 | 340 | 334 | 0 | 6 | 241,172,480 |
+| gemma-3-4b-it-4bit | after | 0 | 0 | 0 | 0 | 340 | 0 | 0 |
+| Meta-Llama-3.1-8B-Instruct-4bit | before (counters build) | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 |
+| Meta-Llama-3.1-8B-Instruct-4bit | after | 0 | 320 | 320 | 318 | 0 | 2 | 234,881,024 |
+
+For Gemma 3 4B the cache held 230 MiB of bf16 weights after prefill and never hit; skipping it lowers active memory after prefill and at the end of the run by 0.24 GB. MLX peak memory does not move, because the peak comes from the per-call dequantized copies inside prefill, which the cache never removed. Prefill and decode differ by less than the run-to-run range. Llama 3.1 8B is the f16 control and reads the same in both arms: its GEMMs never take the ceiling route, and its cache also gets zero hits while holding 224 MiB after prefill; turning its default off is lablup/mlxcel#2232. Decode at batch 1 runs one row, which takes qmv and never reaches the cache.
+
+Logits: `examples/logit_trace` on gemma-3-4b-it-4bit (README.md as the corpus, 4 chunks of 256 tokens, so every projection runs at 256 rows and takes the ceiling route), before and after, gives byte-identical traces (sha256 `e470fd1c...`), and `scripts/compare_logit_traces.py --decided 2.0` reports 0 of 1024 disagreements and an identical perplexity.
+
+Confirmation on the change rebased onto main `ad844354` (2026-10-08), one guarded run per model: gemma-3-4b-it-4bit 2885 tok/s prefill, 62.14 tok/s decode, MLX peak 4.52 GB, active after prefill 2.56 GB, 340 bypasses and an empty cache; Meta-Llama-3.1-8B-Instruct-4bit 1119 tok/s prefill, 36.66 tok/s decode, MLX peak 6.92 GB, active 4.75 GB, and the same 320 misses, 0 hits and 2 cached entries. Both agree with the after runs above. Gemma's active memory at the end of the run reads 2.56 GB there against 2.73 GB on `87538835`; main moved between the two bases and that difference was not isolated.

docs/environment-variables.md

@@ -259,7 +259,12 @@ default applies (lablup/mlxcel#2152). An unset or empty variable takes the
 default silently. `MLX_ROCM_QMM_DEQUANT_M_THRESHOLD` (1 to 2147483647, default
 the built-in crossover) is the row count from which a quantized GEMM prefers
 dequantize + GEMM. `MLX_ROCM_QMM_DEQUANT_CACHE_SIZE` (0 to 2147483647, default
-8, `0` turns it off) is the number of dequantized weights kept for reuse.
+8, `0` turns it off) is the number of dequantized weights kept for reuse; the
+bf16 GEMMs sent above `MLX_ROCM_WMMA_QMM_MAX_M` never use that cache, since a
+forward pass at those row counts cycles it without a hit (lablup/mlxcel#2151).
+`MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1` prints the cache's hit, miss, insert,
+eviction and bypass counts, its entry count and its bytes on stderr when the
+process exits.
 `MLX_ROCM_GEMM_F32_SOLUTION_INDEX` and `MLX_ROCM_GEMM_BF16_SOLUTION_INDEX` (0
 to 2147483647, default 0) pick a rocBLAS solution, and their `_BATCHED_`
 variants (same range) fall back to them. `MLX_ROCM_QMV_TILE_N` (1 to 32,

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/no_rocm.cpp

@@ -67,6 +67,10 @@ bool quantized_matmul_runs_dequant_gemm(
   return false;
 }
 
+DequantCacheStats dequant_cache_stats() {
+  return DequantCacheStats{};
+}
+
 } // namespace rocm
 
 namespace fast {

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip

@@ -9,6 +9,7 @@
 #include "mlx/backend/rocm/kernel_utils.hpp"
 #include "mlx/backend/rocm/quantized/quantized.h"
 #include "mlx/backend/rocm/quantized/qmv_tiled_kernel.hip"
+#include "mlx/backend/rocm/rocm.h"
 #include "mlx/primitives.h"
 
 #include <hip/hip_bfloat16.h>
@@ -35,6 +36,8 @@
 #endif
 #include <algorithm>
 #include <atomic>
+#include <cstdio>
+#include <cstdlib>
 #include <cstring>
 #include <limits>
 #include <list>
@@ -808,6 +811,13 @@ enum class QmmRoute {
   WmmaDense, // qmm_wmma_dense_kernel, fused dequant + rocWMMA
   DequantFp8Gemm, // dequant_fp8_gemm (weight and activation in e4m3)
   DequantGemm, // dequantize (cached) + dequant_rocblas_gemm in the x dtype
+  // The same GEMM for a bf16 shape the fused kernel would take below the row
+  // ceiling (wmma_qmm_max_m). The weight is dequantized into a temporary on
+  // every call and never enters the dequantized-weight cache: a forward pass
+  // runs more distinct projections at these row counts than the cache holds,
+  // so it never hit, and its last entries stayed alive through decode
+  // (lablup/mlxcel#2151).
+  DequantGemmAboveWmmaCeiling,
   Other, // qmv, batched qmv and the tiled kernels below
 };
 
@@ -905,7 +915,7 @@ inline QmmRoute select_qmm_route(const QmmRouteInputs& in, rocm::Device& d) {
     if (!hipblaslt_is_faster) {
       return QmmRoute::WmmaDense;
     }
-    return QmmRoute::DequantGemm;
+    return QmmRoute::DequantGemmAboveWmmaCeiling;
   }
   if (fp8()) {
     return QmmRoute::DequantFp8Gemm;
@@ -978,6 +988,97 @@ inline size_t dequant_cache_max_bytes() {
   return max_bytes;
 }
 
+// The dequantized-weight LRU of the DequantGemm route and its counters
+// (lablup/mlxcel#2151). The counters are relaxed atomics, bumped once per
+// GEMM on the route and never read on the hot path; rocm::dequant_cache_stats
+// reads them with the entry count and bytes. MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1
+// prints them on stderr when the process exits.
+struct DequantCacheCounters {
+  std::atomic<uint64_t> hits{0};
+  std::atomic<uint64_t> misses{0};
+  std::atomic<uint64_t> inserts{0};
+  std::atomic<uint64_t> evictions{0};
+  std::atomic<uint64_t> bypasses{0};
+};
+
+DequantCacheCounters g_dequant_cache_counters;
+
+inline void count_dequant_cache(std::atomic<uint64_t>& counter) {
+  counter.fetch_add(1, std::memory_order_relaxed);
+}
+
+struct DequantCache {
+  std::mutex mutex;
+  std::list<DequantCacheKey> lru;
+  size_t cached_bytes = 0;
+  std::unordered_map<DequantCacheKey, DequantCacheEntry, DequantCacheKeyHasher>
+      entries;
+
+  // Drops least recently used entries until both limits hold. Call with
+  // `mutex` held.
+  void evict_to(int capacity, size_t max_bytes) {
+    while (!lru.empty() &&
+           (static_cast<int>(entries.size()) > capacity ||
+            cached_bytes > max_bytes)) {
+      auto evict_it = entries.find(lru.back());
+      if (evict_it != entries.end()) {
+        cached_bytes -= evict_it->second.bytes;
+        entries.erase(evict_it);
+        count_dequant_cache(g_dequant_cache_counters.evictions);
+      }
+      lru.pop_back();
+    }
+  }
+};
+
+void print_dequant_cache_stats();
+
+// The process-wide cache. The exit report is registered after the cache is
+// constructed, so it runs before the cache is destroyed.
+DequantCache& dequant_cache() {
+  static DequantCache cache;
+  static const bool report_registered = [] {
+    const char* e = std::getenv("MLX_ROCM_QMM_DEQUANT_CACHE_STATS");
+    if (e && e[0] == '1' && e[1] == '\0') {
+      std::atexit(print_dequant_cache_stats);
+      return true;
+    }
+    return false;
+  }();
+  (void)report_registered;
+  return cache;
+}
+
+rocm::DequantCacheStats read_dequant_cache_stats() {
+  rocm::DequantCacheStats stats{};
+  const auto& c = g_dequant_cache_counters;
+  stats.hits = c.hits.load(std::memory_order_relaxed);
+  stats.misses = c.misses.load(std::memory_order_relaxed);
+  stats.inserts = c.inserts.load(std::memory_order_relaxed);
+  stats.evictions = c.evictions.load(std::memory_order_relaxed);
+  stats.bypasses = c.bypasses.load(std::memory_order_relaxed);
+  auto& cache = dequant_cache();
+  std::lock_guard<std::mutex> lock(cache.mutex);
+  stats.entries = cache.entries.size();
+  stats.bytes = cache.cached_bytes;
+  return stats;
+}
+
+void print_dequant_cache_stats() {
+  const auto s = read_dequant_cache_stats();
+  std::fprintf(
+      stderr,
+      "[ROCm] qmm dequant cache: hits=%llu misses=%llu inserts=%llu "
+      "evictions=%llu bypassed=%llu entries=%zu bytes=%zu\n",
+      static_cast<unsigned long long>(s.hits),
+      static_cast<unsigned long long>(s.misses),
+      static_cast<unsigned long long>(s.inserts),
+      static_cast<unsigned long long>(s.evictions),
+      static_cast<unsigned long long>(s.bypasses),
+      s.entries,
+      s.bytes);
+}
+
 inline rocblas_operation to_rocblas_op(bool transpose) {
   return transpose ? rocblas_operation_transpose : rocblas_operation_none;
 }
@@ -3411,7 +3512,13 @@ bool quantized_matmul_runs_dequant_gemm(
   // and fall back to different rocBLAS calls, so only the hipBLASLt case is
   // claimed. The claim assumes the hipBLASLt launch itself does not throw,
   // which would send each side to its own rocBLAS fallback.
-  return route == QmmRoute::DequantGemm && is_hipblaslt_available();
+  return (route == QmmRoute::DequantGemm ||
+          route == QmmRoute::DequantGemmAboveWmmaCeiling) &&
+      is_hipblaslt_available();
+}
+
+DequantCacheStats dequant_cache_stats() {
+  return read_dequant_cache_stats();
 }
 
 } // namespace rocm
@@ -3524,7 +3631,9 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   // Dequant + rocBLAS GEMM path
   // Disable with MLX_ROCM_QMM_DEQUANT_GEMM=0 if needed
-  if (route == QmmRoute::DequantGemm || route == QmmRoute::DequantFp8Gemm) {
+  if (route == QmmRoute::DequantGemm ||
+      route == QmmRoute::DequantGemmAboveWmmaCeiling ||
+      route == QmmRoute::DequantFp8Gemm) {
     if (!((x_batch_count == 1) || (x_batch_count == batch_count))) {
       throw std::runtime_error(
           "Unsupported x batch shape for dequant GEMM fallback");
@@ -3567,15 +3676,16 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
     bool cache_hit = false;
     int cache_cap = dequant_cache_capacity();
     size_t cache_max_bytes = dequant_cache_max_bytes();
-    if (cache_cap > 0 && cache_max_bytes > 0) {
-      static std::mutex cache_mutex;
-      static std::list<DequantCacheKey> lru;
-      static size_t cached_bytes = 0;
-      static std::unordered_map<
-          DequantCacheKey,
-          DequantCacheEntry,
-          DequantCacheKeyHasher>
-          cache;
+    // Above the WMMA ceiling the cache is skipped: the weight is dequantized
+    // into a temporary below, as with the cache turned off.
+    const bool use_cache = route == QmmRoute::DequantGemm && cache_cap > 0 &&
+        cache_max_bytes > 0;
+    if (route == QmmRoute::DequantGemmAboveWmmaCeiling) {
+      dequant_cache(); // registers the exit report, if requested
+      count_dequant_cache(g_dequant_cache_counters.bypasses);
+    }
+    if (use_cache) {
+      auto& c = dequant_cache();
 
       DequantCacheKey key{
           reinterpret_cast<std::uintptr_t>(gpu_ptr<void>(w)),
@@ -3589,14 +3699,18 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
           x.dtype()};
 
       {
-        std::lock_guard<std::mutex> lock(cache_mutex);
-        auto it = cache.find(key);
-        if (it != cache.end() && it->second.weight.shape() == w_dequant_shape) {
-          lru.splice(lru.begin(), lru, it->second.lru_it);
+        std::lock_guard<std::mutex> lock(c.mutex);
+        auto it = c.entries.find(key);
+        if (it != c.entries.end() &&
+            it->second.weight.shape() == w_dequant_shape) {
+          c.lru.splice(c.lru.begin(), c.lru, it->second.lru_it);
           w_dequant = it->second.weight;
           cache_hit = true;
         }
       }
+      count_dequant_cache(
+          cache_hit ? g_dequant_cache_counters.hits
+                    : g_dequant_cache_counters.misses);
 
       if (!cache_hit) {
         w_dequant.set_data(mlx::core::rocm::malloc_async(w_dequant.nbytes(), enc));
@@ -3608,61 +3722,44 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
           fp_dequantize(w, scales, w_dequant, group_size_, bits_, enc, s);
         }
 
-        std::lock_guard<std::mutex> lock(cache_mutex);
-        auto it = cache.find(key);
-        if (it == cache.end()) {
+        std::lock_guard<std::mutex> lock(c.mutex);
+        auto it = c.entries.find(key);
+        if (it == c.entries.end()) {
           size_t entry_bytes = w_dequant.nbytes();
           if (entry_bytes <= cache_max_bytes) {
-            lru.push_front(key);
-            cache.emplace(
+            c.lru.push_front(key);
+            c.entries.emplace(
                 key,
                 DequantCacheEntry{
                     w_dequant,
                     w,
                     scales,
                     has_bias ? std::optional<array>(*biases) : std::nullopt,
                     entry_bytes,
-                    lru.begin()});
-            cached_bytes += entry_bytes;
-
-            while (static_cast<int>(cache.size()) > cache_cap ||
-                   cached_bytes > cache_max_bytes) {
-              auto evict = lru.back();
-              auto evict_it = cache.find(evict);
-              if (evict_it != cache.end()) {
-                cached_bytes -= evict_it->second.bytes;
-                cache.erase(evict_it);
-              }
-              lru.pop_back();
-            }
+                    c.lru.begin()});
+            c.cached_bytes += entry_bytes;
+            count_dequant_cache(g_dequant_cache_counters.inserts);
+            c.evict_to(cache_cap, cache_max_bytes);
           }
         } else {
           size_t entry_bytes = w_dequant.nbytes();
           if (entry_bytes > cache_max_bytes) {
-            cached_bytes -= it->second.bytes;
-            lru.erase(it->second.lru_it);
-            cache.erase(it);
+            c.cached_bytes -= it->second.bytes;
+            c.lru.erase(it->second.lru_it);
+            c.entries.erase(it);
+            count_dequant_cache(g_dequant_cache_counters.evictions);
           } else {
-            cached_bytes -= it->second.bytes;
+            c.cached_bytes -= it->second.bytes;
             it->second.w_source = w;
             it->second.scales_source = scales;
             it->second.biases_source =
                 has_bias ? std::optional<array>(*biases) : std::nullopt;
             it->second.weight = w_dequant;
             it->second.bytes = entry_bytes;
-            cached_bytes += it->second.bytes;
-            lru.splice(lru.begin(), lru, it->second.lru_it);
-
-            while (static_cast<int>(cache.size()) > cache_cap ||
-                   cached_bytes > cache_max_bytes) {
-              auto evict = lru.back();
-              auto evict_it = cache.find(evict);
-              if (evict_it != cache.end()) {
-                cached_bytes -= evict_it->second.bytes;
-                cache.erase(evict_it);
-              }
-              lru.pop_back();
-            }
+            c.cached_bytes += it->second.bytes;
+            c.lru.splice(c.lru.begin(), c.lru, it->second.lru_it);
+            count_dequant_cache(g_dequant_cache_counters.inserts);
+            c.evict_to(cache_cap, cache_max_bytes);
           }
         }
       }

src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/rocm.h

@@ -8,6 +8,7 @@
 #include "mlx/utils.h"
 
 #include <cstddef>
+#include <cstdint>
 #include <optional>
 #include <vector>
 
@@ -79,4 +80,26 @@ MLX_API bool quantized_matmul_runs_dequant_gemm(
     int group_size,
     int bits);
 
+// Counters of QuantizedMatmul's dequantized-weight cache, process-wide since
+// start (lablup/mlxcel#2151). `hits` and `misses` count lookups on the cached
+// dequantize + GEMM route, `inserts` the entries stored (a refresh of an entry
+// under the same key counts), `evictions` the entries dropped for the
+// MLX_ROCM_QMM_DEQUANT_CACHE_SIZE or MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES
+// limit, and `bypasses` the bf16 GEMMs above the fused kernel's row ceiling,
+// which dequantize into a temporary without touching the cache. `entries` and
+// `bytes` are what the cache holds now. All zero without ROCm.
+// MLX_ROCM_QMM_DEQUANT_CACHE_STATS=1 prints the same numbers on stderr at
+// process exit.
+struct DequantCacheStats {
+  uint64_t hits;
+  uint64_t misses;
+  uint64_t inserts;
+  uint64_t evictions;
+  uint64_t bypasses;
+  size_t entries;
+  size_t bytes;
+};
+
+MLX_API DequantCacheStats dequant_cache_stats();
+
 } // namespace mlx::core::rocm