[Inductor cutlass backend] Enable StreamK and nonzero workspace sizes #119005
Conversation
This file contains 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
This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/119005
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 New FailuresAs of commit 8019b3a with merge base 520771d ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Feb 2, 2024
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
…space sizes" This change introduces temporary workspace buffers, which can be used to provide workspace memory to Inductor ops. This is a required feature in order to enable the Cutlass StreamK tile scheduler, which should bring performance benefits and more reliable GEMM performance for many tile shapes. See StreamK paper: https://arxiv.org/abs/2301.03598 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 muchulee8 aakhundov ColinPeppler [ghstack-poisoned]
peterbell10
reviewed
Feb 6, 2024
| @@ -3028,6 +3028,41 @@ class InputBuffer(Buffer): | |||
| pass | |||
|
|
|||
|
|
|||
| @dataclasses.dataclass | |||
| class WorkspaceBuffer(Buffer): | |||
There was a problem hiding this comment.
In #117992 I also add workspace allocations but take a slightly different approach. The workspace size is an implementation detail of the code generation, so I don't represent it in the inductor IR and instead put it in KernelArgs and have the Kernel method create the args.
wdyt?
There was a problem hiding this comment.
@peterbell10 - Just looked at it, looks good to me. If the PR you mention is blocked through this ROCm issue (Is that blocking?),maybe you could split the workspace allocation code off into a separate PR, so we can get it merged quickly and I can adapt this to it.
There was a problem hiding this comment.
ROCm issues are fixed and I'm just waiting on perf results before merging.
|
Had to be squashed into #119007 |
OnlyFor
pushed a commit
to OnlyFor/pytorch
that referenced
this pull request
May 3, 2024
Enable nonzero workspace and Cutlass StreamK for Inductor Cutlass GEMM ops. This is a simpler rewrite of my original version of pytorch#119005 using peterbell10 's workspace allocation mechanism from pytorch#117992 Test Plan: - Additional unit test in test_cutlass_backend.py which specifically tests StreamK GEMM with workspace requirement - CI ghstack-source-id: 24d06299f90a1e31af6b097316b76689e4944df2 Pull Request resolved: pytorch#125406
OnlyFor
pushed a commit
to OnlyFor/pytorch
that referenced
this pull request
May 4, 2024
Enable nonzero workspace and Cutlass StreamK for Inductor Cutlass GEMM ops. This is a simpler rewrite of my original version of pytorch#119005 using peterbell10 's workspace allocation mechanism from pytorch#117992 Test Plan: - Additional unit test in test_cutlass_backend.py which specifically tests StreamK GEMM with workspace requirement - CI ghstack-source-id: 6b2a29b3d2754b1981b503939f79f7bc5889216e Pull Request resolved: pytorch#125406
pytorchmergebot
pushed a commit
that referenced
this pull request
May 5, 2024
…amK (#125406) Enable nonzero workspace and Cutlass StreamK for Inductor Cutlass GEMM ops. This is a simpler rewrite of my original version of #119005 using @peterbell10 's workspace allocation mechanism from #117992 Test Plan: - Additional unit test in test_cutlass_backend.py which specifically tests StreamK GEMM with workspace requirement - CI Pull Request resolved: #125406 Approved by: https://github.com/jansel
kadeng
added a commit
that referenced
this pull request
May 5, 2024
Enable nonzero workspace and Cutlass StreamK for Inductor Cutlass GEMM ops. This is a simpler rewrite of my original version of #119005 using peterbell10 's workspace allocation mechanism from #117992 Test Plan: - Additional unit test in test_cutlass_backend.py which specifically tests StreamK GEMM with workspace requirement - CI ghstack-source-id: eccd80aedce633b345a88eb25ca1d1149bb12756 Pull Request resolved: #125406
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.
Stack from ghstack (oldest at bottom):
This change introduces temporary workspace buffers, which can be used to
provide workspace memory to Inductor ops. This is a required feature
in order to enable the Cutlass StreamK tile scheduler, which should
bring performance benefits and more reliable GEMM performance for
many tile shapes.
See StreamK paper: https://arxiv.org/abs/2301.03598
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @muchulee8 @aakhundov @ColinPeppler