Skip to content
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

Replace exp2 with exp #25984

Merged
merged 2 commits into from
Jan 20, 2025
Merged

Replace exp2 with exp #25984

merged 2 commits into from
Jan 20, 2025

Conversation

Rifur13
Copy link
Collaborator

@Rifur13 Rifur13 commented Jan 20, 2025

Use the exp2 instruction instead of the exp operation for a 10-20% speedup to the attention kernel.
This change works because of the identity: e^x = 2^(x * log2(e)).

The softmax scale and the log2(e) rescaling are also folded in one multiplication.

@Rifur13 Rifur13 requested a review from sharadmv January 20, 2025 04:48
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jan 20, 2025
@copybara-service copybara-service bot merged commit 7f19b34 into jax-ml:main Jan 20, 2025
23 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants