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.

Sign up for GitHub

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

Jump to bottom

update Gemma attention for TPU #2130

Merged
merged 19 commits into from
Mar 8, 2025
Merged

update Gemma attention for TPU #2130

merged 19 commits into from
Mar 8, 2025

Conversation

divyashreepathihalli
Copy link
Collaborator

No description provided.

@github-actions github-actions bot added the Gemma Gemma model specific issues label Mar 7, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 7, 2025
mattdangerw
mattdangerw reviewed Mar 7, 2025
View reviewed changes
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
mattdangerw
mattdangerw approved these changes Mar 8, 2025
View reviewed changes
Copy link
Member

@mattdangerw mattdangerw left a comment

Choose a reason for hiding this comment

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

LGTM! Once we have tested this to our satisfaction, sounds like this belongs in a point release.

@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
@divyashreepathihalli divyashreepathihalli added the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
@kokoro-team kokoro-team removed the kokoro:force-run Runs Tests on GPU label Mar 8, 2025
@divyashreepathihalli divyashreepathihalli merged commit 7a7a6bd into keras-team:master Mar 8, 2025
10 checks passed
rtg0795 pushed a commit that referenced this pull request Mar 17, 2025
@divyashreepathihalli @rtg0795
update Gemma attention for TPU (#2130)
ae3d126 
* update Gemma attention for TPU

* add default fallback for GPU and CPU

* add fallback option if not running with JAX and TPU

* address review comments

* check input signature

* remove checking q length

* code reformat

* handle case when soft cap support is not needed

* fix format

* add tests for FA calls

* fix test

* update tests

* fix code format

* address review comments

* Update requirements-jax-cuda.txt

* Update gemma_causal_lm_test.py

* Update requirements-jax-cuda.txt
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mattdangerw mattdangerw mattdangerw approved these changes

Assignees
No one assigned
Labels
Gemma Gemma model specific issues
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@divyashreepathihalli @mattdangerw @kokoro-team