Support runtime GEMM dimensions for dynamic shapes#11
Open
harz05 wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
cublasLtMatrixLayoutSetAttributebefore 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
inferovern_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 viaSetAttribute. 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.