docs(mxfp): correct the outdated scaled_matmul TPU comment#331
Open
lokic233 wants to merge 1 commit into
Open
Conversation
jax.nn.scaled_matmul now has a TPU lowering (via the xla.scaled_dot composite, jax PR #39056), so the 'Once scaled_matmul supports TPU we will enable it' comment is outdated. But on TPU the composite expands to a dequantize->bf16 dot_general, which is slower than qwix's existing native fp8 tiled path. Update the comment to state that scaled_matmul is used only on GPU by design, and TPU/CPU intentionally fall through to the faster fp8 path. No behavior change.
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
1 participant
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.
Summary
jax.nn.scaled_matmulnow has a TPU lowering (via thexla.scaled_dotcomposite,jax PR #39056), so the existing comment — "Once scaled_matmul supports TPU, we will
enable it for TPU as well" — is outdated.
However, enabling it on TPU would be a regression: on TPU the composite expands to
a dequantize → bf16
dot_general, which is slower than qwix's existing native fp8tiled path (
_fast_dot_general). This PR updates the comment to state the accurate,intentional design —
scaled_matmulis used only on GPU, and TPU/CPU deliberatelyfall through to the faster native fp8 path.
No behavior change (comment only).
Authored with Navi on behalf of @lokic233.