add static LHS scaling support for GMM v2 FWD by CaptainO5 · Pull Request #4735 · AI-Hypercomputer/maxtext

CaptainO5 · 2026-08-04T23:51:56Z

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_scale helper to extract the static activation scale directly from the provided Qwix quantization_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 extracted lhs_scale down into the underlying GMM v2 kernel to seamlessly orchestrate FP8 LHS quantization. Finally, this change also updates the underlying gmm_v2 kernel fork to accept static LHS scale.

BUGS: 538171764, 540900176
FIXES: 538171764

Tests

I added unit tests in tests/unit/quantizations_test.py.

E2E Test

Loss is similar and perf is slightly better between after (fwd lhs quant internal static) vs. before (fwd lhs quant internal dynamic), more details: b/538171764#comment2

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

codecov · 2026-08-05T00:12:55Z

Codecov Report

❌ Patch coverage is 92.85714% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/quantizations.py	72.22%	4 Missing and 1 partial ⚠️
...ernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py	98.18%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

shuningjin

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.py and quantization_test.py.

shuningjin

Great work! LGTM.

github-actions · 2026-08-08T00:09:49Z

🤖 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.

github-actions

## 📋 Review Summary

This pull request introduces static LHS (activation) scaling support for the GMM v2 forward pass under the Megablox and Pallas TPU v2 backends in MaxText. The changes are very high quality, exceptionally well-structured, and include robust unit tests verifying both scale extraction and boundary/dynamic cases.

🔍 General Feedback

Well-Designed Abstractions: Wrapping lhs and optional scale into the LhsRef registered dataclass is highly elegant and aligns perfectly with existing codebase conventions like WeightsRef.
Robust Scale Validation: Checking for underflow and scale values close to 0 protects the kernel from division-by-zero or numeric instability.
Thorough Test Coverage: Unit tests cover all key helper methods, including validation of symmetric static scaling and dynamic calibration fallback.

RissyRan

LGTM, and have a few comments. I think we could follow up static scaling inside of kernel to perf reason. Happy to chat more!

RissyRan · 2026-08-08T01:19:58Z

