Uh oh!
There was an error while loading. Please reload this page.
Use 4B vector loads/stores in cast-transpose kernel for small matrices - #101
Merged
Conversation
Signed-off-by: Tim Moon <tmoon@nvidia.com>
timmoon10force-pushed
the
small-cast-transpose
branch
from
March 14, 2023 23:01
c2a95c5 to
3a931adComparetimmoon10
commented
Mar 14, 2023
MemberAuthor
/te-ci |
timmoon10
commented
Mar 17, 2023
MemberAuthor
/te-ci |
ptrendx
approved these changes
Mar 17, 2023
ptrendx
left a comment
Member
There was a problem hiding this comment.
LGTM with this small bandaid and revisit this (also for the other functions like regular transpose etc.) when we move to runtime compilation.
cyanguwa pushed a commit
to cyanguwa/TransformerEngine
that referenced
this pull request
Mar 31, 2023
NVIDIA#101) Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Charlene Yang <charleney@nvidia.com>
cyanguwa pushed a commit
to cyanguwa/TransformerEngine
that referenced
this pull request
Apr 1, 2023
NVIDIA#101) Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Charlene Yang <charleney@nvidia.com>
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 freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
@vasunvidia has observed poor performance when running the cast-transpose kernel with small matrices. The issue is that we use a per-CTA tiling scheme with hard-coded 8B vector loads and stores, and if the number of CTAs drops below the number of SMs (132 for H100) we can expect reduced throughput. This PR adds a very rough heuristic: if the data is small enough that we would run with fewer than 128 CTAs, reduce the vector size to 4B to increase the number of CTAs by a factor of 4.
Running on a H100 SXM with FP16 input and FP8 output:
The timings are not quite apples-to-apples since the baseline runs use the 22.12 PyTorch container while the optimized runs use 23.01. The runs for 2048 x 2048 and larger should all use the same kernel.
We could optimize even further for even smaller matrices, but it would require a more intricate performance model that considers the reduced cache efficiency. In this scale, I'd also expect the kernel launch cost to be the dominant factor in runtime.