Skip to content

Support runtime GEMM dimensions for dynamic shapes #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
harz05 wants to merge 3 commits into
base: dev
Choose a base branch
from
Open

Support runtime GEMM dimensions for dynamic shapes #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1b51a1a
fix: support runtime GEMM dimensions for dynamic shapes
harz05 Jul 24, 2026
b4fc25e
feat: resolve GEMM algorithms per call site instead of per shape
harz05 Aug 4, 2026
21f938d
test: cover multi-size use of one instance and bounded algorithm cache
harz05 Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 155 additions & 60 deletions include/sofieBLAS/backends/cuda/sofieBLAS_cublas.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,11 @@
#include <cstdlib>
#include <functional>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <unordered_map>
#include <utility>
#include <vector>

#include "sofieBLAS/core.hpp"
#include <alpaka/alpaka.hpp>
Expand All @@ -31,22 +33,6 @@
} while (0)


struct PairHash {
std::size_t
operator()(const std::pair<std::size_t, std::size_t> &p) const noexcept {
std::size_t h1 = std::hash<std::size_t>{}(p.first);
std::size_t h2 = std::hash<std::size_t>{}(p.second);
return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
}
};

struct PairEq {
bool operator()(const std::pair<std::size_t, std::size_t> &a,
const std::pair<std::size_t, std::size_t> &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;
Expand All @@ -69,8 +55,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
Expand All @@ -90,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
Expand All @@ -98,16 +95,27 @@ class BlasCuda {
size_t workspaceSize = 1u << 25; // 32 MB (was 4 MB)
cudaStream_t stream = nullptr;

std::unordered_map<std::pair<std::size_t, std::size_t>,
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<DescKey, cublasLtMatmulDesc_t, DescKeyHash> descStore;

std::unordered_map<AlgoKey, cublasLtMatmulHeuristicResult_t, AlgoKeyHash>
algoCache;

// call-site envelopes declared by addLayoutConfig
std::vector<ShapeEnvelope> 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;
Expand All @@ -129,8 +137,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);
Expand All @@ -149,21 +157,30 @@ class BlasCuda {
}
}

// 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 lda, std::size_t ldb, std::size_t ldc,
std::size_t, std::size_t, std::size_t,
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);
const auto kA = layoutKeyA(transa, m, k);
const auto kB = layoutKeyB(transb, k, n);
const std::pair<std::size_t, std::size_t> 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 <typename T, typename TIdx>
Expand Down Expand Up @@ -370,14 +387,51 @@ 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<std::size_t, std::size_t> &key) {
const uint64_t rows = key.first, cols = key.second;
const int64_t ld = static_cast<int64_t>(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;
}

// 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<std::size_t, std::size_t> &kA,
const std::pair<std::size_t, std::size_t> &kB,
const std::pair<std::size_t, std::size_t> &kC) const {
const ShapeEnvelope *best = nullptr;
std::size_t bestExcess = std::numeric_limits<std::size_t>::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,
Expand Down Expand Up @@ -408,36 +462,59 @@ 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<std::size_t, std::size_t> &kA,
const std::pair<std::size_t, std::size_t> &kB,
const std::pair<std::size_t, std::size_t> &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<std::size_t, std::size_t> &kA,
const std::pair<std::size_t, std::size_t> &kB,
const std::pair<std::size_t, std::size_t> &kC) {
const std::pair<std::size_t, std::size_t> &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);
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));
++stats.heuristicQueries;
if (returnedResults == 0) {
if (!required)
return nullptr;
std::cerr << "[sofieBLAS] No suitable cuBLASLt algorithm found for "
<< "transA=" << transA << " transB=" << transB
<< " epilogue=" << epilogue
<< " A=[" << kA.first << "x" << kA.second << "]"
<< " 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,
Expand All @@ -448,24 +525,42 @@ class BlasCuda {
const std::pair<std::size_t, std::size_t> &kA,
const std::pair<std::size_t, std::size_t> &kB,
const std::pair<std::size_t, std::size_t> &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(
desc, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
&bias_ptr, sizeof(bias_ptr)));
}

// 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<std::size_t, std::size_t>
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);
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),
&h.algo, d_workspace, workspaceSize, stream));
&alpha, A, lA,
B, lB,
&beta, D_in, lC,
C_out, lC,
&h->algo, d_workspace, workspaceSize, stream));
}
};

Expand Down
Loading