From d322d47e468114064bb7b39bbd7cb7051d6cd9d0 Mon Sep 17 00:00:00 2001 From: maxtext authors Date: Wed, 2 Sep 2026 20:37:10 -0700 Subject: [PATCH] [megablox] Remove redundant expert-weight transpose copy in unquantized gmm_v2 DLHS The tokamax gmm_v2 kernel has no native transpose_rhs support, so the V2 DLHS gradient path (`_dlhs_run_tokamax_v2`) manually materializes `rhs.swapaxes(1, 2)` on every backward step. For the MoE expert weight this is a full-tensor HBM copy of a large `[num_experts, k, n]` array (e.g. `bf16[16, 7168, 2048]`), emitted by XLA as a standalone `copy` with `op_name="rhs"`. The V1 tokamax path already computes the identical DLHS with no transpose: it calls `tokamax.ragged_dot_general` with `DLHS_RAGGED_DOT_DIM_NUMS`, contracting over the `n` dimension directly. This change routes the unquantized gmm_v2 DLHS (quantization_rule is None) through that transpose-free V1 kernel. Quantized paths are unchanged and keep the V2 kernel, because their scale handling (`_bwd_prepare_inputs` / `_dlhs_scale_grad_by_rhs_scale`) is gmm_v2-specific. Verified on TPU with a DeepSeek-scale shape (lhs[m=16384, k=7168], weight[g=16, k=7168, n=2048], bf16): the full expert-weight transpose copy is removed (1 -> 0), gradients are bit-identical (max relative error 0.0), and the full gmm fwd+bwd is ~3% faster. The output is bit-identical, so this is a pure efficiency change with no numerical or convergence impact. PiperOrigin-RevId: 975477390 --- src/maxtext/kernels/megablox/ops.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index f88e0aa0ab..e761e754b5 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -501,6 +501,7 @@ def _gmm_bwd( interpret, lhs_vma_axes, use_gmm_v2_heuristic_tiling, + quantization_rule, ) # 4. DRHS Gradient Execution @@ -630,6 +631,7 @@ def _compute_dlhs( interpret: bool, lhs_vma_axes: tuple, use_gmm_v2_heuristic_tiling: bool, + quantization_rule: qwix.QtRule | None = None, ) -> jnp.ndarray: """Routes execution of DLHS based on backend choices.""" if use_tokamax_backend and not use_gmm_v2: @@ -642,6 +644,24 @@ def _compute_dlhs( use_manual_quantization, ) elif use_tokamax_backend and use_gmm_v2: + # The gmm_v2 kernel lacks native transpose_rhs support, so its DLHS + # (_dlhs_run_tokamax_v2) materializes a full transpose of the weight + # (rhs.swapaxes(1, 2)) every backward step -- a pure HBM copy of the (large) + # expert weight. For the unquantized path the V1 tokamax.ragged_dot_general + # DLHS computes the identical result with no transpose (via + # DLHS_RAGGED_DOT_DIM_NUMS) and is measurably faster on TPU (bit-identical + # output). Quantized paths keep the V2 kernel because their scale handling + # (see _bwd_prepare_inputs / _dlhs_scale_grad_by_rhs_scale) is gmm_v2 + # specific. + if quantization_rule is None: + return _dlhs_run_tokamax_v1( + dlhs_dout, + rhs, + group_sizes, + lhs_dtype, + transpose_rhs, + use_manual_quantization, + ) return _dlhs_run_tokamax_v2( dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, use_gmm_v2_heuristic_tiling, transpose_rhs )