Support runtime GEMM dimensions for dynamic shapes by harz05 · Pull Request #11 · ML4EP/sofieBLAS

harz05 · 2026-07-24T10:46:37Z

Implements #10

Edit:

The PR has been brought up to date with the dev branch and the approach has grown from the older description, described below.

The layout part is unchanged: one cublasLtMatrixLayout_t per role, runtime dims stamped in before each multiply operation. The change is on the algorithm caching. Re-querying the heuristic per new size meant a dynamic model pays one heuristic query for every distinct size it meets and the cache grows with the sweep. Now addLayoutConfig is no longer a no-op: it records the call site's construct-time shape as its envelope and resolves the algorithm for it once, up front. At infer time any size covered by an envelope reuses that entry; so a size sweep does zero heuristic queries and adds zero cache entries. In any case if cuBLASLt cannot run the envelope's algorithm at the actual size the call falls back to resolving that exact shape.
LayoutStats and algoCacheSize() expose the counts that the tests assert through them.

The cache is unbounded by default. Passing a limit to the constructor enables LRU eviction, for the case where many above-envelope shapes get resolved individually.

Benchmark (Nvidia H100) through exisiting, bench_cuda:

Original description:

sofieBLAS registers cuBLASLt matrix layouts at construct size (addLayoutConfig) and the matmul path looks them up by the runtime (rows, cols). With dynamic shapes the two disagree i.e. a Session constructed at one size and inferred at another misses the layout and throws std::out_of_range, so one Session only works at a single size.

This PR resolves each matrix's layout at matmul time instead: keep one cublasLtMatrixLayout_t per role (A/B/C) and stamp the runtime dims into it via cublasLtMatrixLayoutSetAttribute before each multiply. The descriptor is host-side metadata consumed by cublasLtMatmul at the call, so reusing one object across shapes is safe.

sanjibansg · 2026-08-18T07:43:07Z

/runtest h100

github-actions · 2026-08-18T07:43:18Z

/runtest (h100): triggered - view run

sanjibansg · 2026-08-18T07:43:45Z

/runtest mi300x

github-actions · 2026-08-18T07:43:54Z

/runtest (mi300x): triggered - view run

github-actions · 2026-08-18T07:44:57Z

/runtest (h100): GPU Unit Tests ✅ passed - view run

sanjibansg · 2026-08-18T08:53:16Z

/runtest w7900

github-actions · 2026-08-18T08:53:26Z

/runtest (w7900): triggered - view run

github-actions · 2026-08-18T08:55:00Z

/runtest (w7900): GPU Unit Tests ✅ passed - view run

sanjibansg

LGTM!
One last comment, can you add a bit of documentation mentioning details like how to initialize the cache, etc.

harz05 · 2026-08-18T11:35:51Z

As requested, added a README section covering the dynamic-shape behavior and how to initialize the cache limit
@sanjibansg

harz05 · 2026-08-19T09:19:34Z

hipBLASLt was missing the same implementation that was done for cuBLASLt, so this PR includes that as well. The only difference is that hipBLASLt has no hipblasLtMatmulAlgoCheck, so the exact-shape fallback uses hipblaslt_ext::matmulIsAlgoSupported instead.

Tests pass on AMD, MI100

sanjibansg

Some more comments on more descriptive methods and algorithmic reasoning. Comments mostly addresses the cuBLASLt code, but stand similar for the hipBLASLt ones as well.

sanjibansg · 2026-08-19T09:28:14Z

+
+## Dynamic GEMM shapes and the algorithm cache
+
+Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply.


Using Both GPU Backends is ambiguous, better state clearly what it infers here.

sanjibansg · 2026-08-19T09:31:19Z

+
+Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply.
+
+Algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution.


This paragraph is not very clear as to what it addresses. We need to state clearly what Algorithm selection here means, what envelope here is, etc. Code documentation are better descriptive of the algorithms, features, and declarations it entails

sanjibansg · 2026-08-19T18:31:27Z

+    if (!L) {
+      CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&L, CUDA_R_32F, rows, cols, ld));
+    } else {
+      CHECK_CUBLAS(cublasLtMatrixLayoutSetAttribute(


what if there is only a change in rows, but the other fields remain unaffected, do we still need to update all the attributes?

Right now all three attributes are rewritten on every call even when nothing changed. I'm planning to keep track of the dimensions each descriptor currently holds and update only what differs i.e. a change in rows updates rows and ld together (ld=rows), a change in cols updates only cols and a repeated shape writes nothing.

However the above approach will only avoid rewrites when the same size repeats in succession, so a better approach could be that we give each call site declared through addLayoutConfig its own three descriptors instead of sharing one per matrix role. Sites then stop overwriting each other, so when sizes repeat across events nothing is written at all and the ordering constraint between the validity check and the final stamp mostly goes away; so I think we can try this instead

sanjibansg · 2026-08-19T18:32:18Z

+  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 {


explain what this method is exactly doing

sanjibansg · 2026-08-19T18:35:01Z

+
+    if (env && !algoUsable(desc, h.algo, kA, kB, kC)) {
+      ++stats.envelopeRejects;
+      h = *getOrComputeAlgo(transA, transB, epilogue, kA, kB, kC);


why do we need to compute the algo twice here?

it's not computed twice in the normal path. The second resolution happens only when cuBLASLt rejects the declared shape's algorithm at this call's exact size (returns NOT_SUPPORTED, m=1 was found to do this for cuBLASLt) and then it falls back to resolving at the exact shape, which also gets cached. Added a comment explaining the same

sanjibansg · 2026-08-19T18:35:46Z

-  std::unordered_map<std::pair<std::size_t, std::size_t>,
-                     hipblasLtMatrixLayout_t, PairHash, PairEq>
-      layoutStore;
+  enum LayoutRole { ROLE_A = 0, ROLE_B = 1, ROLE_C = 2 };


add a comment as to what LayoutRole means and what those individual roles mean here.

harz05 · 2026-08-28T19:00:27Z

@sanjibansg
Following are the benchmarking results for the new changes on both CUDA and HIP (using static constexpr members for cuda and hip)-

All the existing tests present in dev are also passing

Nvidia: H100
AMD: MI100

sanjibansg · 2026-09-02T14:06:04Z

/runtest h100-47gb

github-actions · 2026-09-02T14:06:15Z

/runtest (h100-47gb): triggered - view run

sanjibansg · 2026-09-02T14:06:26Z

/runtest mi300x

github-actions · 2026-09-02T14:06:38Z

/runtest (mi300x): triggered - view run

harz05 · 2026-09-04T11:43:00Z

Nvidia H100

harz05 · 2026-09-04T11:44:15Z

AMD MI100

sanjibansg · 2026-09-04T13:12:27Z

/runtest w7900

github-actions · 2026-09-04T13:12:36Z

/runtest (w7900): triggered - view run

sanjibansg · 2026-09-04T13:13:49Z

/runtest l40s

github-actions · 2026-09-04T13:13:53Z

/runtest (w7900): GPU Unit Tests ✅ passed - view run

github-actions · 2026-09-04T13:13:58Z

/runtest (l40s): triggered - view run

sanjibansg · 2026-09-06T08:32:42Z

/runtest h100-47gb

github-actions · 2026-09-06T08:32:51Z

/runtest (h100-47gb): triggered - view run

sanjibansg

Very close, couple of more comments and the PR is ready to be merged! thanks for the efforts!

sanjibansg · 2026-09-06T08:21:30Z

+// The cuBLASLt spellings of everything the shared BlasLt implementation in
+// backends/gpu/detail uses: types, enum values and functions.


Suggested change

// The cuBLASLt spellings of everything the shared BlasLt implementation in

// backends/gpu/detail uses: types, enum values and functions.

// The cuBLASLt forwarding of the shared BlasLt implementation in backends/gpu/detail

sanjibansg · 2026-09-06T08:23:33Z

-                    static_cast<std::size_t>(k.epilogue);
-    return h ^ (h >> 16);
-  }
+// The hipBLASLt spellings of everything the shared BlasLt implementation in


Suggested change

// The hipBLASLt spellings of everything the shared BlasLt implementation in

// The hipBLASLt forwarding of the shared BlasLt implementation in

sanjibansg · 2026-09-06T08:26:37Z

 ```

-The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addLayoutConfig` (used to pre-register cuBLASLt/hipBLASLt matrix layouts for a given shape before the first `matmul`/`gemm` call on that shape).
+The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addOperationConfig` (creates the matrix layouts and resolves the multiply algorithm for a call site's shape ahead of its first call, see below).


list these out in bullets to make them easier to understand

Suggested change

The GPU backends (`BlasCuda`, `BlasHip`) additionally expose `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues), `gemmStridedBatched`, and `addOperationConfig` (creates the matrix layouts and resolves the multiply algorithm for a call site's shape ahead of its first call, see below).

The GPU backends (`BlasCuda`, `BlasHip`) additionally expose

- `gemmrelu`/`gemmgelu` (fused bias + activation via cuBLASLt/hipBLASLt epilogues)

- `gemmStridedBatched` for batched gemm operations through strides

- `addOperationConfig` that creates the matrix layouts and resolves the multiply algorithm for a call site's shape ahead of its first call, see below).

sanjibansg · 2026-09-06T08:28:22Z

+### Warming the cache with addOperationConfig
+
+`addOperationConfig(m, n, k, lda, ldb, ldc, transa, transb, epilogue)` fills the cache for one call site ahead of its first call: it creates the three layouts and resolves the algorithm for the given dimensions. The `epilogue` argument is the `Epilogue` enum from `sofieBLAS/core.hpp` and names which call the site will make, because the fused epilogue is part of the selected kernel:


I would not mention warmup step here, since this is more inference workflow related, there could be situations where a warmup is not possible in its traditional sense, sofieBLAS works even in that case by just caching in the information in its first call.

sanjibansg · 2026-09-06T08:29:09Z

+| `Epilogue::ReluBias` | `gemmrelu` (bias, then ReLU) |
+| `Epilogue::GeluBias` | `gemmgelu` (bias, then GELU) |
+
+A generated Session constructor calls it once per GEMM call site with the construction-time dimensions, so the first inference pays no heuristic queries at those sizes.


this is SOFIE-related text, I would not mention this here. sofieBLAS should work independent of SOFIE, therefore could be utilized in other projects as well

sanjibansg · 2026-09-06T08:31:59Z

+blas.matmul('N', 'N', 37, 3, 5, 1.0f, dA, dB, 0.0f, dC); // new size: created on first use
+```
+
+### One implementation for both GPU backends


this comment is confusing and could hint that we need to instantiate blas methods differently for CPU and GPU backends, which is not correct. The BLAS methods are instantiated in identical ways with only difference being the template dispatch variable.

github-actions · 2026-09-06T08:34:34Z

/runtest (h100-47gb): GPU Unit Tests ✅ passed - view run

harz05 · 2026-09-06T13:41:10Z

@sanjibansg thanks for the review, i've addressed the comments and made changes

harz05 added 3 commits July 24, 2026 15:55

fix: support runtime GEMM dimensions for dynamic shapes

1b51a1a

feat: resolve GEMM algorithms per call site instead of per shape

b4fc25e

test: cover multi-size use of one instance and bounded algorithm cache

21f938d

harz05 force-pushed the fix/dynamic-gemm-layouts branch from 6edd3e3 to 21f938d Compare August 4, 2026 10:58

harz05 added 4 commits August 17, 2026 19:14

Merge upstream/dev into fix/dynamic-gemm-layouts

ddc773e

feat: optional LRU limit on the algorithm cache

5c45651

style: clang-format

09d3123

chore: cleanup, cache limit via constructor

e5b362c

sanjibansg approved these changes Aug 18, 2026

View reviewed changes

docs: layout and caching behavior

d71ddf5

harz05 requested a review from sanjibansg August 18, 2026 11:38

harz05 mentioned this pull request Aug 18, 2026

Initial dynamic input support for gpu codegen ML4EP/SOFIE#45

Open

feat: dynamic gemm shapes and algorithm cache for hipBLASLt

1b0ea97

sanjibansg requested changes Aug 19, 2026

View reviewed changes

harz05 added 3 commits August 20, 2026 21:31

docs: describe layout roles, envelope matching and shape naming

11ff88e

docs: explain addLayoutConfig, algo resolution and stats counters

9816590

fix: cache layouts and algorithms, one shared BLASLt backend

4bcb1e2

harz05 commented Aug 28, 2026

View reviewed changes

Comment thread include/sofieBLAS/backends/hip/sofieBLAS_hipblaslt.hpp

harz05 added 2 commits September 4, 2026 15:58

feat: addOperationConfig with a epilogue, tests for runtime sizes

8b6c3e5

docs: describe addOperationConfig, the epilogue enum and the cache limit

d639575

sanjibansg requested changes Sep 6, 2026

View reviewed changes

sanjibansg mentioned this pull request Sep 6, 2026

[tests] remove cuda/hip test duplication with a templated fixture #13

Open

docs: address review comments

cd5a4e0


		## Dynamic GEMM shapes and the algorithm cache

		Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply.


		Both GPU backends behave the same way here. One instance serves GEMM calls at sizes that vary at runtime: matrix layouts are not tied to a shape, and each call stamps its dimensions into a shared per-role descriptor right before the multiply.

		Algorithm selection is cached. `addLayoutConfig(m, n, k, lda, ldb, ldc, transa, transb)` declares the largest shape a call site will use (its envelope) and resolves the algorithm for it once, up front. Any later call at a size covered by an envelope reuses that entry, so sweeping sizes does not re-query the heuristic or grow the cache. A call with no covering envelope is resolved at its exact shape and cached per shape, and if an envelope's algorithm cannot run a particular size the call falls back to exact-shape resolution.

		// The cuBLASLt spellings of everything the shared BlasLt implementation in
		// backends/gpu/detail uses: types, enum values and functions.

	// The cuBLASLt spellings of everything the shared BlasLt implementation in
	// backends/gpu/detail uses: types, enum values and functions.
	// The cuBLASLt forwarding of the shared BlasLt implementation in backends/gpu/detail

	// The hipBLASLt spellings of everything the shared BlasLt implementation in
	// The hipBLASLt forwarding of the shared BlasLt implementation in

		### Warming the cache with addOperationConfig

		`addOperationConfig(m, n, k, lda, ldb, ldc, transa, transb, epilogue)` fills the cache for one call site ahead of its first call: it creates the three layouts and resolves the algorithm for the given dimensions. The `epilogue` argument is the `Epilogue` enum from `sofieBLAS/core.hpp` and names which call the site will make, because the fused epilogue is part of the selected kernel:

Uh oh!

Conversation

harz05 commented Jul 24, 2026 • edited by sanjibansg Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Edit:

Original description:

Uh oh!

sanjibansg commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

sanjibansg commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

sanjibansg commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

sanjibansg left a comment

Choose a reason for hiding this comment

Uh oh!

harz05 commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

harz05 commented Aug 19, 2026

Uh oh!

sanjibansg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

harz05 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sanjibansg commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

sanjibansg commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

harz05 commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

harz05 commented Sep 4, 2026

Uh oh!

sanjibansg commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

sanjibansg commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

sanjibansg commented Sep 6, 2026

harz05 commented Jul 24, 2026 •

edited by sanjibansg

Loading

harz05 commented Aug 18, 2026 •

edited

Loading

harz05 commented Aug 28, 2026 •

edited

Loading

harz05 commented Sep 4, 2026 •

edited

Loading

sanjibansg Sep 6, 2026 •

edited

Loading