[BACKEND] Implement multiCTA support for TMA gather/scatter#9977
Merged
Conversation
Contributor
Author
|
@peterbell10 can you review this one? |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e63aee2ccb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
lezcano
force-pushed
the
scatter_gather_multicast
branch
from
e63aee2 to
93b0fb1
Compare
Collaborator
|
I might get a chance to try this out actually |
lezcano
force-pushed
the
scatter_gather_multicast
branch
from
9152ee1 to
52744e4
Compare
lezcano
changed the title
peterbell10
requested changes
Apr 13, 2026
lezcano
force-pushed
the
scatter_gather_multicast
branch
from
43bbba8 to
bfe7657
Compare
Contributor
Author
|
done, added these ops to the comprehensive TMA into mma test. |
Contributor
Author
|
ugh, the tests are failing. Let me look into those. |
lezcano
force-pushed
the
scatter_gather_multicast
branch
from
3967b38 to
d55f753
Compare
Contributor
Author
|
found a real latent multicta issue and fixed it. Thank you for pushing for more comprehensive tests. |
peterbell10
reviewed
Apr 15, 2026
| mbarrier.init(bar, count=1) | ||
|
|
||
| gather_offsets = ttgl.load(gather_idx_ptr + ttgl.arange(0, BLOCK_M, layout=x_offsets_layout)) | ||
| mbarrier.expect(bar, blackwell_tma.nbytes_per_cta_gather(in_desc, gather_offsets)) |
Contributor
There was a problem hiding this comment.
Maybe we should have smem.nbytes_per_cta instead?
Contributor
Author
There was a problem hiding this comment.
Yep, I was thinking about that the other day. Let me do that.
Contributor
Author
There was a problem hiding this comment.
Added, for SharedLinearEncoding we compute a pseudo cga_layout of sorts to divide the shape by it.
We also add tighter invariants for the cga_layout part of all TMA ops
lezcano
force-pushed
the
scatter_gather_multicast
branch
2 times, most recently
from
c1148c1 to
6a54bcc
Compare
lezcano
added a commit
that referenced
this pull request
Apr 15, 2026
tbh, I think this is also the way we should represent this in our IR. I will add an end-to-end test once #9977 is merged (I will extend the test in that PR to test this op)
lezcano
force-pushed
the
scatter_gather_multicast
branch
from
89e495a to
8523533
Compare
peterbell10
reviewed
Apr 16, 2026
peterbell10
approved these changes
Apr 16, 2026
raymondtay
pushed a commit
to raymondtay/triton
that referenced
this pull request
Apr 18, 2026
tbh, I think this is also the way we should represent this in our IR. I will add an end-to-end test once triton-lang#9977 is merged (I will extend the test in that PR to test this op)
raymondtay
pushed a commit
to raymondtay/triton
that referenced
this pull request
Apr 18, 2026
…ang#9977) We also add tighter invariants for gather/scatter ops as well
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
4 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.
We also add tighter invariants for gather/scatter ops as well