Fix performance regression on transposes on Llama 3 8B FP8 #434
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
i-chaochen
approved these changes
Nov 17, 2025
Collaborator
i-chaochen
left a comment
i-chaochen
left a comment
There was a problem hiding this comment.
Nice! I like the renaming it as SHMEM_GROUP
Also calculate their count as kNumThreadsPerBlock / kNumShmemBanks to avoid inconsistency when manually specified. This change is NFC for non-AMD GPUs. For AMD GPUs, it fixes the performance regression caused by inconsistency between shmem_group size, kNumThreadsPerBlock and kNumShmemBanks. It ended up in a situation downstream where half of the launched threads per block were not utilized at all. Updated packed transpose tests to verify correct thread utilization.
nurmukhametov
force-pushed
the
anurmukh/fix-packed-transpose-threads-0.7.1
branch
from
1bb6c4d to
7db6f55
Compare
zoranjovanovic-ns
approved these changes
Nov 17, 2025
Collaborator
|
Hi @nurmukhametov all CI are green, I guess we can merge it now? |
Author
Yes, we can if nobody else wants to review it. |
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.
[XLA:GPU] Rename
warptoshmem_groupinPackedTranspose.Also calculate their count as
kNumThreadsPerBlock / kNumShmemBanksto avoid inconsistency when manually specified.This change is NFC for non-AMD GPUs. For AMD GPUs, it fixes the performance regression caused by inconsistency between
shmem_groupsize,kNumThreadsPerBlockandkNumShmemBanks. It ended up in a situation downstream where half of the launched threads per block were not utilized at all. Updated packed transpose tests to verify correct thread utilization.