FP8 support for tokamax gmm kernel - #2706
Merged
copybara-service[bot] merged 1 commit intoNov 20, 2025
Merged
Conversation
Copy link
Copy Markdown
Collaborator
Sorry, something went wrong.
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
2 times, most recently
from
9d50c6a to
33cddf8
Compare
November 17, 2025 22:33
khatwanimohit
marked this pull request as ready for review
November 17, 2025 22:41
khatwanimohit
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
michelle-yooh,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
November 17, 2025 22:41
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
from
33cddf8 to
64268d9
Compare
November 17, 2025 22:47
suexu1025
approved these changes
Nov 17, 2025
Comment thread
src/MaxText/layers/moe.py
RissyRan
reviewed
Nov 17, 2025
Comment thread
src/MaxText/configs/base.yml
Outdated
Copy link
Copy Markdown
Contributor
|
🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
Sorry, something went wrong.
Copy link
Copy Markdown
Contributor
There was a problem hiding this comment.
📋 Review Summary
This pull request adds support for FP8 quantization using the tokamax GMM kernel. The changes are well-structured and introduce fine-grained configuration for tiling dimensions. The core logic is updated to switch between the existing megablox backend and the new tokamax backend.
🔍 General Feedback
- The introduction of specific tiling configurations for forward and backward passes is a good improvement for tunability.
- The changes in moe.py to support the tokamax backend are well-integrated.
- I've left a few minor suggestions for improving code clarity and fixing a potential typo. There is also a question about a change in implementation detail that might affect performance.
Sorry, something went wrong.
Comment thread
src/MaxText/configs/types.py
Outdated
Comment thread
src/MaxText/kernels/megablox/ops.py
RissyRan
reviewed
Nov 18, 2025
Comment thread
src/MaxText/configs/types.py
Outdated
Comment thread
src/MaxText/configs/types.py
Outdated
Comment thread
src/MaxText/configs/types.py
Outdated
Comment thread
src/MaxText/kernels/megablox/ops.py
Comment thread
src/MaxText/kernels/megablox/ops.py
Comment thread
src/MaxText/layers/moe.py
Comment thread
src/MaxText/layers/moe.py
Outdated
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
from
64268d9 to
1bc7561
Compare
November 18, 2025 00:30
RissyRan
approved these changes
Nov 18, 2025
RissyRan
left a comment
Copy link
Copy Markdown
Collaborator
There was a problem hiding this comment.
LGTM! Just one comment to move proper assertion in moe.py to the pyconfig file
Sorry, something went wrong.
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
from
1bc7561 to
e9d47c6
Compare
November 18, 2025 02:45
Comment thread
src/MaxText/layers/moe.py
Outdated
Comment thread
src/MaxText/kernels/megablox/backend.py
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
4 times, most recently
from
7cf0bc9 to
4f761e0
Compare
November 19, 2025 01:05
Comment thread
src/MaxText/configs/base.yml
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
from
4f761e0 to
ffda110
Compare
November 19, 2025 18:37
khatwanimohit
force-pushed
the
mohit/tokamax_quant_gmm
branch
from
ffda110 to
67dc1f3
Compare
November 19, 2025 18:37
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Reviewers
4 participants
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.
Description
This PR calls tokamax_api.gmm in megablox/ops.py for supporting FP8 quantization using tokamax's gmm kernel. It also extends the tiling configs to 18 values, 9 for wi-* and 9 for wo tensors.
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/458461066
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Command to repro: https://paste.googleplex.com/6669707558125568
Xprof: http://shortn/_HoU0Xj2IZ9
Checklist
Before submitting this PR, please make sure (put X in square brackets):