PR #18838: [NVIDIA GPU] Support multi-operand collective-permute #20302
+664
−167
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
05397ed to
d7dadbf
Compare
Imported from GitHub PR #18838 For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because 1. it gets rid of some kernel launch overhead, and allows NCCL to do some message fusion; 2. fewer collectives make it easier for LHS to make better decision. In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands. Copybara import of the project: -- 5e10aba by Terry Sun <[email protected]>: support multi-operand cp -- 170fead by Terry Sun <[email protected]>: minor refactoring -- 0d85070 by Terry Sun <[email protected]>: update python interface -- 9812a10 by Terry Sun <[email protected]>: polish python interface -- 3a1552c by Terry Sun <[email protected]>: formatting -- d3657f8 by Terry Sun <[email protected]>: formatting -- 9caacb4 by Terry Sun <[email protected]>: refactor overloading -- 0aff5e0 by Terry Sun <[email protected]>: minor refactor -- 20a0e3d by Terry Sun <[email protected]>: add parser test Merging this change closes #18838 FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp a32d5aa PiperOrigin-RevId: 704206076
d7dadbf to
3e1c094
Compare
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.
PR #18838: [NVIDIA GPU] Support multi-operand collective-permute
Imported from GitHub PR #18838
For collective-permutes with small message sizes, it is beneficial to combine them into a single collective because
In order to support combining collective-permutes, we need to support multi-operand collective-permute first, a.k.a. the combined collective-permute. This PR extends the existing CP interface by overloading it, so that a CP can have multiple operands.
Copybara import of the project:
--
5e10aba by Terry Sun [email protected]:
support multi-operand cp
--
170fead by Terry Sun [email protected]:
minor refactoring
--
0d85070 by Terry Sun [email protected]:
update python interface
--
9812a10 by Terry Sun [email protected]:
polish python interface
--
3a1552c by Terry Sun [email protected]:
formatting
--
d3657f8 by Terry Sun [email protected]:
formatting
--
9caacb4 by Terry Sun [email protected]:
refactor overloading
--
0aff5e0 by Terry Sun [email protected]:
minor refactor
--
20a0e3d by Terry Sun [email protected]:
add parser test
Merging this change closes #18838
FUTURE_COPYBARA_INTEGRATE_REVIEW=#18838 from terryysun:terryysun/grouped_cp a32d5aa