From 1b51a1a7c779282c003c7ee1e427072684c446bf Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Fri, 24 Jul 2026 15:55:02 +0530 Subject: [PATCH 1/3] fix: support runtime GEMM dimensions for dynamic shapes --- .../backends/cuda/sofieBLAS_cublas.hpp | 97 +++++++++---------- 1 file changed, 44 insertions(+), 53 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index f0b9e57..7063448 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -31,22 +31,6 @@ } while (0) -struct PairHash { - std::size_t - operator()(const std::pair &p) const noexcept { - std::size_t h1 = std::hash{}(p.first); - std::size_t h2 = std::hash{}(p.second); - return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); - } -}; - -struct PairEq { - bool operator()(const std::pair &a, - const std::pair &b) const noexcept { - return a.first == b.first && a.second == b.second; - } -}; - struct DescKey { int transA; // CUBLAS_OP_N / CUBLAS_OP_T encoded as int int transB; @@ -69,8 +53,8 @@ struct DescKeyHash { struct AlgoKey { DescKey dk; - std::size_t rowsA, colsA; // physical dimensions of A in layoutStore - std::size_t rowsB, colsB; // physical dimensions of B in layoutStore + std::size_t rowsA, colsA; // physical dimensions of A + std::size_t rowsB, colsB; // physical dimensions of B bool operator==(const AlgoKey &o) const noexcept { return dk == o.dk && rowsA == o.rowsA && colsA == o.colsA @@ -98,9 +82,12 @@ class BlasCuda { size_t workspaceSize = 1u << 25; // 32 MB (was 4 MB) cudaStream_t stream = nullptr; - std::unordered_map, - cublasLtMatrixLayout_t, PairHash, PairEq> - layoutStore; + // One persistent layout descriptor per matrix role, re-stamped with the + // runtime dimensions before each matmul. The descriptor is host-side metadata + // consumed by cublasLtMatmul at the call, so a single object can be reused + // across shapes - this is what lets one Session serve dynamic (runtime) sizes. + enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 }; + cublasLtMatrixLayout_t roleLayout[3] = {}; std::unordered_map descStore; @@ -129,8 +116,8 @@ class BlasCuda { } ~BlasCuda() { - for (auto &[key, layout] : layoutStore) - if (layout) cublasLtMatrixLayoutDestroy(layout); + for (auto L : roleLayout) + if (L) cublasLtMatrixLayoutDestroy(L); for (auto &[key, desc] : descStore) if (desc) cublasLtMatmulDescDestroy(desc); if (preference) cublasLtMatmulPreferenceDestroy(preference); @@ -149,22 +136,9 @@ class BlasCuda { } } - void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, - std::size_t lda, std::size_t ldb, std::size_t ldc, - char transa, char transb) { - // Physical A: (m×k) if NoTrans, (k×m) if Trans - if (transa == 'N' || transa == 'n') - checkAndAddLayout(m, k, lda); - else - checkAndAddLayout(k, m, lda); - // Physical B: (k×n) if NoTrans, (n×k) if Trans - if (transb == 'N' || transb == 'n') - checkAndAddLayout(k, n, ldb); - else - checkAndAddLayout(n, k, ldb); - // C is always (m×n) - checkAndAddLayout(m, n, ldc); - } + // No-op kept for the generated ctor's API; layouts are created lazily now. + void addLayoutConfig(std::size_t, std::size_t, std::size_t, + std::size_t, std::size_t, std::size_t, char, char) {} template inline void @@ -370,14 +344,25 @@ class BlasCuda { : std::make_pair(n, k); } - void checkAndAddLayout(std::size_t rows, std::size_t cols, std::size_t ld) { - auto key = std::make_pair(rows, cols); - if (layoutStore.find(key) == layoutStore.end()) { - cublasLtMatrixLayout_t layout = nullptr; - CHECK_CUBLAS( - cublasLtMatrixLayoutCreate(&layout, CUDA_R_32F, rows, cols, ld)); - layoutStore.emplace(key, layout); + // Resolve a matrix role's layout at the runtime dims: create the descriptor + // once, then overwrite its dims in place on later calls. ld = rows (dense, + // column-major, as the generated calls produce). + cublasLtMatrixLayout_t stampLayout(LayoutRole role, + const std::pair &key) { + const uint64_t rows = key.first, cols = key.second; + const int64_t ld = static_cast(key.first); + cublasLtMatrixLayout_t &L = roleLayout[role]; + if (!L) { + CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&L, CUDA_R_32F, rows, cols, ld)); + } else { + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_ROWS, &rows, sizeof(rows))); + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_COLS, &cols, sizeof(cols))); + CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute( + L, CUBLASLT_MATRIX_LAYOUT_LD, &ld, sizeof(ld))); } + return L; } cublasLtMatmulDesc_t &getOrCreateDesc(cublasOperation_t transA, @@ -421,12 +406,13 @@ class BlasCuda { return it->second; auto &desc = getOrCreateDesc(transA, transB, epilogue); + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); // C and D share the same layout cublasLtMatmulHeuristicResult_t h{}; int returnedResults = 0; CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( - ltHandle, desc, - layoutStore.at(kA), layoutStore.at(kB), - layoutStore.at(kC), layoutStore.at(kC), + ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); if (returnedResults == 0) { std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " @@ -459,12 +445,17 @@ class BlasCuda { &bias_ptr, sizeof(bias_ptr))); } + // Re-stamp the shared role layouts to this call's shape right before the + // matmul (the algo-cache hit path in getOrComputeAlgo skips stamping). + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); CHECK_CUBLAS(cublasLtMatmul( ltHandle, desc, - &alpha, A, layoutStore.at(kA), - B, layoutStore.at(kB), - &beta, D_in, layoutStore.at(kC), - C_out, layoutStore.at(kC), + &alpha, A, lA, + B, lB, + &beta, D_in, lC, + C_out, lC, &h.algo, d_workspace, workspaceSize, stream)); } }; From b4fc25e5eed2a0177617f4933360f04e525dd9a6 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Tue, 4 Aug 2026 16:26:59 +0530 Subject: [PATCH 2/3] feat: resolve GEMM algorithms per call site instead of per shape --- .../backends/cuda/sofieBLAS_cublas.hpp | 134 ++++++++++++++++-- 1 file changed, 119 insertions(+), 15 deletions(-) diff --git a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp index 7063448..324668c 100644 --- a/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp +++ b/include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp @@ -5,9 +5,11 @@ #include #include #include +#include #include #include #include +#include #include "sofieBLAS/core.hpp" #include @@ -74,6 +76,17 @@ struct AlgoKeyHash { } }; +// A call site's maximum shape, as declared by addLayoutConfig from the +// generated Session constructor. +struct ShapeEnvelope { + std::size_t rowsA, colsA, rowsB, colsB, rowsC, colsC; +}; + +struct LayoutStats { + std::size_t heuristicQueries = 0; + std::size_t envelopeRejects = 0; // envelope algorithm unusable at the call +}; + class BlasCuda { cublasLtHandle_t ltHandle = nullptr; cublasHandle_t handle = nullptr; // legacy cuBLAS for batched ops @@ -94,7 +107,15 @@ class BlasCuda { std::unordered_map algoCache; + // call-site envelopes declared by addLayoutConfig + std::vector envelopes; + + LayoutStats stats; + public: + const LayoutStats &layoutStats() const { return stats; } + std::size_t algoCacheSize() const { return algoCache.size(); } + BlasCuda(const BlasCuda &) = delete; BlasCuda &operator=(const BlasCuda &) = delete; BlasCuda(BlasCuda &&) = delete; @@ -136,9 +157,31 @@ class BlasCuda { } } - // No-op kept for the generated ctor's API; layouts are created lazily now. - void addLayoutConfig(std::size_t, std::size_t, std::size_t, - std::size_t, std::size_t, std::size_t, char, char) {} + // Records the call site's envelope. The generated constructor evaluates its + // shape expressions with its own parameters, so for a dynamic model these are + // the largest dims the call site will ever use. Layouts are created lazily, + // so nothing is registered here beyond the envelope and, with warmup on, the + // algorithm resolved for it. + void addLayoutConfig(std::size_t m, std::size_t n, std::size_t k, + std::size_t, std::size_t, std::size_t, + char transa, char transb) { + const auto kA = layoutKeyA(transa, m, k); + const auto kB = layoutKeyB(transb, k, n); + const std::pair kC{m, n}; + envelopes.push_back({kA.first, kA.second, kB.first, kB.second, m, n}); + + // The constructor does not know which epilogue this call site uses, so + // resolve all three. Unused ones cost one heuristic query each, off the + // inference path. + const cublasOperation_t tA = charToCuBlasTranspose(transa); + const cublasOperation_t tB = charToCuBlasTranspose(transb); + const cublasLtEpilogue_t eps[] = {CUBLASLT_EPILOGUE_DEFAULT, + CUBLASLT_EPILOGUE_BIAS, + CUBLASLT_EPILOGUE_RELU_BIAS}; + for (cublasLtEpilogue_t ep : eps) { + getOrComputeAlgo(tA, tB, ep, kA, kB, kC, /*required=*/false); + } + } template inline void @@ -365,6 +408,32 @@ class BlasCuda { return L; } + // Tightest declared envelope covering this call, or null if none does. + // Tightest matters: several envelopes may cover a small shape, but only the + // call site's own matches its weight dims exactly and so has zero excess + // on those axes. + const ShapeEnvelope * + findEnvelope(const std::pair &kA, + const std::pair &kB, + const std::pair &kC) const { + const ShapeEnvelope *best = nullptr; + std::size_t bestExcess = std::numeric_limits::max(); + for (const auto &e : envelopes) { + // colsA and rowsB are both the contraction dimension k, which comes from + // the weight tensor and never varies at runtime. Requiring an exact match + // on it stops one call site's envelope from serving another's shapes. + if (e.colsA != kA.second || e.rowsB != kB.first) + continue; + if (e.rowsA < kA.first || e.colsB < kB.second || + e.rowsC < kC.first || e.colsC < kC.second) + continue; + const std::size_t ex = (e.rowsA - kA.first) + (e.colsA - kA.second) + + (e.rowsB - kB.first) + (e.colsB - kB.second); + if (ex < bestExcess) { bestExcess = ex; best = &e; } + } + return best; + } + cublasLtMatmulDesc_t &getOrCreateDesc(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue) { @@ -393,17 +462,37 @@ class BlasCuda { return descStore.at(key); } - cublasLtMatmulHeuristicResult_t & + // Whether an algorithm can actually run this shape. cuBLASLt rejects some + // combinations, so an algorithm resolved at a call site's envelope is not + // guaranteed to work at every smaller shape it serves. + bool algoUsable(cublasLtMatmulDesc_t desc, const cublasLtMatmulAlgo_t &algo, + const std::pair &kA, + const std::pair &kB, + const std::pair &kC) { + auto lA = stampLayout(ROLE_A, kA); + auto lB = stampLayout(ROLE_B, kB); + auto lC = stampLayout(ROLE_C, kC); + cublasLtMatmulHeuristicResult_t chk{}; + return cublasLtMatmulAlgoCheck(ltHandle, desc, lA, lB, lC, lC, &algo, + &chk) == CUBLAS_STATUS_SUCCESS && + chk.workspaceSize <= workspaceSize; + } + + // required=false is used by constructor warmup, which speculatively resolves + // epilogues the call site may never use: those may legitimately have no + // algorithm and must not abort. + cublasLtMatmulHeuristicResult_t * getOrComputeAlgo(cublasOperation_t transA, cublasOperation_t transB, cublasLtEpilogue_t epilogue, const std::pair &kA, const std::pair &kB, - const std::pair &kC) { + const std::pair &kC, + bool required = true) { AlgoKey key{{(int)transA, (int)transB, (int)epilogue}, kA.first, kA.second, kB.first, kB.second}; auto it = algoCache.find(key); if (it != algoCache.end()) - return it->second; + return &it->second; auto &desc = getOrCreateDesc(transA, transB, epilogue); auto lA = stampLayout(ROLE_A, kA); @@ -414,7 +503,10 @@ class BlasCuda { CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic( ltHandle, desc, lA, lB, lC, lC, preference, 1, &h, &returnedResults)); + ++stats.heuristicQueries; if (returnedResults == 0) { + if (!required) + return nullptr; std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for " << "transA=" << transA << " transB=" << transB << " epilogue=" << epilogue @@ -422,8 +514,7 @@ class BlasCuda { << " B=[" << kB.first << "x" << kB.second << "]\n"; exit(EXIT_FAILURE); } - algoCache.emplace(key, h); - return algoCache.at(key); + return &algoCache.emplace(key, h).first->second; } void executeMatmul(cublasOperation_t transA, cublasOperation_t transB, @@ -434,10 +525,6 @@ class BlasCuda { const std::pair &kA, const std::pair &kB, const std::pair &kC) { - // Retrieve (or lazily compute) the cached algorithm for this shape - auto &h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); - - // Retrieve the cached descriptor and patch the real bias pointer in-place auto &desc = getOrCreateDesc(transA, transB, epilogue); if (bias_ptr) { CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( @@ -445,8 +532,25 @@ class BlasCuda { &bias_ptr, sizeof(bias_ptr))); } - // Re-stamp the shared role layouts to this call's shape right before the - // matmul (the algo-cache hit path in getOrComputeAlgo skips stamping). + // Resolve at this call site's declared envelope, so every runtime size it + // produces shares one cache entry. + const ShapeEnvelope *env = findEnvelope(kA, kB, kC); + const std::pair + aA = env ? std::make_pair(env->rowsA, env->colsA) : kA, + aB = env ? std::make_pair(env->rowsB, env->colsB) : kB, + aC = env ? std::make_pair(env->rowsC, env->colsC) : kC; + auto *h = getOrComputeAlgo(transA, transB, epilogue, aA, aB, aC); + + // Fall back to the exact shape when the envelope's algorithm cannot run + // it. cuBLASLt returns CUBLAS_STATUS_NOT_SUPPORTED for at least some + // shape/algorithm combinations; m=1 was the first observed. + if (env && !algoUsable(desc, h->algo, kA, kB, kC)) { + ++stats.envelopeRejects; + h = getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC); + } + + // Re-stamp the shared role layouts to this call's exact shape. Resolving + // the algorithm leaves them at the envelope size, so this happens last. auto lA = stampLayout(ROLE_A, kA); auto lB = stampLayout(ROLE_B, kB); auto lC = stampLayout(ROLE_C, kC); @@ -456,7 +560,7 @@ class BlasCuda { B, lB, &beta, D_in, lC, C_out, lC, - &h.algo, d_workspace, workspaceSize, stream)); + &h->algo, d_workspace, workspaceSize, stream)); } }; From 21f938da4eca43e2388a4a9f9e93137fd2b19ad0 Mon Sep 17 00:00:00 2001 From: Harsh Chauhan Date: Tue, 4 Aug 2026 16:26:59 +0530 Subject: [PATCH 3/3] test: cover multi-size use of one instance and bounded algorithm cache --- tests/test.cc | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) diff --git a/tests/test.cc b/tests/test.cc index e106275..0f29561 100644 --- a/tests/test.cc +++ b/tests/test.cc @@ -567,6 +567,95 @@ static void runCudaTests() { } } +// One instance used at many sizes, which is what a dynamic-shape model does. +// earlier, any size other than the one registered at construction threw std::out_of_range from the layout lookup +// so this is the case the rest of the suite never covered. +static void runDynamicShapeTests() { + std::cout << "\n=== CUDA Dynamic-Shape Tests ===\n"; + + alpaka::PlatformCudaRt platform{}; + auto dev = alpaka::getDevByIdx(platform, 0u); + alpaka::Queue queue{dev}; + sofieBLAS blas(queue); + + alpaka::PlatformCpu hostPlatform{}; + auto hostDev = alpaka::getDevByIdx(hostPlatform, 0u); + + // Buffers hold MCAP rows while the call site declares MENV, so sizes above + // MENV stay in bounds and exercise the path where no envelope covers them. + constexpr int MCAP = 96, MENV = 64, N = 3, K = 5; + + auto hA = alpaka::allocBuf(hostDev, static_cast(MCAP * K)); + auto hB = alpaka::allocBuf(hostDev, static_cast(K * N)); + auto hC = alpaka::allocBuf(hostDev, static_cast(MCAP * N)); + float *A = alpaka::getPtrNative(hA); + float *B = alpaka::getPtrNative(hB); + float *C = alpaka::getPtrNative(hC); + fillSeq(A, MCAP * K, 0.5f, 0.25f); + fillSeq(B, K * N, 1.f, 0.5f); + + auto dA = alpaka::allocAsyncBuf(queue, static_cast(MCAP * K)); + auto dB = alpaka::allocAsyncBuf(queue, static_cast(K * N)); + auto dC = alpaka::allocAsyncBuf(queue, static_cast(MCAP * N)); + alpaka::memcpy(queue, dA, hA); + alpaka::memcpy(queue, dB, hB); + alpaka::wait(queue); + + // Declare the call site's largest shape, as a generated Session constructor + // does with its own arguments. + blas.addLayoutConfig(MENV, N, K, ldaFor('N', MENV, K), ldbFor('N', K, N), + MENV, 'N', 'N'); + + std::vector ref; + auto runAt = [&](int m, const std::string &name) { + ref.assign(static_cast(m) * N, 0.f); + refMatmul(ref.data(), A, B, m, N, K, 1.f, 0.f, false, false); + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::memcpy(queue, hC, dC); + alpaka::wait(queue); + checkClose(C, ref.data(), m * N, name); + }; + + // At, below and above the declared envelope, all on one instance. + for (int m : {MENV, 37, 8, 51, 1, MENV, MCAP}) + runAt(m, "cuda::dynamic m=" + std::to_string(m)); + + // The cache must not grow one entry per size. Entries are added only where + // the envelope's algorithm is rejected. Keying by shape fails only here. + const int nSizes = MENV - 1; + const std::size_t cacheBefore = blas.algoCacheSize(); + const std::size_t rejBefore = blas.layoutStats().envelopeRejects; + const std::size_t searchBefore = blas.layoutStats().heuristicQueries; + for (int m = 2; m <= MENV; ++m) + blas.matmul('N', 'N', static_cast(m), static_cast(N), + static_cast(K), 1.f, dA, dB, 0.f, dC); + alpaka::wait(queue); + const std::size_t added = blas.algoCacheSize() - cacheBefore; + const std::size_t rejected = blas.layoutStats().envelopeRejects - rejBefore; + + std::cout << " " << nSizes << " sizes added " << added + << " cache entries, " << rejected << " rejected\n"; + if (added < static_cast(nSizes)) { + std::cout << " PASS cuda::cache bounded\n"; + } else { + std::cerr << " FAIL [cuda::cache bounded] one entry per size\n"; + ++gFailures; + } + + // The constructor resolves every declared envelope, so a size it covers must + // not trigger a search. Only a rejected envelope algorithm may. + const std::size_t searched = + blas.layoutStats().heuristicQueries - searchBefore; + if (searched <= rejected) { + std::cout << " PASS cuda::no search during inference\n"; + } else { + std::cerr << " FAIL [cuda::no search during inference] " << searched + << " searches, only " << rejected << " explained by rejects\n"; + ++gFailures; + } +} + #endif // ALPAKA_ACC_GPU_CUDA_ENABLED // --------------------------------------------------------------------------- @@ -579,6 +668,7 @@ int main() { #endif #ifdef ALPAKA_ACC_GPU_CUDA_ENABLED runCudaTests(); + runDynamicShapeTests(); #endif std::cout << "\n";