Cache the hipBLASLt int8 matmul algo per problem on ROCm

[Abdennacer-Badaoui]

On ROCm, igemmlt calls hipblasLtMatmulAlgoGetHeuristic on every int8 matmul. On MI300X that query takes ~430 ms (~6.8 s on the first call), while the matmul itself takes under 0.2 ms. Every forward and backward of every Linear8bitLt layer paid it, so 8-bit (LLM.int8) training and inference on AMD GPUs were several hundred times slower than they should be.

This PR caches the algo that ran successfully, so the heuristic is queried once per problem instead of on every call:

• The cache is keyed on (handle, m, k, lda, ldb, ldc) and a bucket of n, the token count (same buckets as aiter's getPaddedM). The first algo found for a bucket is reused by every n in it, so a new sequence length usually doesn't query again. A shared algo is checked with hipblaslt_ext::matmulIsAlgoSupported before it runs. An n it doesn't fit gets its own entry from a full query.
• The existing "try several candidates and keep the first one that runs" logic is kept for every query.
• Problems where none of the workspace-free algos run are cached too, and keep returning ERR_NOT_IMPLEMENTED for the fp32 fallback without querying again. A shared algo that fails at run time triggers a fresh query instead of an error.
• The heuristic preference is now destroyed (it leaked), and the fallback path no longer returns before the layout and matmul descriptors are destroyed.

Only the HIP path changes. CUDA is untouched.

Results

MI300X (gfx942), ROCm 7.2, torch 2.13, built with -DCOMPUTE_BACKEND=hip -DBNB_ROCM_ARCH=gfx942.

Per call (torch.ops.bitsandbytes.int8_linear_matmul, Qwen2-0.5B layer shapes):

Shape (tokens × in → out)	Before	After	torch._int_mm
256 × 896 → 896	424 ms	0.038 ms	0.029 ms
256 × 896 → 4864	450 ms	0.038 ms	0.049 ms
256 × 4864 → 896	434 ms	0.171 ms	0.085 ms
256 × 896 → 128	421 ms	0.023 ms	0.024 ms
1024 × 896 → 4864	428 ms	0.042 ms	0.056 ms

End to end, this makes 8-bit + LoRA training with TRL about 100× faster on MI300X: SFT of Qwen2-0.5B went from ~75 s per step to 34.7 s for 50 steps (~0.7 s per step, including the one-time heuristic queries). The loss curve matches 4-bit QLoRA.

Sharing algos across n buckets matters for workloads with many distinct sequence lengths: sweeping n = 1..300 over three Qwen2-0.5B layer shapes takes 66 heuristic queries instead of 900, with unchanged per-call latency.

Tests

• Outputs are bit-exact against an int32 reference on all tested shapes, including small (1–5 tokens) and non-multiple-of-8 ones.
• The n = 1..300 sweep above is bit-exact for all 900 calls.
• pytest tests/test_ops.py tests/test_linear8bitlt.py tests/test_autograd.py tests/test_functional.py -k "int8 or igemmlt or 8bit or matmullt or Linear8bit": 748 passed, 2 skipped, 2 xfailed.
• Also tested by @kashif on a Radeon 890M (gfx1150, ROCm 10, torch 2.12): bit-exact outputs and the same int8 test subset passes.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[kashif]

Nice find, thanks! I tested this locally on a Radeon 890M (gfx1150, ROCm 10, torch 2.12): outputs are bit-exact and the int8 test subset passes (748 passed). One thing I noticed: n is the token count, so with variable sequence lengths every new length still pays for a full heuristic query per layer shape.

I tried a small follow-up on top of your branch, borrowing two ideas from aiter's hipBLASLt GEMM code:

• Bucket n (same buckets as aiter's getPaddedM: next power of two below 16, then multiples of 16/32/64/128). The first algo found for a bucket is reused by every n in that bucket.
• Check with hipblaslt_ext::matmulIsAlgoSupported before launching, both for heuristic candidates and for reused algos, instead of finding out through a failed launch. If the bucket's algo doesn't fit a given n, or nothing works for it, that n gets its own exact entry, so each n is queried at most once. A cached algo that stops working now re-queries instead of raising RuntimeError.

Sweeping n = 1..300 over three Qwen2-0.5B layer shapes (900 calls) took 66 heuristic queries instead of 900. Per-call latency is unchanged, and outputs stay bit-exact. I also tested with an LD_PRELOAD shim that forces hipblasLtMatmul to fail for chosen n: the fp32 fallback kicks in for exactly those n, there's no re-querying on later calls, and a one-off failure recovers without falling back. I couldn't measure the MI300X numbers here since the heuristic is cheap on gfx1150, so it would be great if you could check it there. It also needs hipblaslt-ext.hpp; ROCm 10 has it, but I haven't built against 6.4.

Happy to open this as a separate PR instead if you'd rather keep this one minimal.

Patch (on top of 66b12f8)

diff --git c/csrc/ops.cu w/csrc/ops.cu
index 26b1a86..f6f8904 100644
--- c/csrc/ops.cu
+++ w/csrc/ops.cu
@@ -16,6 +16,9 @@
 #include <map>
 #include <mutex>
 #include <tuple>
+#ifndef NO_HIPBLASLT
+#include <hipblaslt/hipblaslt-ext.hpp>
+#endif
 
 // NOTE: This queries device 0 once and caches the result. On mixed RDNA+CDNA
 // systems (warp size 32 vs 64) this will return the wrong value for whichever
@@ -284,16 +287,61 @@ int roundoff(int v, int d) { return (v + d - 1) / d * d; }
 
 #if BNB_HIP && !defined(NO_HIPBLASLT)
 // hipBLASLt's heuristic query takes hundreds of milliseconds per call on MI300X, far longer than the matmul
-// itself, so the algo that ran successfully is cached per problem. A missing algo (has_algo = false) means no
-// workspace-free algo ran for that problem and the fp32 fallback is used.
+// itself, so the algo that ran successfully is cached. n is the token count, the only dimension that changes
+// between calls for a given layer, so the first algo found is shared by every n in its bucket (same buckets as
+// aiter's getPaddedM) and a new sequence length does not query again. A shared algo is checked with
+// matmulIsAlgoSupported before it runs. When it does not fit, the n gets its own entry from a full query, which
+// takes precedence over the bucket. A missing algo (has_algo = false) means no workspace-free algo ran for that n
+// and the fp32 fallback is used.
 struct igemmlt_algo_entry {
     bool has_algo;
     hipblasLtMatmulAlgo_t algo;
 };
 
-using igemmlt_algo_key = std::tuple<bnb_blasLt_handle_t, int, int, int, int, int, int>;
+// (handle, m, k, lda, ldb, ldc, n or its bucket, whether the entry is for that exact n)
+using igemmlt_algo_key = std::tuple<bnb_blasLt_handle_t, int, int, int, int, int, int, bool>;
 static std::map<igemmlt_algo_key, igemmlt_algo_entry> igemmlt_algo_cache;
 static std::mutex igemmlt_algo_cache_mutex;
+
+static int igemmlt_bucket_n(int n) {
+    if (n < 16) {
+        int p = 1;
+        while (p < n)
+            p <<= 1;
+        return p;
+    }
+    if (n <= 256)
+        return roundoff(n, 16);
+    if (n <= 1024)
+        return roundoff(n, 32);
+    if (n <= 4096)
+        return roundoff(n, 64);
+    return roundoff(n, 128);
+}
+
+// Runs the matmul with algo if hipBLASLt reports it supports this problem without a workspace.
+static bool igemmlt_run_algo(
+    bnb_blasLt_handle_t ltHandle, bnb_blasLt_matmul_desc_t matmulDesc, const int* alpha, const int8_t* A,
+    bnb_blasLt_layout_t aDesc, const int8_t* B, bnb_blasLt_layout_t bDesc, const int* beta, void* C,
+    bnb_blasLt_layout_t cDesc, hipblasLtMatmulAlgo_t algo, bnb_stream_t stream
+) {
+    size_t workspace_size = 0;
+    if (hipblaslt_ext::matmulIsAlgoSupported(
+            ltHandle, matmulDesc, alpha, aDesc, bDesc, beta, cDesc, cDesc, algo, workspace_size
+        ) != HIPBLAS_STATUS_SUCCESS ||
+        workspace_size != 0)
+        return false;
+
+    if (bnb_blasLtMatmul(
+            ltHandle, matmulDesc, alpha, A, aDesc, B, bDesc, beta, (int32_t*)C, cDesc, (int32_t*)C, cDesc, &algo, NULL,
+            0, stream
+        ) == BNB_BLAS_STATUS_SUCCESS)
+        return true;
+
+    // Drain the HIP last-error flag a failed launch sets, otherwise the next unrelated HIP call inherits it.
+    (void)hipGetLastError();
+    return false;
+}
 #endif
 
 template <int DTYPE_OUT, int SCALE_ROWS>
@@ -336,18 +384,28 @@ int igemmlt(
     if (DTYPE_OUT == 32) {
 #if BNB_HIP
         int alpha = 1, beta = 0;
-        const igemmlt_algo_key key{ltHandle, m, n, k, lda, ldb, ldc};
+        const igemmlt_algo_key exact_key{ltHandle, m, k, lda, ldb, ldc, n, true};
+        const igemmlt_algo_key bucket_key{ltHandle, m, k, lda, ldb, ldc, igemmlt_bucket_n(n), false};
         igemmlt_algo_entry entry{};
-        bool cached;
+        bool exact = false, shared = false;
         {
             std::lock_guard<std::mutex> lock(igemmlt_algo_cache_mutex);
-            auto it = igemmlt_algo_cache.find(key);
-            cached = it != igemmlt_algo_cache.end();
-            if (cached)
+            auto it = igemmlt_algo_cache.find(exact_key);
+            if (it != igemmlt_algo_cache.end()) {
+                exact = true;
                 entry = it->second;
+            } else if ((it = igemmlt_algo_cache.find(bucket_key)) != igemmlt_algo_cache.end()) {
+                shared = true;
+                entry = it->second;
+            }
         }
 
-        if (!cached) {
+        bool ran =
+            entry.has_algo &&
+            igemmlt_run_algo(ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, C, cDesc, entry.algo, stream);
+
+        // An entry for the exact n is final, so a problem without a usable algo is only queried once.
+        if (!ran && !exact) {
             // HIP requires heuristic algo selection
             const int64_t max_workspace_size = 0; // set to 0 to avoid choosing GSU kernel
 
@@ -359,7 +417,7 @@ int igemmlt(
 
             // hipBLASLt's first heuristic algo can be unusable for small-n int8 gemms
             // (e.g. n=4 on MI300X) and fails at matmul time with INVALID_VALUE. Request
-            // several candidates and use the first one that actually runs successfully.
+            // several candidates and use the first one that is supported and runs.
             const int request_solutions = 8;
             bnb_blasLt_heuristic_t heuristicResult[request_solutions];
             int returnedAlgoCount = 0;
@@ -373,34 +431,24 @@ int igemmlt(
                 has_error = 1;
                 fprintf(stderr, "Error: Matmul Algo Heuristic didn't return algorithms\n");
             } else {
-                for (int i = 0; i < returnedAlgoCount; ++i) {
-                    if (bnb_blasLtMatmul(
-                            ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int32_t*)C, cDesc, (int32_t*)C,
-                            cDesc, &heuristicResult[i].algo, NULL, 0, stream
-                        ) == BNB_BLAS_STATUS_SUCCESS) {
+                entry = {false, {}};
+                for (int i = 0; i < returnedAlgoCount && !ran; ++i) {
+                    ran = igemmlt_run_algo(
+                        ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, C, cDesc, heuristicResult[i].algo,
+                        stream
+                    );
+                    if (ran)
                         entry = {true, heuristicResult[i].algo};
-                        break;
-                    }
-                }
-                if (!entry.has_algo) {
-                    // Every workspace-free algo hipBLASLt offered failed at runtime
-                    // (seen on MI300X for some small-n int8 gemms). Drain the HIP
-                    // last-error flag the failed launches set, otherwise the next
-                    // unrelated HIP call will inherit it.
-                    (void)hipGetLastError();
                 }
+                // The bucket keeps the first algo that ran for it. An n it does not fit, or one without any
+                // usable algo, gets its own entry instead of replacing the bucket's.
                 std::lock_guard<std::mutex> lock(igemmlt_algo_cache_mutex);
-                igemmlt_algo_cache[key] = entry;
+                igemmlt_algo_cache[shared || !ran ? exact_key : bucket_key] = entry;
             }
-        } else if (entry.has_algo) {
-            has_error |= checkBlasLtStatus(bnb_blasLtMatmul(
-                ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int32_t*)C, cDesc, (int32_t*)C, cDesc,
-                &entry.algo, NULL, 0, stream
-            ));
         }
 
         // Signal the Python wrapper to take the fp32 fallback via ERR_NOT_IMPLEMENTED.
-        if (!has_error && !entry.has_algo)
+        if (!has_error && !ran)
             has_error = ERR_NOT_IMPLEMENTED;
 #else
         int alpha = 1, beta = 0;

[Abdennacer-Badaoui]

Thanks @kashif