Skip to content

CUDA backend: softmax #2272

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 11, 2025
Merged

CUDA backend: softmax #2272

merged 1 commit into from
Jun 11, 2025

Conversation

zcbenz
Copy link
Contributor

@zcbenz zcbenz commented Jun 11, 2025

This PR is split from #1983. This PR relies on changes in #2269 and must be merged later.

Ports the logsumexp and softmax kernels from Metal. The eval_gpu has a lot of duplicate code that should be revisited later.

Copy link
Member

@awni awni left a comment

Choose a reason for hiding this comment

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

I rebased on main and will merge once the tests clear!

@awni awni merged commit c371baf into ml-explore:main Jun 11, 2025
5 of 6 checks passed
@zcbenz zcbenz deleted the cuda-softmax branch June 12, 2025 02:32
nastya236 pushed a commit to nastya236/mlx that referenced this pull request Jun 16, 2025
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.

2 participants