Skip to content

Fix CUDA all-reduce planning for large inputs - #3603

Merged
zcbenz merged 1 commit into
ml-explore:mainfrom
sofinvalery:fix-all-reduce
May 31, 2026
Merged

zcbenz merged 1 commit into
ml-explore:mainfrom
sofinvalery:fix-all-reduce

Conversation

@sofinvalery

Copy link
Copy Markdown
Contributor

Fixes #3602.
Not sure if I should add a test.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for fixing this! I'm good without a test since it is going to require allocating a very large buffer.

Comment thread mlx/backend/cuda/reduce/all_reduce.cu Resolved
@zcbenz
zcbenz merged commit 790be8a into ml-explore:main May 31, 2026
16 checks passed
megacpp pushed a commit to DatasunriseOU/mlx that referenced this pull request Jun 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] All_reduce over >INT_MAX elements hits RuntimeError

2 participants