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-28T14:21:18Z

+  using Operation = hipblasOperation_t;
+  using Epilogue = hipblasLtEpilogue_t;
+
+  static constexpr auto OpN = HIPBLAS_OP_N;


constexpr benchamark

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


		## 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.

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!

Choose a reason for hiding this comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

harz05 commented Jul 24, 2026 •

edited by sanjibansg

Loading

harz05 commented Aug 18, 2026 •

edited

Loading

harz05 commented Aug 28, 2026 •

edited

Loading