Skip to content

Support runtime GEMM dimensions for dynamic shapes#11

Open
harz05 wants to merge 1 commit into
ML4EP:devfrom
harz05:fix/dynamic-gemm-layouts
Open

Support runtime GEMM dimensions for dynamic shapes#11
harz05 wants to merge 1 commit into
ML4EP:devfrom
harz05:fix/dynamic-gemm-layouts

Conversation

@harz05

@harz05 harz05 commented Jul 24, 2026

Copy link
Copy Markdown

Implements #10
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. addLayoutConfig becomes a no-op; the algorithm cache is unchanged (keyed by shape, re-queried per new size).

Approaches tried

Two ways to handle runtime dimensions were tried and benchmarked on ParticleNet onnx model, running one Session and sweeping infer over n_pf = 100, 80, 60, 40, 20 on a T4. The first was register-on-miss: create an immutable layout per distinct (rows, cols) on a miss and cache it in the map. The second was mutate-in-place (this PR): keep one layout per role and re-stamp the runtime dims via SetAttribute. Both give correct results (a construct-at-A / infer-at-B GEMM matches a CPU reference to ~1e-5), and they perform very close. mutate-in-place is marginally faster and a bit lighter on host memory.

image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant