[CUDA] Implement Scan kernel #2347
Merged
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
|
Looks awesome! I'm wondering if those logaddexp tests will fail on Metal / CPU as well? |
awni
approved these changes
Jul 10, 2025
|
The exp test fails on Metal, should I port the |
Yes that would be super useful.. it looks like a bug as is. |
|
Updated to fix the exp on Metal. |
1 task
Jckwind
pushed a commit
to TheProxyCompany/mlx
that referenced
this pull request
Aug 28, 2025
* Contiguous scan * Strided scan * Enable tests * Fix failing logaddexp test * Use cexpf in Metal
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 PR ports the Scan kernels from Metal to CUDA.
There are also a few changes to fix the
LogAddExpop for complex numbers which caused the scan tests to fail:cexpf.cuhfile is added to port the complex numberexpfrom thrust, we can not usethrust::expdirectly unfortunately because it does not compile in JIT kernels when building with old CUDA versions (<= 12.2).log1pis moved fromutils.cuhtounary_ops.cuhso it is put together with other complex number functions.LogAddExp's implementation is updated to use PyTorch's algorithm.