Repository navigation
PR #41808: [XLA:GPU] Migrate CollectivePermute command buffer support to thunk - #118827
Merged
Merged
Conversation
… to thunk Imported from GitHub PR openxla/xla#41808 PR description: 📝 Summary of Changes Migrates CollectivePermute command-buffer support from `CollectivePermuteCmd` to `CollectivePermuteThunk` by appending the thunk directly as a command. Removes the old `CollectivePermuteCmd` implementation and command type. Adds command-buffer create/update coverage and real NCCL multi-GPU execution tests for `CollectivePermuteThunk`. 🎯 Justification This continues the migration of GPU collective command-buffer handling into thunks, reducing duplicated command wrapper logic and aligning CollectivePermute with the thunk-owned command-buffer path. Workloads using CollectivePermute under GPU command buffers benefit from exercising the same thunk implementation for eager execution and command-buffer recording. 🚀 Kind of Contribution 🧪 Tests ♻️ Cleanup 📊 Benchmark (for Performance Improvements) N/A. This PR is not a performance improvement. 🧪 Unit Tests: Updated `//xla/backends/gpu/runtime:collective_permute_thunk_test` to cover CollectivePermute thunk command-buffer create/update behavior and source-target remapping logic. 🧪 Execution Tests: Added `//xla/backends/gpu/runtime:collective_permute_thunk_multigpu_test`, covering real NCCL CollectivePermute execution on 2 GPUs through: - `ExecuteOnStream` - command-buffer create mode with warm-up - command-buffer update mode Verified with: `bazel test //xla/backends/gpu/runtime:collective_permute_thunk_test //xla/backends/gpu/runtime:collective_permute_thunk_multigpu_test` Copybara import of the project: -- 7e5aa8f1d0ba7764836dbf8aaf6a811e9a8baaf5 by Shawn Wang <[email protected]>: [xla:gpu] Migrate CollectivePermuteCmd to thunk command -- 2a81fcd0df75baaf6b855e4700a1afb4bbb9799d by Shawn Wang <[email protected]>: [xla:gpu] Add CollectivePermuteThunk multigpu tests -- 3091a945c2cdc7001f6fd082c03643d3c41d072e by Shawn Wang <[email protected]>: Fix clang-format for collective permute test -- 2d16ad745025e77b116716c6129b8bfc84dccfb4 by Shawn Wang <[email protected]>: Refactor CollectivePermuteThunk multi-GPU test -- e49a9476cf4eee52eafc0f444b95c2cf2f7ceac5 by Shawn Wang <[email protected]>: Fix clang-format -- 97e7f54d3d75a572d802e2e6cefd82f6264dc6c1 by Shawn Wang <[email protected]>: Fix clang-tidy issues in collective-permute tests -- 51a14bd549ab8282e78d11881615f2b83af6c989 by Shawn Wang <[email protected]>: Fix ASSERT_OK include-cleaner issue in multigpu test -- 546b99468761decd0d4c1408d9cc15690fa617ed by Shawn Wang <[email protected]>: Fix collective permute multigpu lambda return types Merging this change closes #41808 PiperOrigin-RevId: 917613648
copybara-service
Bot
force-pushed
the
exported_pr_917598170
branch
from
May 19, 2026 05:49
8570ff2 to
d91f231
Compare
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
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 #41808: [XLA:GPU] Migrate CollectivePermute command buffer support to thunk
Imported from GitHub PR openxla/xla#41808
PR description:
📝 Summary of Changes
Migrates CollectivePermute command-buffer support from
CollectivePermuteCmdtoCollectivePermuteThunkby appending the thunk directly as a command. Removes the oldCollectivePermuteCmdimplementation and command type. Adds command-buffer create/update coverage and real NCCL multi-GPU execution tests forCollectivePermuteThunk.🎯 Justification
This continues the migration of GPU collective command-buffer handling into thunks, reducing duplicated command wrapper logic and aligning CollectivePermute with the thunk-owned command-buffer path. Workloads using CollectivePermute under GPU command buffers benefit from exercising the same thunk implementation for eager execution and command-buffer recording.
🚀 Kind of Contribution
🧪 Tests
♻️ Cleanup
📊 Benchmark (for Performance Improvements)
N/A. This PR is not a performance improvement.
🧪 Unit Tests:
Updated
//xla/backends/gpu/runtime:collective_permute_thunk_testto cover CollectivePermute thunk command-buffer create/update behavior and source-target remapping logic.🧪 Execution Tests:
Added
//xla/backends/gpu/runtime:collective_permute_thunk_multigpu_test, covering real NCCL CollectivePermute execution on 2 GPUs through:ExecuteOnStreamVerified with:
bazel test //xla/backends/gpu/runtime:collective_permute_thunk_test //xla/backends/gpu/runtime:collective_permute_thunk_multigpu_testCopybara import of the project:
--
7e5aa8f1d0ba7764836dbf8aaf6a811e9a8baaf5 by Shawn Wang [email protected]:
[xla:gpu] Migrate CollectivePermuteCmd to thunk command
--
2a81fcd0df75baaf6b855e4700a1afb4bbb9799d by Shawn Wang [email protected]:
[xla:gpu] Add CollectivePermuteThunk multigpu tests
--
3091a945c2cdc7001f6fd082c03643d3c41d072e by Shawn Wang [email protected]:
Fix clang-format for collective permute test
--
2d16ad745025e77b116716c6129b8bfc84dccfb4 by Shawn Wang [email protected]:
Refactor CollectivePermuteThunk multi-GPU test
--
e49a9476cf4eee52eafc0f444b95c2cf2f7ceac5 by Shawn Wang [email protected]:
Fix clang-format
--
97e7f54d3d75a572d802e2e6cefd82f6264dc6c1 by Shawn Wang [email protected]:
Fix clang-tidy issues in collective-permute tests
--
51a14bd549ab8282e78d11881615f2b83af6c989 by Shawn Wang [email protected]:
Fix ASSERT_OK include-cleaner issue in multigpu test
--
546b99468761decd0d4c1408d9cc15690fa617ed by Shawn Wang [email protected]:
Fix collective permute multigpu lambda return types
Merging this change closes #41808
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#41808 from shawnwang18:shawnw/migrate_collective_permute_cmd_to_thunk 546b99468761decd0d4c1408d9cc15690fa617ed