From 66b12f87030681beef7f809bb652470fcc0ca9b7 Mon Sep 17 00:00:00 2001 From: Abdennacer-Badaoui Date: Fri, 16 Oct 2026 09:18:09 +0000 Subject: [PATCH 1/2] Cache the hipBLASLt int8 matmul algo per problem on ROCm hipBLASLt's heuristic query takes ~430 ms per call on MI300X, while the int8 matmul itself takes well under 1 ms, and igemmlt ran it on every call. Cache the algo that ran successfully per (handle, m, n, k, lda, ldb, ldc), including the problems where no workspace-free algo runs, which keep taking the fp32 fallback. Also destroy the heuristic preference and the descriptors on the fallback path, which leaked before. --- csrc/compat.cuh | 1 + csrc/ops.cu | 118 ++++++++++++++++++++++++++++++++---------------- 2 files changed, 80 insertions(+), 39 deletions(-) diff --git a/csrc/compat.cuh b/csrc/compat.cuh index f8c307c2a..fa92ef917 100644 --- a/csrc/compat.cuh +++ b/csrc/compat.cuh @@ -165,6 +165,7 @@ using bnb_blasLt_preference_t = hipblasLtMatmulPreference_t; #define bnb_blasLtMatmul hipblasLtMatmul #define bnb_blasLtPrefCreate hipblasLtMatmulPreferenceCreate #define bnb_blasLtPrefSetAttr hipblasLtMatmulPreferenceSetAttribute +#define bnb_blasLtPrefDestroy hipblasLtMatmulPreferenceDestroy #define bnb_blasLtAlgoGetHeuristic hipblasLtMatmulAlgoGetHeuristic #define BNB_BLASLT_DESC_TRANSA HIPBLASLT_MATMUL_DESC_TRANSA diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..26b1a86cb 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -13,6 +13,9 @@ #if BNB_HIP #include #include +#include +#include +#include // 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 @@ -279,6 +282,20 @@ void strided_gemmex( 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. +struct igemmlt_algo_entry { + bool has_algo; + hipblasLtMatmulAlgo_t algo; +}; + +using igemmlt_algo_key = std::tuple; +static std::map igemmlt_algo_cache; +static std::mutex igemmlt_algo_cache_mutex; +#endif + template int igemmlt( bnb_blasLt_handle_t ltHandle, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, @@ -318,50 +335,73 @@ int igemmlt( if (DTYPE_OUT == 32) { #if BNB_HIP - // HIP requires heuristic algo selection - const int64_t max_workspace_size = 0; // set to 0 to avoid choosing GSU kernel + int alpha = 1, beta = 0; + const igemmlt_algo_key key{ltHandle, m, n, k, lda, ldb, ldc}; + igemmlt_algo_entry entry{}; + bool cached; + { + std::lock_guard lock(igemmlt_algo_cache_mutex); + auto it = igemmlt_algo_cache.find(key); + cached = it != igemmlt_algo_cache.end(); + if (cached) + entry = it->second; + } - bnb_blasLt_preference_t pref; - checkBlasLtStatus(bnb_blasLtPrefCreate(&pref)); - checkBlasLtStatus( - bnb_blasLtPrefSetAttr(pref, BNB_BLASLT_PREF_MAX_WORKSPACE, &max_workspace_size, sizeof(max_workspace_size)) - ); + if (!cached) { + // HIP requires heuristic algo selection + const int64_t max_workspace_size = 0; // set to 0 to avoid choosing GSU kernel - // 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. - const int request_solutions = 8; - bnb_blasLt_heuristic_t heuristicResult[request_solutions]; - int returnedAlgoCount = 0; - checkBlasLtStatus(bnb_blasLtAlgoGetHeuristic( - ltHandle, matmulDesc, aDesc, bDesc, cDesc, cDesc, pref, request_solutions, heuristicResult, - &returnedAlgoCount - )); + bnb_blasLt_preference_t pref; + checkBlasLtStatus(bnb_blasLtPrefCreate(&pref)); + checkBlasLtStatus(bnb_blasLtPrefSetAttr( + pref, BNB_BLASLT_PREF_MAX_WORKSPACE, &max_workspace_size, sizeof(max_workspace_size) + )); - if (returnedAlgoCount == 0) { - has_error = 1; - fprintf(stderr, "Error: Matmul Algo Heuristic didn't return algorithms\n"); - } else { - int alpha = 1, beta = 0; - bnb_blas_status_t matmul_status = BNB_BLAS_STATUS_SUCCESS; - for (int i = 0; i < returnedAlgoCount; ++i) { - matmul_status = bnb_blasLtMatmul( - ltHandle, matmulDesc, &alpha, A, aDesc, B, bDesc, &beta, (int32_t*)C, cDesc, (int32_t*)C, cDesc, - &heuristicResult[i].algo, NULL, 0, stream - ); - if (matmul_status == BNB_BLAS_STATUS_SUCCESS) - break; - } - if (matmul_status != BNB_BLAS_STATUS_SUCCESS) { - // 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. Then signal the Python - // wrapper to take the fp32 fallback via ERR_NOT_IMPLEMENTED. - (void)hipGetLastError(); - return ERR_NOT_IMPLEMENTED; + // 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. + const int request_solutions = 8; + bnb_blasLt_heuristic_t heuristicResult[request_solutions]; + int returnedAlgoCount = 0; + checkBlasLtStatus(bnb_blasLtAlgoGetHeuristic( + ltHandle, matmulDesc, aDesc, bDesc, cDesc, cDesc, pref, request_solutions, heuristicResult, + &returnedAlgoCount + )); + checkBlasLtStatus(bnb_blasLtPrefDestroy(pref)); + + if (returnedAlgoCount == 0) { + 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 = {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(); + } + std::lock_guard lock(igemmlt_algo_cache_mutex); + igemmlt_algo_cache[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) + has_error = ERR_NOT_IMPLEMENTED; #else int alpha = 1, beta = 0; has_error |= checkBlasLtStatus(bnb_blasLtMatmul( From 6ec2ab19fd39810e1f4d80c593264cb287c8078d Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 16 Oct 2026 13:38:32 +0200 Subject: [PATCH 2/2] Share hipBLASLt int8 algos across token-count buckets Check algos with matmulIsAlgoSupported before launching them and key the cache on a bucket of n, so a new sequence length reuses an algo instead of running the heuristic again. An n the shared algo doesn't fit gets its own entry, and a cached algo that fails re-queries instead of erroring. Signed-off-by: Kashif Rasul --- csrc/ops.cu | 110 +++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 79 insertions(+), 31 deletions(-) diff --git a/csrc/ops.cu b/csrc/ops.cu index 26b1a86cb..f6f890455 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -16,6 +16,9 @@ #include #include #include +#ifndef NO_HIPBLASLT +#include +#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; +// (handle, m, k, lda, ldb, ldc, n or its bucket, whether the entry is for that exact n) +using igemmlt_algo_key = std::tuple; static std::map 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 @@ -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 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 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;