Skip to content

Update the megablox kernel utility function to support float16 #1423

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Update the megablox kernel utility function to support float16 #1423

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7859a11
Update the megablox kernel to support float16
wang2yn84 Mar 19, 2025
9234754
Merge branch 'main' into wang2yn84-patch-2
wang2yn84 Mar 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions MaxText/kernels/megablox/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,5 +55,7 @@ def select_input_dtype(lhs: jnp.ndarray, rhs: jnp.ndarray) -> jnp.dtype:
# input precision, we need to convert bf16 argument to fp32 beforehand.
if supports_bfloat16_matmul() and lhs.dtype == jnp.bfloat16 and rhs.dtype == jnp.bfloat16:
return jnp.bfloat16
elif lhs.dtype == jnp.float16 and rhs.dtype == jnp.float16:
Copy link
Collaborator

Choose a reason for hiding this comment

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

I am don't think any current test will cover newly added float16 type?

return jnp.float16
else:
return jnp.float32
Loading