Refactor CCL APIs to align with torch.distributed conventions #326
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
This refactor reorders parameters and adds support for process groups: API Changes: - all_reduce: (out, in, op=SUM, group=None, async_op=False, config=None, workspace=None) - reduce_scatter: (out, in, op=SUM, group=None, async_op=False, config=None) - all_gather: (out, in, group=None, async_op=False, config=None) - all_to_all: (out, in, group=None, async_op=False, config=None) New Features: - Add ReduceOp enum (SUM, PRODUCT, MIN, MAX, etc.) matching torch.distributed - Add extract_group_info() helper to extract rank_start/rank_stride from ProcessGroup - Support strided process groups (e.g., TP groups [0,1,2,3] or DP groups [0,4,8,12]) - op parameter validates only SUM is used (other ops to be added later) Kernel Changes: - All CCL kernels now accept rank_start and rank_stride constexpr parameters - Kernel loops updated to iterate using group-aware rank calculation - Ring all-reduce computes next_rank on host side for group support Backward Compatibility: - Existing code using keyword arguments (config=...) continues to work - torch.distributed compatible parameter ordering (group before config)
github-actions
bot
added
in-progress
Contributor
There was a problem hiding this comment.
Pull request overview
This pull request refactors the CCL (Collective Communication Library) APIs to align with torch.distributed conventions by reordering parameters and adding support for process groups. However, the implementation contains several critical bugs that prevent process groups from working correctly.
Changes:
- Adds ReduceOp enum matching torch.distributed semantics
- Reorders API parameters to match torch.distributed: (out, in, op, group, async_op, config)
- Adds extract_group_info() helper to extract rank/stride information from ProcessGroup
- Updates all CCL kernels to accept rank_start and rank_stride parameters for group-aware rank calculation
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 14 comments.
Show a summary per file
| File | Description |
|---|---|
| iris/iris.py | Updated all CCL method signatures to add op and group parameters with reordered arguments |
| iris/experimental/iris_gluon.py | Updated CCL method signatures (missing group parameter in all_gather) |
| iris/ccl/init.py | Added ReduceOp to exports |
| iris/ccl/utils.py | Added ReduceOp enum and extract_group_info() helper function |
| iris/ccl/all_reduce.py | Updated kernels and function to support group parameters with rank_start/rank_stride |
| iris/ccl/reduce_scatter.py | Updated kernel and function to support group parameters |
| iris/ccl/all_gather.py | Updated kernel and function to support group parameters |
| iris/ccl/all_to_all.py | Updated Triton and Gluon kernels and function to support group parameters |
mawad-amd
reviewed
Jan 27, 2026
Collaborator
mawad-amd
left a comment
mawad-amd
left a comment
There was a problem hiding this comment.
Can we find better names for cur_rank and cur_rank_global? They are really confusing throughout the code.
Comment on lines
+29
to
+30
| rank_start: tl.constexpr, | ||
| rank_stride: tl.constexpr, |
Collaborator
There was a problem hiding this comment.
This is a valid concern. Is an iris instance over the entire world or just the GPUs belonging to the group?
Co-authored-by: Muhammad Awad <112003944+mawad-amd@users.noreply.github.com>
Member
Author
Let me think about that. |
mawad-amd
approved these changes
Jan 27, 2026
Collaborator
mawad-amd
left a comment
mawad-amd
left a comment
There was a problem hiding this comment.
Looks good to me. Thanks!
Add optional parameter to barrier methods to support process group-specific synchronization, aligning with the torch.distributed.barrier(group=None) API convention.
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 refactor reorders parameters and adds support for process groups:
API Changes:
New:
Kernel Changes: