add static LHS scaling support for GMM v2 FWD - #4735
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
| # ============================================================================== | ||
| # Forked from: | ||
| # https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py | ||
| # https://github.com/openxla/tokamax/blob/a1105e7513c4cc8604bad5627d099dcf09430ca1/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py |
There was a problem hiding this comment.
To clarify the PR description, are there any divergent changes to this file in this PR?
There was a problem hiding this comment.
Good point, while there are no new divergent changes introduced in this pr, there is a divergence related to the partial_sum feature added earlier.
There was a problem hiding this comment.
Thanks for the carefully merging tokamax changes with previous maxtext changes! I have double checked. This aligns with my attempt, diff.
There was a problem hiding this comment.
Thanks for adding the support! Overall looks good.
- Also thanks for the carefully merging new tokamax changes with previous maxtext changes! I have double checked. This aligns with my attempt, diff.
- In the PR description, can add E2E test result, e.g., loss is similar and perf is slightly better between after (fwd lhs quant internal static) vs. before (fwd lhs quant internal dynamic), with a bug link for future reference.
- Maybe the function and test can be moved to
quantization.pyandquantization_test.py.
| # ============================================================================== | ||
| # Forked from: | ||
| # https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py | ||
| # https://github.com/openxla/tokamax/blob/a1105e7513c4cc8604bad5627d099dcf09430ca1/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py |
There was a problem hiding this comment.
Thanks for the carefully merging tokamax changes with previous maxtext changes! I have double checked. This aligns with my attempt, diff.
- Consolidate scale parsing from `ops.py` and `manual_quantize` into a new `quantizations.get_static_scale` utility. - Streamline `manual_quantize` string validation to delegate directly to the new utility. - Ensure `ops._fwd_prepare_lhs_scale` preserves its original `(1, 1) float32` and `None` fallback constraints. - Migrate scale extraction and exception tests from `moe_test.py` to `quantizations_test.py`.
| if args[0] + args[1] == 0: | ||
| qmax = float(numerics.get_symmetric_bound(qtype)) | ||
| scale_val = args[1] / qmax | ||
| else: | ||
| qmin, qmax = numerics.get_asymmetric_bound(qtype) | ||
| scale_val = (args[1] - args[0]) / float(qmax - qmin) |
There was a problem hiding this comment.
Seems gmm2 only supports in-kernel lhs scale with symmetric bound.
pallas_mosaic_tpu_v2_gmm_kernel.py performs symmetric scale-and-clip: block_lhs_q = jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype). here
The kernel does not implement zero-point shifts. If an asymmetric range is supplied, the result can be mathematically incorrect.
There was a problem hiding this comment.
Should use
flat_args_preceding = (group_sizes, group_offset, lhs_in, rhs_weights)
otherwise can be incorrect when partial_sum is not None and lhs_scale is not None
There was a problem hiding this comment.
Great catch, thanks for the headsup
3 participants
Description
This change introduces static LHS (activation) scaling support to the GMM v2 integration in
ops.py.Specifically, it adds a new
_fwd_prepare_lhs_scalehelper to extract the static activation scale directly from the provided Qwixquantization_rule(triggering when fixed calibration bounds are utilized). The GMM v2 forward pass (_fwd_run_tokamax_v2) is then updated to accept the quantization rule and pass the extractedlhs_scaledown into the underlying GMM v2 kernel to seamlessly orchestrate FP8 LHS quantization. Finally, this change also updates the underlyinggmm_v2kernel fork to accept static LHS scale.BUGS: 538171764, 540900176
FIXES: 538171764
Tests
I added a couple of unit tests in
tests/unit/moe_test.py.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.