+def _fwd_prepare_lhs_scale(quantization_rule: qwix.QtRule | None) -> jax.Array | None:
+  """Extracts the static LHS (activation) scale for the GMM v2 forward pass.
+
+  GMM v2 only supports lhs_scale from symmetric fixed range calibration


Nit: Why we have "only"? It see you also mentioned if we set as None, it will support dynamic scaling.

GMM v2 does not support asymmetric fixed range calibration, i.e., when min != -max in fixed,min,max. I will rewrite the comment to make it clearer.

RissyRan · 2026-08-08T01:22:27Z

+  Unlike `rhs`, the lhs is passed to the kernel *unquantized*. When
+  `scale` is provided, the kernel uses it to quantize the lhs (i.e.
+  `qvalue = clip(lhs / scale)` and the result is multiplied back by `scale`).
+  The scale's shape encodes the granularity (per-tensor `[1, 1]`; extensible to


Does this mean we could support sub-channel when passing [M, num_blocks]?

In validate_inputs() we explicitly check for a [1, 1] shape now, but this design is extensible to sub-channel quantization, correct?

True, currently only per-tensor is supported

RissyRan · 2026-08-08T01:27:12Z

  has_scale: bool = False

+  @property
+  def should_use_external_scale(self) -> bool:


Wondering for perf reason, should we keep quantization inside of gmm v2 kernel? So that user doesn't need to handle static scaling before the call.

suhithr

LGTM, thank you!

suhithr · 2026-08-10T19:04:51Z

+  # Whether a scale array accompanies this input. The *direction* is inferred
+  # from the dtype relationship: when the input already arrives quantized
+  # (dtype == quant_dtype) the scale dequantizes it (rhs); when it arrives
+  # unquantized (dtype != quant_dtype) the scale quantizes it online (lhs).


Thanks for the clarifying comment

suhithr · 2026-08-10T19:15:45Z

-          block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale)
-          # Convert lhs into quantized dtype.
-          block_lhs_q = (block_lhs * block_scale_inv).astype(lhs_q_dtype)
+          if should_use_external_scale:


The above comment indicates that the previous version of the code depends upon a compiler optimization to avoid repeated re-computation of the block_scale.
Given your performance test results it seems we are getting the same CSE and avoiding recomputation with this branch?

I believe so, yes

suhithr · 2026-08-10T19:23:29Z

+  Unlike `rhs`, the lhs is passed to the kernel *unquantized*. When
+  `scale` is provided, the kernel uses it to quantize the lhs (i.e.
+  `qvalue = clip(lhs / scale)` and the result is multiplied back by `scale`).
+  The scale's shape encodes the granularity (per-tensor `[1, 1]`; extensible to


In validate_inputs() we explicitly check for a [1, 1] shape now, but this design is extensible to sub-channel quantization, correct?

Reconciles the ahead-of-time activation quantization (quantize before EP all-gather / ragged-sort so both move fp8 instead of bf16) with main's independent history since this branch's stale base: - ops.py / pallas_mosaic_tpu_v2_gmm_kernel.py: supersedes PR #4735's static-LHS-scaling mechanism (_fwd_prepare_lhs_scale/LhsRef), which assumes lhs always arrives unquantized and has no QArray handling -- incompatible with activations arriving pre-quantized. This branch's ahead-of-time quantization already covers #4735's target case (qwix's own qpl.quantize() natively supports "fixed" calibration), so nothing is lost. - moe.py: re-applied permute()/unpermute()/roe_ag_and_route/extract_vma QArray handling and the use_single_sparsecore removal on top of main's current moe.py (which independently gained DSV4 aux-loss-free routing and unrelated forward-pass bug fixes since this branch's base). - ragged_gather_reduce_v2.py: also removes disabled dead code (`if False and ...`) and a silent broad-except around qwix.dequantize() in the SparseCore fallback path. - Deliberately did not re-apply: stale vllm_batched_rpa/get_logical_axis_rules reversions (main already has this content; branch was just behind), or the sparse_matmul dispatch restructuring + two quantization-guard weakenings (kept main's existing self.config.quantization and ... guards intact). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Follow-up to the previous rebase commit -- audited every hunk against "is this strictly required for quantize-before-EP-all-gather" and reverted what wasn't: - moe.py, ragged_sort.py: restore use_single_sparsecore (a SparseCore core-count knob for the ragged kernels, unrelated to activation dtype) that had been dropped as part of the original PR's own cleanup. main already carries it unmodified; only the QArray dual-dispatch is new. - ragged_gather_reduce_v2.py: revert to main verbatim. The SparseCore capacity/tiling fixes (num_cores==2 handling, assert->fallback, etc.) aren't quantization-specific -- they'd matter for bf16 runs under the same hardware conditions. The QArray-dequant fallback in _fallback_implementation is also dead code: checked every call site of ragged_gather_reduce and none ever pass it a QArray (only gradients and GMM outputs, never quantized activations flow through the reduce path). - tests/unit/pallas_mosaic_tpu_v2_kernel_test.py: revert to main verbatim. test_gmm_prequantized_activation exercises gmm_v2's lhs_scale parameter, which belongs to PR #4735's superseded static-scale mechanism and doesn't exist in this branch's kernel signature. - ops.py / pallas_mosaic_tpu_v2_gmm_kernel.py: confirmed already minimal, no #4735 mechanism present, only the lhs_q_dtype detection / inner_kernel bypass / out_dtype bf16 floor needed for pre-quantized lhs to work. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Per feedback: don't delete a colleague's in-kernel quantization logic just because it's currently unreachable for our specific call path -- restore it as an additional case rather than replacing it, and revisit any actual overlap in a separate PR/session. _fwd_run_tokamax_v2 (and make_gmm_configs / inner_kernel in the kernel) now handle three cases: 1. lhs is a pre-quantized QArray (ours): unwrap to .qvalue, no in-kernel scale, real scale reapplied externally after the kernel call. 2. lhs is raw with a static "fixed"-calibration scale available (#4735): _fwd_prepare_lhs_scale supplies lhs_scale, kernel quantizes internally using it instead of a dynamic per-block absmax. 3. lhs is raw with neither: existing dynamic in-kernel quantization, unchanged. Verified: AOT compile check (train_compile_671b_v6e64.sh, deepseek3-671b fp8, v6e-64) still compiles cleanly with this layered version, same memory footprint as before. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

CaptainO5 requested review from Shuwen-Fang and suhithr August 5, 2026 01:15

suhithr reviewed Aug 5, 2026

View reviewed changes

Comment thread src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py Resolved

shuningjin reviewed Aug 6, 2026

View reviewed changes

Comment thread src/maxtext/kernels/megablox/ops.py Resolved

Comment thread tests/unit/moe_test.py Resolved

Comment thread src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py Resolved

CaptainO5 requested review from jshin1394 and liudangyi as code owners August 6, 2026 20:42

shuningjin reviewed Aug 6, 2026

View reviewed changes

Comment thread src/maxtext/layers/quantizations.py Resolved

Comment thread src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py Resolved

shuningjin reviewed Aug 7, 2026

View reviewed changes

Comment thread tests/unit/quantizations_test.py Resolved

Comment thread src/maxtext/kernels/megablox/ops.py Resolved

Comment thread src/maxtext/layers/quantizations.py Resolved

Comment thread src/maxtext/kernels/megablox/ops.py Resolved

Comment thread src/maxtext/kernels/megablox/ops.py Resolved

shuningjin approved these changes Aug 7, 2026

View reviewed changes

RissyRan added the gemini-review label Aug 8, 2026

github-actions Bot reviewed Aug 8, 2026

View reviewed changes

Comment thread src/maxtext/kernels/megablox/ops.py Resolved

CaptainO5 force-pushed the lhs_scale_import branch from e152632 to 305e9cd Compare August 8, 2026 00:20

RissyRan approved these changes Aug 8, 2026

View reviewed changes

add static LHS scaling support for GMM v2

5fd507a

CaptainO5 force-pushed the lhs_scale_import branch from 379db16 to 5fd507a Compare August 10, 2026 16:13

CaptainO5 added the pull ready label Aug 10, 2026

suhithr approved these changes Aug 10, 2026

View reviewed changes

copybara-service Bot merged commit ef71a22 into main Aug 10, 2026
127 of 128 checks passed

copybara-service Bot deleted the lhs_scale_import branch August 10, 2026 21:20

brightcoder01 mentioned this pull request Aug 16, 2026

Update expected threads to fix maxtext profiler failures #799

Closed

Conversation

CaptainO5 commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

E2E Test

Checklist

Uh oh!

codecov Bot commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

shuningjin left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

shuningjin left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Aug 8, 2026

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

🔍 General Feedback

Uh oh!

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CaptainO5 Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

suhithr left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

CaptainO5 commented Aug 4, 2026 •

edited

Loading

codecov Bot commented Aug 5, 2026 •

edited

Loading

shuningjin left a comment •

edited

Loading

CaptainO5 Aug 10, 2026 •

edited

Loading