[AMD] Always swap operands of mfma and use mfma.transposed layout #4767
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
zhanglx13
force-pushed
the
fix_order_mfma_transposed
branch
from
2065595
to
c525ee0
Compare
antiagainst
requested changes
Sep 20, 2024
zhanglx13
force-pushed
the
fix_order_mfma_transposed
branch
from
c525ee0
to
b3a9fcd
Compare
For original mfmaLayout, all elements for each thread are along the M dim. Threfore, when storing results to global memory, each thread cannot do vectorized global store since the result tensor is always N-minor. This PR swaps the operands of mfma instructions, and its effect is that in the result tensor, all elements for each thread are along the N dim. Now threads can do vectorized global store. And this can reduce the time of the epilogue. We already enabled swapping mfma operands for flash attention kernels so that the results of the first dot can be kept in register and used as the 1st operand of the second dot. For more details about swapping operands and how it works, please check the presentation about AMD backend at last year's triton conference: Bringing Triton to AMD GPUs: https://www.youtube.com/watch?v=8o7Jhbv8xek&t=1s
In general, we should use getThreadOrder in most places where getOrder is called. Note that order and threadOrder can be different, and this is the case for mfma.transposed layout.
zhanglx13
force-pushed
the
fix_order_mfma_transposed
branch
from
b3a9fcd
to
7ce8775
Compare
antiagainst
approved these changes
Sep 28, 2024
|
The macos failure seems to be relating to infra issues. Also I verified locally compilation on macos is fine. So merging. |
|
macOS build fix at #4827 |
Luosuu
pushed a commit
to Luosuu/triton
that referenced
this pull request
Nov 13, 2024
…iton-lang#4767) This helps to improve writeout to use `global_store_dwordx2`. Along the way this PR - Fixed the issue with getOrder for mfma layout - Fixed the issue with reduceOp when dealing with mfma.transposed layout In general, getOrder and getThreadOrder can return different values, and this is the case for mfma.transposed layout. Therefore, we shouldn't assume order and threadOrder are always the same.
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.
This PR
In general, getOrder and getThreadOrder can return different values, and this is the case for mfma.transposed layout. Therefore, we shouldn't assume order and threadOrder are always the same.