From 1fb67bff2c48513b7be39815b91df561e41f526a Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 00:03:33 +0000 Subject: [PATCH 01/11] Rebase quantize_sort_2 onto latest main 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 --- src/maxtext/kernels/megablox/ops.py | 89 ++++----- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 182 ++++-------------- .../kernels/ragged/ragged_gather_reduce_v2.py | 41 ++-- src/maxtext/kernels/ragged/ragged_sort.py | 149 +++++++------- src/maxtext/layers/moe.py | 141 ++++++++++---- .../unit/pallas_mosaic_tpu_v2_kernel_test.py | 45 +++++ 6 files changed, 315 insertions(+), 332 deletions(-) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index bf533e353c..c295e58e1c 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -27,6 +27,7 @@ from maxtext.layers import quantizations import qwix import qwix.pallas as qpl +from qwix._src.core.qarray import call_with_generic_broadcast import tokamax @@ -102,13 +103,18 @@ def gmm( act_calibration_method="absmax", ) + lhs_scale = None + if isinstance(lhs, qpl.QArray): + lhs_scale = lhs.scale + lhs = lhs.qvalue + gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0] # pylint: disable=C3001 gmm_fwd_bwd = jax.custom_vjp( gmm_fwd_bwd, nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15), ) gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype)) - return gmm_fwd_bwd( + out = gmm_fwd_bwd( lhs, rhs, group_sizes, @@ -127,6 +133,9 @@ def gmm( use_gmm_v2, partial_sum, ) + if lhs_scale is not None: + out = call_with_generic_broadcast(jnp.multiply, out, lhs_scale.astype(out.dtype)) + return out # ============================================================================== @@ -202,15 +211,7 @@ def _gmm_fwd( out = _fwd_run_tokamax_v1(lhs, rhs, group_sizes, preferred_element_type, transpose_rhs, use_manual_quantization) elif use_tokamax_backend and use_gmm_v2: out = _fwd_run_tokamax_v2( - lhs, - rhs, - group_sizes, - preferred_element_type, - tiling, - group_offset, - partial_sum, - transpose_rhs, - quantization_rule, + lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs ) else: out = _fwd_run_megablox( @@ -238,7 +239,11 @@ def _fwd_quantize_activation_and_weight( transpose_rhs: bool, ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]: """Handles act and weight quantization for GMM forward inputs.""" - if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and not use_gmm_v2: + if ( + quantization_rule.act_qtype + and not isinstance(lhs, qpl.QArray) + and not (jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn)) + ): lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, @@ -330,36 +335,6 @@ def _fwd_prepare_rhs_scale(rhs: qpl.QArray, transpose_rhs: bool = False) -> jnp. return jnp.broadcast_to(rhs_scale, (G, num_quant_blocks, 1, N)) -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. - - If a static scale is used, GMM v2 requires it to be from a symmetric fixed-range - calibration (e.g., 'fixed,-max,max' or 'fixed,max'). If no static scale is - provided, the kernel will compute a dynamic scale on the fly. - - Enforces a default (1, 1) shape for per-tensor quantization kernels. - - Args: - quantization_rule: The Qwix quantization rule from which to extract the scale. - - Returns: - The extracted static scale array, or None if not using purely fixed calibration. - """ - if quantization_rule is None: - return None - - method = quantization_rule.act_calibration_method - qtype = quantization_rule.act_qtype - - # Use dynamic quantization, gmm_v2 calculates dynamic scale internally - if method is None or qtype is None or not method.lower().startswith("fixed"): - return None - - scale_val = quantizations.get_static_scale(qtype, method) - - return jnp.full((1, 1), scale_val, jnp.float32) - - def _fwd_run_tokamax_v2( lhs: jnp.ndarray | qpl.QArray, rhs: jnp.ndarray | qpl.QArray, @@ -369,7 +344,6 @@ def _fwd_run_tokamax_v2( group_offset: jnp.ndarray | None, partial_sum: jnp.ndarray | None, transpose_rhs: bool, - quantization_rule: qwix.QtRule | None = None, ) -> jnp.ndarray: """Executes the Tokamax GMM V2 backend for forward pass OUT = LHS @ RHS.""" # if transpose_rhs=False, rhs is [g, k, n], remain unchanged @@ -382,24 +356,40 @@ def _fwd_run_tokamax_v2( rhs_operand = rhs_operand.qvalue rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs) + lhs_operand = lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs + maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and not ( + jnp.issubdtype(lhs_operand.dtype, jnp.integer) or jnp.issubdtype(lhs_operand.dtype, jnp.float8_e4m3fn) + ) + custom_fwd_tiling = gmm_v2.TileSizes( tile_m=tiling[0], tile_k=tiling[1], tile_n=tiling[2], ) - return gmm_v2.gmm_v2( - lhs=lhs, # pyrefly: ignore[bad-argument-type] + eff_pref_dtype = ( + jnp.bfloat16 + if (jnp.issubdtype(lhs_operand.dtype, jnp.integer) or jnp.issubdtype(lhs_operand.dtype, jnp.float8_e4m3fn)) + else preferred_element_type + ) + + out = gmm_v2.gmm_v2( + lhs=lhs_operand, # pyrefly: ignore[bad-argument-type] rhs=rhs_operand, # pyrefly: ignore[bad-argument-type] group_sizes=group_sizes, rhs_scale=rhs_scale, tile_info=custom_fwd_tiling, - preferred_element_type=preferred_element_type, + preferred_element_type=eff_pref_dtype, partial_sum=partial_sum, group_offset=group_offset, - lhs_scale=_fwd_prepare_lhs_scale(quantization_rule), + maybe_quantize_lhs=maybe_quantize_lhs, ) + if isinstance(lhs, qpl.QArray): + out *= lhs.scale.astype(out.dtype) + + return out + def _fwd_run_megablox( lhs: jnp.ndarray, @@ -560,7 +550,12 @@ def _bwd_prepare_inputs( # GMM2 FWD performs lhs quantization inside kernel, lhs is stored as unquantized dtype # in the residual tuple. In BWD, we explicitly quantize lhs. - if quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray): + if ( + quantization_rule + and quantization_rule.act_qtype + and not isinstance(lhs, qpl.QArray) + and not (jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn)) + ): lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 1826ef576b..24567d6272 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== # Forked from: -# https://github.com/openxla/tokamax/blob/a1105e7513c4cc8604bad5627d099dcf09430ca1/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py +# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py """GMM kernel implemented using Pallas.""" from abc import ABC, abstractmethod @@ -141,29 +141,6 @@ def get_bias(self) -> jax.Array: return jnp.concatenate([b_gate, b_up], axis=-1) -@jax.tree_util.register_dataclass -@dataclasses.dataclass(frozen=True) -class LhsRef: - """Dataclass for the lhs value and its optional quantization scale. - - 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 - per-channel `[M, 1]` and sub-channel `[M, num_blocks]`). - """ - - value: Any - scale: Any | None - - def get_value(self) -> jax.Array: - return self.value[...] - - def get_scale(self) -> jax.Array: - assert self.scale is not None - return self.scale[...] - - @jax.tree_util.register_dataclass @dataclasses.dataclass(frozen=True) class MetadataRef: @@ -196,21 +173,8 @@ class InputConfigs: quant_block_size: int | None dtype: jnp.dtype has_bias: bool = False - # 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). has_scale: bool = False - @property - def should_use_external_scale(self) -> bool: - # A scale is present but the input is not yet quantized - # (dtype != quant_dtype). The kernel uses it to quantize the input online - # and multiply the result by the scale after. This differs from an already - # quantized input (dtype == quant_dtype), whose scale only dequantizes after - # the matmul. - return self.has_scale and self.quant_dtype is not None and self.dtype != self.quant_dtype - @property def should_bitcast(self) -> bool: bits = jax.dtypes.itemsize_bits(self.dtype) @@ -275,14 +239,6 @@ def lhs_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): return (pl.ds(row_start, row_size), 0, k_id) - def lhs_scale_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): - # Per-tensor scale: a single [1, 1] value shared across every tile, so the - # block always reads index 0. Extension point: when the scale is per-channel - # or sub-channel, tile the row axis like `lhs_index_map` (using gm_id) and - # index the K-block axis from `k_id`. - del gm_id, k_id - return (0, 0) - def rhs_weight_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array): group_id = self.metadata_ref.gm_id_to_group_id[gm_id] return (group_id, k_id, n_id) @@ -327,23 +283,16 @@ def ps_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): def generate_block_specs( metadata_ref: MetadataRef, cfgs: GmmConfigs -) -> Tuple[Tuple[LhsRef, WeightsRef, pl.BlockSpec | None], pl.BlockSpec]: +) -> Tuple[Tuple[pl.BlockSpec, WeightsRef, pl.BlockSpec | None], pl.BlockSpec]: """Generates block specs for the given lhs, rhs, and out refs.""" index_map = IndexMaps(metadata_ref, cfgs) bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane) - lhs_value_spec = pl.BlockSpec( + lhs_block_spec = pl.BlockSpec( (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), index_map.lhs_index_map, ) - lhs_scale_spec = None - if cfgs.lhs_cfgs.has_scale: - lhs_scale_spec = pl.BlockSpec( - (1, 1), - index_map.lhs_scale_index_map, - ) - lhs_block_spec = LhsRef(value=lhs_value_spec, scale=lhs_scale_spec) tile_k_rhs = cfgs.tiles.tile_k if cfgs.rhs_cfgs.should_bitcast: @@ -392,7 +341,7 @@ def generate_block_specs( def inner_kernel( # In - tiled_lhs_ref: LhsRef, + tiled_lhs_ref: jax.Array, # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k] tiled_rhs_ref: RhsRef, # [tile_k, tile_n] # Partial Sum @@ -433,7 +382,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): mxu_size = tpu_info.mxu_column_size # Step 1: Input pre-processing. - tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] + tiled_lhs = tiled_lhs_ref.reshape(-1, cfgs.tiles.tile_k)[...] tiled_rhs = tiled_rhs_ref.get_weight() # When rhs is packed (quantized dtype packed into uint32), unpack it # back to the original dtype using pltpu.bitcast which operates on K @@ -497,14 +446,6 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): dtype_max = float(jnp.iinfo(lhs_q_dtype).max) preferred_element_type = jnp.int32 - # When the caller supplies a quantization scale, use it directly instead - # of computing a dynamic per-block absmax. - lhs_scale = lhs_scale_inv = None - should_use_external_scale = cfgs.lhs_cfgs.should_use_external_scale - if should_use_external_scale: - lhs_scale = tiled_lhs_ref.get_scale().astype(acc_ref.dtype) - lhs_scale_inv = 1.0 / lhs_scale - # Without n outer loop, result of quantized matmul becomes available only # at the last iteration of the loop. This means [tile_m, tile_n] value # needs to be stored until the last iteration. By adding n outer loop, @@ -525,11 +466,9 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): # Perform lhs quantization. Note that for every block_lhs, # same computation will be performed tiles_n//mxu_size times. # But we can let compiler perform CSE and avoid recomputation. - if should_use_external_scale: - assert lhs_scale is not None - assert lhs_scale_inv is not None - block_lhs_q = jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype) - block_scale = lhs_scale # [1, 1] + if jnp.issubdtype(tiled_lhs.dtype, jnp.integer) or jnp.issubdtype(tiled_lhs.dtype, jnp.float8_e4m3fn): + block_lhs_q = block_lhs + block_scale = jnp.array(1.0, dtype=acc_ref.dtype) else: block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True) block_scale = block_abs_max / dtype_max @@ -831,7 +770,7 @@ def kernel_main( lhs_group_sizes_ref: jax.Array, # int32[size_lhs_group] group_offset_ref: jax.Array, # int32[1] # In - lhs_ref: LhsRef, # value: [size_m, size_k] + lhs_ref: jax.Array, # [size_m, size_k] rhs_ref: WeightsRef, # [size_group, size_k, size_n] partial_sum_ref: jax.Array, # [size_m, size_n] # Out @@ -920,10 +859,8 @@ def kernel_main( ) # Bounded slice requires second last dim to be aligned to the sublane size. - # rhs_ref uses static tiling thus reshape is not needed. The lhs quant scale - # (when present) is small and statically tiled, so it is passed through as-is. - lhs_value_in = lhs_ref.value.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.value.shape[-1]) - lhs_in = LhsRef(value=lhs_value_in, scale=lhs_ref.scale) + # rhs_ref uses static tiling thus reshape is not needed. + lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.shape[-1]) ps_in = None if cfgs.has_partial_sum: ps_in = partial_sum_ref.reshape(-1, cfgs.dims.size_lhs_sublane, partial_sum_ref.shape[-1]) @@ -1056,8 +993,6 @@ def validate_inputs( group_sizes: jax.Array, group_offset: jax.Array, fuse_act: str | None = None, - maybe_quantize_lhs: bool = True, - lhs_scale: jax.Array | None = None, ) -> Dimensions: """Validates the inputs for the GMM kernel.""" @@ -1076,18 +1011,9 @@ def validate_inputs( assert partial_sum.shape[0] <= size_m if rhs_scale is not None: num_quant_blocks = rhs_scale.shape[1] - assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n), ( - f"rhs_scale shape {rhs_scale.shape}. Expecting ({size_group}," f" {num_quant_blocks}, 1, {size_n})" - ) + assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n) assert size_k % num_quant_blocks == 0 - if lhs_scale is not None: - assert maybe_quantize_lhs, "lhs_scale requires maybe_quantize_lhs=True." - # Only per-tensor scales are supported for now. The current implementation generalizes to per-channel [M, 1] and - # sub-channel [M, num_k_blocks]; extend the validation and the block spec / - # index map together when adding those. - assert lhs_scale.shape == (1, 1), "Only per-tensor lhs_scale of shape (1, 1) is supported, got " f"{lhs_scale.shape}." - assert group_offset.shape == (1,) size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype) @@ -1170,22 +1096,10 @@ def make_gmm_configs( maybe_quantize_lhs: bool, zero_initialize: bool, fuse_act: str | None = None, - lhs_scale: jax.Array | None = None, ): """Fills the GMM config for the GMM kernel.""" - dims = validate_inputs( - lhs, - rhs, - rhs_scale, - rhs_bias, - partial_sum, - group_sizes, - group_offset, - fuse_act, - maybe_quantize_lhs, - lhs_scale, - ) + dims = validate_inputs(lhs, rhs, rhs_scale, rhs_bias, partial_sum, group_sizes, group_offset, fuse_act) if rhs_scale is not None: has_scale = True @@ -1206,7 +1120,9 @@ def make_gmm_configs( ) lhs_q_dtype = None - if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul: + if jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn): + lhs_q_dtype = lhs.dtype + elif maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul: # Choose lhs quantization dtype based on TPU hardware support. is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) # pyrefly: ignore[bad-argument-type] tpu_info = pltpu.get_tpu_info() @@ -1222,14 +1138,6 @@ def make_gmm_configs( if not is_rhs_float: lhs_q_dtype = jnp.int8.dtype - if lhs_scale is not None: - assert lhs_q_dtype is not None, ( - "lhs_scale requires lhs quantization to engage, but no lhs quant " - "dtype was selected. Ensure rhs is quantized and the hardware supports " - "fp8/int8 matmul." - ) - has_lhs_scale = lhs_scale is not None and lhs_q_dtype is not None - lhs_cfgs = InputConfigs( quant_dtype=lhs_q_dtype, # Input quantization involves reading all elements in a block to compute @@ -1238,11 +1146,10 @@ def make_gmm_configs( # enough to minimize compute overhead of quantization. quant_block_size=512, dtype=lhs.dtype, - has_scale=has_lhs_scale, ) - if out_dtype is None: - out_dtype = lhs.dtype + if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn): + out_dtype = jnp.bfloat16.dtype if acc_dtype is None: if lhs_cfgs.quant_dtype is None: @@ -1302,7 +1209,6 @@ def gmm_v2( rhs_bias: jax.Array | None = None, # [size_group, 1, out_size] partial_sum: jax.Array | None = None, # [size_m, size_n] group_offset: jax.Array | None = None, # int32[1] - lhs_scale: jax.Array | None = None, # [1, 1] (per-tensor) *, tile_info: TileSizes | TileFn = calculate_tiling, # pyrefly: ignore[bad-function-definition] vmem_limit_bytes: int | None = None, @@ -1327,12 +1233,6 @@ def gmm_v2( rhs_bias: The rhs bias of shape [size_group, 1, out_size]. partial_sum: Optional. Per-token partial sums of shape [size_m, size_n]. group_offset: Optional. The group offset of shape [1,]. - lhs_scale: Optional scale used to quantize the (unquantized) lhs - inside the kernel and the result is multiplied back by `scale`. The shape - encodes granularity; currently only per-tensor `[1, 1]` is supported. When - None, a quantized lhs uses the default dynamic per-block absmax - calibration. Only takes effect when maybe_quantize_lhs is True and rhs is - quantized. tile_info: The tile sizes or tile function to use. vmem_limit_bytes: Optional vmem limit in bytes. precision: Unused. Exists for compatibility reasons. @@ -1372,20 +1272,11 @@ def gmm_v2( maybe_quantize_lhs=maybe_quantize_lhs, zero_initialize=zero_initialize, fuse_act=fuse_act, - lhs_scale=lhs_scale, ) dims = cfgs.dims tiles = cfgs.tiles # Prepare block specs. - lhs_scale_spec = None - if cfgs.lhs_cfgs.has_scale: - assert lhs_scale is not None - lhs_scale = lhs_scale.astype(jnp.float32) - lhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) - else: - lhs_scale = None - rhs_scale_spec = rhs_bias_spec = None if rhs_scale is not None: rhs_scale = rhs_scale.astype(jnp.float32) @@ -1437,15 +1328,33 @@ def gmm_v2( aligned_n = align_to(cfgs.out_size_n, num_lanes) out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n), cfgs.out_dtype) - lhs_in = LhsRef(value=lhs, scale=lhs_scale) rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), + WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ), + ] + partial_sum_spec = None if partial_sum is not None: + in_specs.append(pl.BlockSpec(memory_space=pltpu.HBM)) partial_sum_spec = pl.BlockSpec(memory_space=pltpu.HBM) + in_specs = [ + pl.BlockSpec(memory_space=pltpu.HBM), # lhs + WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ), # rhs_weights + partial_sum_spec, # partial_sum + ] input_output_aliases = {} if partial_sum is not None: - flat_args_preceding = (group_sizes, group_offset, lhs_in, rhs_weights) + flat_args_preceding = (group_sizes, group_offset, lhs, rhs_weights) leaves = jax.tree_util.tree_leaves(flat_args_preceding) partial_sum_idx = sum(1 for x in leaves if x is not None) input_output_aliases = {partial_sum_idx: 0} @@ -1455,18 +1364,7 @@ def gmm_v2( out_shape=out_init, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, - in_specs=[ - LhsRef( - value=pl.BlockSpec(memory_space=pltpu.HBM), - scale=lhs_scale_spec, - ), - WeightsRef( - weight=pl.BlockSpec(memory_space=pltpu.HBM), - scale=rhs_scale_spec, - bias=rhs_bias_spec, - ), - partial_sum_spec, - ], + in_specs=in_specs, out_specs=pl.BlockSpec(memory_space=pltpu.HBM), scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), @@ -1478,4 +1376,4 @@ def gmm_v2( cost_estimate=get_cost_estimate(cfgs), metadata=get_metadata(cfgs), input_output_aliases=input_output_aliases, - )(group_sizes, group_offset, lhs_in, rhs_weights, partial_sum)[:, : cfgs.out_size_n] + )(group_sizes, group_offset, lhs, rhs_weights, partial_sum)[:, : cfgs.out_size_n] diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index 0594676b74..c6a1a14f9e 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -18,6 +18,7 @@ import dataclasses import functools +import math from typing import Any import jax @@ -175,6 +176,10 @@ def _fallback_implementation( reduce_group_size: int, ) -> jax.Array: """Fallback implementation using JAX ops for non-SparseCore TPU or small inputs.""" + if hasattr(x, "scale") and hasattr(x, "qvalue"): + import qwix # pylint: disable=import-outside-toplevel + + x = qwix.dequantize(x) out = x[indices] * topk_weights[:, None].astype(jnp.float32) out = jnp.where(valid_rows_mask[:, None], out, 0) out = out.reshape(-1, reduce_group_size, out.shape[-1]) @@ -196,27 +201,17 @@ def _calculate_num_column_partitions( preferred_num_stages = 4 num_column_partitions = 1 while ( - num_cores % (num_column_partitions * 2) == 0 + (num_cores == 2 or num_cores % (num_column_partitions * 2) == 0) and hidden_size % (num_lanes * num_column_partitions * 2) == 0 and hidden_size // (num_column_partitions * 2 * num_lanes) >= preferred_num_stages ): next_candidate = num_column_partitions * 2 - next_row_partitions = num_cores // next_candidate - - # Calculate exactly how many pipeline invocations (outer loop) - _, row_chunk_size = _calculate_row_tiling(input_size, num_simd_lanes, next_row_partitions) - num_iterations = input_size // (row_chunk_size * next_row_partitions) # Ensure we satisfy the hardware constraint (num_row_partitions <= num_simd_lanes) first. if num_cores // num_column_partitions > num_simd_lanes: num_column_partitions = next_candidate continue - # Too many iterations cause high cumulative pipeline overhead. Set the - # limit based on empirical data. - if num_iterations > _CostModelConstants.MAX_ITERATIONS: - break - num_column_partitions = next_candidate return num_column_partitions @@ -228,8 +223,8 @@ def _calculate_row_tiling( num_row_partitions: int, ) -> tuple[int, int]: """Calculates the number of row subchunks and row chunk size.""" - base_block_size = num_simd_lanes * num_row_partitions - num_row_subchunks = max(1, min(4, pl.cdiv(input_size, base_block_size))) + base_block_size = max(1, num_simd_lanes * num_row_partitions) + num_row_subchunks = max(1, min(2, math.ceil(int(input_size) / base_block_size))) row_chunk_size = num_simd_lanes * num_row_subchunks return num_row_subchunks, row_chunk_size @@ -625,14 +620,7 @@ def col_loop(col_compute_offset): @functools.partial( - jax.jit, - static_argnames=( - "reduce_group_size", - "enforce_fallback", - "flops_override", - "bytes_accessed_override", - "use_single_sparsecore", - ), + jax.jit, static_argnames=("reduce_group_size", "enforce_fallback", "flops_override", "bytes_accessed_override") ) def ragged_gather_reduce( x: jax.Array, @@ -643,7 +631,6 @@ def ragged_gather_reduce( enforce_fallback: bool = False, flops_override: int = -1, bytes_accessed_override: int = -1, - use_single_sparsecore: bool = False, ) -> jax.Array: """Gathers ``x`` by ``indices``, weights and masks, then reduces by group. @@ -681,12 +668,12 @@ def ragged_gather_reduce( input_size = indices.size num_simd_lanes = sc_info.num_lanes num_lanes = pltpu.get_tpu_info().num_lanes - num_sc_cores = 1 if use_single_sparsecore else sc_info.num_cores - num_cores = num_sc_cores * sc_info.num_subcores + num_cores = sc_info.num_cores * sc_info.num_subcores num_column_partitions = _calculate_num_column_partitions(hidden_size, input_size, num_cores, num_lanes, num_simd_lanes) num_row_partitions = num_cores // num_column_partitions - assert num_row_partitions <= num_simd_lanes, f"{num_row_partitions=} must be <= {num_simd_lanes=}" + if num_row_partitions > num_simd_lanes: + return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size) num_row_subchunks, row_chunk_size = _calculate_row_tiling(input_size, num_simd_lanes, num_row_partitions) aligned_hidden_size = _align_to(hidden_size, 128 * num_column_partitions) @@ -720,7 +707,7 @@ def ragged_gather_reduce( # Step 4: Launch the SparseCore kernel. vector_mesh = plsc.VectorSubcoreMesh( - num_cores=num_sc_cores, + num_cores=sc_info.num_cores, num_subcores=sc_info.num_subcores, core_axis_name="core", subcore_axis_name="subcore", @@ -760,7 +747,7 @@ def ragged_gather_reduce( flops_override=flops_override, bytes_accessed_override=bytes_accessed_override, ), - scratch_types=( # pyrefly: ignore[bad-argument-type] + scratch_types=( _Scratch( num_rows_per_row_partition_vmem=pltpu.VMEM((num_simd_lanes,), jnp.int32), prev_iter_last_row_vmem=pltpu.VMEM((col_size // col_chunk_size, col_chunk_size), jnp.float32), diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index f3def646b9..62b8e3ea5a 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -18,6 +18,7 @@ import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce +import qwix.pallas as qpl def ring_ragged_sort( @@ -34,7 +35,6 @@ def ring_ragged_sort( gather_reduce_flops_override=-1, gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, - use_single_sparsecore=False, ): """Ragged-gather variant for AG-RS Expert Parallelism token routing. @@ -104,16 +104,36 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): if buffer_size is None or buffer_size >= num_tokens_local * topk: local_buffer_size = num_tokens_local * topk - x = ragged_gather( - hidden_states_local, - token_indices_sorted, - shard_output_start[None], - shard_output_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + if isinstance(hidden_states_local, qpl.QArray): + x_qval = ragged_gather( + hidden_states_local.qvalue, + token_indices_sorted, + shard_output_start[None], + shard_output_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) + x_scale = ragged_gather( + hidden_states_local.scale, + token_indices_sorted, + shard_output_start[None], + shard_output_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) + x = qpl.QArray(qvalue=x_qval, scale=x_scale) + else: + x = ragged_gather( + hidden_states_local, + token_indices_sorted, + shard_output_start[None], + shard_output_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) else: local_buffer_size = buffer_size # We only gather up to the available buffer size or the actual number of @@ -128,16 +148,36 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): local_buffer_size, axis=0, ) - x = ragged_gather( - hidden_states_local, - sliced_indices, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + if isinstance(hidden_states_local, qpl.QArray): + x_qval = ragged_gather( + hidden_states_local.qvalue, + sliced_indices, + jnp.int32(0)[None], + gather_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) + x_scale = ragged_gather( + hidden_states_local.scale, + sliced_indices, + jnp.int32(0)[None], + gather_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) + x = qpl.QArray(qvalue=x_qval, scale=x_scale) + else: + x = ragged_gather( + hidden_states_local, + sliced_indices, + jnp.int32(0)[None], + gather_end[None], + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + ) out = (x, group_sizes_local, topk_argsort_revert_indices) @@ -191,7 +231,6 @@ def _ring_ragged_sort_bwd(res, g_out): enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) else: # Buffering: g_x has size `local_buffer_size` (packed). @@ -217,7 +256,6 @@ def _ring_ragged_sort_bwd(res, g_out): enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) return grad_hidden_states, None @@ -240,7 +278,6 @@ def ring_ragged_unsort( gather_reduce_flops_override=-1, gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, - use_single_sparsecore=False, ): """Dual of :func:`ring_ragged_sort`. @@ -272,27 +309,14 @@ def ring_ragged_unsort( """ @jax.custom_vjp - def _ring_ragged_unsort( - sorted_tokens_local, - group_sizes_local, - topk_argsort_revert_indices, - topk_weights_flat, - ): + def _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat): """Unsort and scatter activations.""" return _ring_ragged_unsort_fwd( - sorted_tokens_local, - group_sizes_local, - topk_argsort_revert_indices, - topk_weights_flat, + sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat )[0] @jax.named_scope("ragged-unsort-fwd") - def _ring_ragged_unsort_fwd( - sorted_tokens_local, - group_sizes_local, - topk_argsort_revert_indices, - topk_weights_flat, - ): + def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat): """Executes unsorting sending tokens back.""" group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) @@ -328,7 +352,6 @@ def _ring_ragged_unsort_fwd( enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) else: # Shift indices so they map to the packed local buffer [0, local_num_tokens). @@ -347,7 +370,6 @@ def _ring_ragged_unsort_fwd( enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) res = ( @@ -406,7 +428,6 @@ def _ring_ragged_unsort_bwd(res, g_out): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) # Mask out gradients that correspond to elements outside the valid shard # output range. @@ -430,7 +451,6 @@ def _ring_ragged_unsort_bwd(res, g_out): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, ) # Mask out gradients for elements beyond the valid limit of the local buffer. limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) @@ -443,22 +463,10 @@ def _ring_ragged_unsort_bwd(res, g_out): # Build the flat weights array from the routing weights. topk_weights_flat = topk_weights.astype(jnp.float32) - return _ring_ragged_unsort( - sorted_tokens_local, - group_sizes_local, - topk_argsort_revert_indices, - topk_weights_flat, - ) + return _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat) -def a2a_ragged_sort( - inputs, - sort_indices, - valid_end, - enforce_gather_fallback=False, - enforce_gather_reduce_fallback=False, - use_single_sparsecore=False, -): +def a2a_ragged_sort(inputs, sort_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False): """Ragged-gather variant for ``local_permute``. Unlike :func:`ring_ragged_sort`, the rows valid for this shard live in @@ -498,13 +506,7 @@ def _a2a_ragged_sort(inputs, sort_indices, valid_end): def _a2a_ragged_sort_fwd(inputs, sort_indices, valid_end): start = jnp.int32(0) end = valid_end.astype(jnp.int32) if hasattr(valid_end, "astype") else jnp.int32(valid_end) - out = ragged_gather( - inputs, - sort_indices, - start[None], - end[None], - use_single_sparsecore=use_single_sparsecore, - ) + out = ragged_gather(inputs, sort_indices, start[None], end[None]) n = sort_indices.shape[0] valid_mask = jnp.arange(n) < end out = jnp.where(valid_mask[:, None], out, 0.0) @@ -528,7 +530,6 @@ def _a2a_ragged_sort_bwd(res, g_out): valid_rows_mask=valid_rows_mask[idx_inv], reduce_group_size=1, enforce_fallback=enforce_gather_reduce_fallback, - use_single_sparsecore=use_single_sparsecore, ) # custom_vjp must return one gradient per primal arg; valid_end is integer # and non-differentiable, so we return None for it. @@ -539,12 +540,7 @@ def _a2a_ragged_sort_bwd(res, g_out): def a2a_ragged_unsort( - sorted_tokens, - revert_indices, - valid_end, - enforce_gather_fallback=False, - enforce_gather_reduce_fallback=False, - use_single_sparsecore=False, + sorted_tokens, revert_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False ): """Dual of :func:`a2a_ragged_sort`. @@ -587,7 +583,6 @@ def _a2a_ragged_unsort_fwd(sorted_tokens, revert_indices, valid_end): valid_rows_mask=valid_rows_mask, reduce_group_size=1, enforce_fallback=enforce_gather_reduce_fallback, - use_single_sparsecore=use_single_sparsecore, ) res = (revert_indices, end, sorted_tokens.shape, start) return out, res @@ -599,13 +594,7 @@ def _a2a_ragged_unsort_bwd(res, g_out): # Because revert_indices is a permutation, build the inverse and use # ragged_gather to pull the per-row gradients to the right positions. idx_inv = jnp.argsort(revert_indices) - grad_sorted = ragged_gather( - g_out, - idx_inv, - start[None], - end[None], - use_single_sparsecore=use_single_sparsecore, - ) + grad_sorted = ragged_gather(g_out, idx_inv, start[None], end[None]) num_rows = sorted_tokens_shape[0] pos = jnp.arange(num_rows) valid = pos < end diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f6e89f204d..3c03bbd2be 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -868,7 +868,10 @@ def permute( # reshape inputs (batch, sequence, emb) to (batch * sequence, emb) inputs_shape = inputs.shape bsz_times_seq_len = inputs_shape[0] * inputs_shape[1] - inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2])) + if isinstance(inputs, qpl.QArray): + inputs_2d = inputs.reshape(bsz_times_seq_len, inputs.shape[-1]) + else: + inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2])) weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids) lb_loss = None if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing: @@ -891,6 +894,17 @@ def permute( inputs_2d = inputs_2d * router_scores.reshape(bsz_times_seq_len, -1) num_expert_parallelism = self.get_expert_parallelism_size() + + # Pre-quantize activations with Qwix before permute / ragged gather if quantization is configured + quantization_rule = qpl.get_current_rule("gmm") + if quantization_rule and quantization_rule.act_qtype and not isinstance(inputs_2d, qpl.QArray): + inputs_2d = qpl.quantize( + inputs_2d, + quantization_rule.act_qtype, + channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], + calibration_method=quantization_rule.act_calibration_method, + ) + # The ragged-kernel path inside permute()/unpermute() is only correct for # the ring-of-experts strategy: each shard's output is masked to its own # [start, end) range within a globally-sorted layout. When ring of experts @@ -915,22 +929,57 @@ def permute( else: buffer_size = None - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( - inputs_2d, - topk_indices_2d, - self.config.num_experts, - self.num_experts_per_tok, - self._expert_parallelism_name, - num_expert_parallelism, - buffer_size=buffer_size, - enforce_gather_fallback=self.config.ragged_gather_fallback, - enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, - gather_flops_override=self.config.ragged_gather_cost_estimate_flops, - gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, - gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, - gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, - ) + if isinstance(inputs_2d, qpl.QArray): + sorted_qvalue, group_size, sorted_selected_experts = ring_ragged_sort( + inputs_2d.qvalue, + topk_indices_2d, + self.config.num_experts, + self.num_experts_per_tok, + self._expert_parallelism_name, + num_expert_parallelism, + buffer_size=buffer_size, + enforce_gather_fallback=self.config.ragged_gather_fallback, + enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, + gather_flops_override=self.config.ragged_gather_cost_estimate_flops, + gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, + gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, + gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + ) + if inputs_2d.scale.shape[0] == inputs_2d.qvalue.shape[0]: + sorted_scale, _, _ = ring_ragged_sort( + inputs_2d.scale, + topk_indices_2d, + self.config.num_experts, + self.num_experts_per_tok, + self._expert_parallelism_name, + num_expert_parallelism, + buffer_size=buffer_size, + enforce_gather_fallback=self.config.ragged_gather_fallback, + enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, + gather_flops_override=self.config.ragged_gather_cost_estimate_flops, + gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, + gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, + gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + ) + else: + sorted_scale = inputs_2d.scale + sorted_inputs = qpl.QArray(qvalue=sorted_qvalue, scale=sorted_scale) + else: + sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( + inputs_2d, + topk_indices_2d, + self.config.num_experts, + self.num_experts_per_tok, + self._expert_parallelism_name, + num_expert_parallelism, + buffer_size=buffer_size, + enforce_gather_fallback=self.config.ragged_gather_fallback, + enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, + gather_flops_override=self.config.ragged_gather_cost_estimate_flops, + gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, + gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, + gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + ) else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -938,10 +987,20 @@ def permute( flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts sorted_selected_experts = jnp.argsort(flatten_selected_experts) # sort inputs for number of selected experts - replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0) - sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype( - self.dtype - ) + if isinstance(inputs_2d, qpl.QArray): + replicated_inputs_2d = qpl.QArray( + qvalue=jnp.repeat(inputs_2d.qvalue, self.num_experts_per_tok, axis=0), + scale=jnp.repeat(inputs_2d.scale, self.num_experts_per_tok, axis=0), + ) + sorted_inputs = qpl.QArray( + qvalue=_sort_activations(replicated_inputs_2d.qvalue, sorted_selected_experts, use_custom_sort_vjp), + scale=_sort_activations(replicated_inputs_2d.scale, sorted_selected_experts, use_custom_sort_vjp), + ) + else: + replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0) + sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype( + self.dtype + ) group_size = jnp.bincount(flatten_selected_experts, length=self.num_experts) num_tokens = bsz_times_seq_len * self.num_experts_per_tok @@ -1018,7 +1077,6 @@ def unpermute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: unsort_intermediate = _sort_activations( @@ -1083,7 +1141,6 @@ def local_permute( use_custom_sort_vjp=True, use_ragged_sort=False, ragged_buffer_factor=-1.0, - use_single_sparsecore=False, ): """Permutes tokens locally within an expert shard. @@ -1167,12 +1224,7 @@ def local_permute( # the worst-case ragged buffer. Restricting the gather to that prefix # makes both forward and backward proportional to the routed token count. valid_end = jnp.sum(local_group_size).astype(jnp.int32) - sorted_inputs = a2a_ragged_sort( - inputs, - sorted_indices, - valid_end, - use_single_sparsecore=use_single_sparsecore, - ) + sorted_inputs = a2a_ragged_sort(inputs, sorted_indices, valid_end) else: sorted_inputs = _sort_activations(inputs, sorted_indices, use_custom_sort_vjp) sorted_experts_ids = expert_indices[sorted_indices] @@ -1477,6 +1529,8 @@ def extract_vma(tensor): # Parses the varying mesh axes from JAX's type string for a tensor inside shard_map. # jax.typeof(t) renders as e.g. 'f32[128,256]{V:(expert, fsdp)}'; this extracts # ('expert', 'fsdp'). Returns () if the tensor has no varying axes. + if isinstance(tensor, qpl.QArray): + tensor = tensor.qvalue type_str = str(jax.typeof(tensor)) if "{V:" in type_str: start = type_str.index("{V:") + 3 @@ -1487,7 +1541,7 @@ def extract_vma(tensor): lhs_vma_axes = extract_vma(inputs) rhs_vma_axes = extract_vma(kernel) - if inputs.shape[0] != expert_assignments.shape[0]: + if (inputs.qvalue if isinstance(inputs, qpl.QArray) else inputs).shape[0] != expert_assignments.shape[0]: raise ValueError("The number of input tokens must match the number of expert assignments!") tokamax_group_sizes = get_tokamax_group_sizes(group_sizes, inputs, kernel) @@ -1663,9 +1717,27 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, # The ring-of-experts strategy first duplicates the inputs to all # expert shards, and then routes within each shard. - # Duplicate inputs to all expert shards. - x, logits, pre_bias_logits = tuple( - jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) for z in (x, logits, pre_bias_logits) + # Duplicate inputs to all expert shards + rule = None + if self.config.quantization and self.config.use_qwix_quantization: + rules = quantizations.get_quantization_rule(self.config) + rule = rules[0] if isinstance(rules, list) and rules else (rules if isinstance(rules, qwix.QtRule) else None) + + if rule and rule.act_qtype and not isinstance(x, qpl.QArray): + x_q = qpl.quantize( + x, + rule.act_qtype, + channelwise_axes=[] if rule.disable_channelwise_axes else [0], + calibration_method=rule.act_calibration_method, + ) + x_qvalue = jax.lax.all_gather(x_q.qvalue, axis_name=self._expert_parallelism_name, tiled=True) + x_scale = jax.lax.all_gather(x_q.scale, axis_name=self._expert_parallelism_name, tiled=True) + x = qpl.QArray(qvalue=x_qvalue, scale=x_scale) + else: + x = jax.lax.all_gather(x, axis_name=self._expert_parallelism_name, tiled=True) + + logits, pre_bias_logits = tuple( + jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) for z in (logits, pre_bias_logits) ) # "Route" tokens within each shard. @@ -1773,7 +1845,6 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: x, local_sorted_indices, group_sizes, selected_experts = RoutedMoE.local_permute( @@ -1786,7 +1857,6 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) return ( @@ -1972,7 +2042,6 @@ def unsort_output_and_ra2a( intermediate_output, jnp.argsort(route_metadata.local_sorted_indices), # pylint: disable=undefined-variable valid_end, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: local_output = _sort_activations( diff --git a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py index 9c16084e24..586a9f8e4b 100644 --- a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py +++ b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py @@ -1068,6 +1068,51 @@ def test_gmm_fused_activation( chex.assert_trees_all_close(actual, expected, atol=atol, rtol=rtol) + def test_gmm_prequantized_activation(self): + """Test GMM v2 with pre-quantized activation vs in-kernel quantization.""" + batch_size = 128 + in_size = 512 + out_size = 512 + num_groups = 4 + group_offset = 0 + + num_local_groups = num_groups - group_offset + key = jax.random.key(0) + k0, k1 = jax.random.split(key, 2) + + lhs_float = jax.random.normal(k0, (batch_size, in_size), dtype=jnp.bfloat16) + rhs_float = jax.random.normal(k1, (num_local_groups, in_size, out_size), dtype=jnp.bfloat16) + + group_sizes = get_group_sizes(batch_size, num_groups) + group_offset_arr = jnp.array([group_offset], dtype=jnp.int32) + + rhs_q, rhs_scale = quantize_tensor(rhs_float, jnp.int8, axis=1, block_size=256) + rhs_scale = jnp.expand_dims(rhs_scale, axis=2) + + # 1. In-kernel dynamic activation quantization + actual_inkernel = gmm_backend.gmm_v2( + lhs_float, + rhs_q, + group_sizes, + rhs_scale=rhs_scale, + group_offset=group_offset_arr, + maybe_quantize_lhs=True, + ) + + # 2. Pre-quantized activation + lhs_q, lhs_scale = quantize_tensor(lhs_float, jnp.int8, axis=1, block_size=in_size) + actual_prequantized = gmm_backend.gmm_v2( + lhs_q, + rhs_q, + group_sizes, + rhs_scale=rhs_scale, + lhs_scale=lhs_scale, + group_offset=group_offset_arr, + maybe_quantize_lhs=False, + ) + + assert_arrays_all_close(actual_prequantized, actual_inkernel, atol=1e-2, rtol=1e-2) + if __name__ == "__main__": absltest.main() From eb1b7871068c2f747a0f5eff8445ab7acee4a944 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 00:25:12 +0000 Subject: [PATCH 02/11] Tighten scope: keep use_single_sparsecore, drop unrelated capacity fixes 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 --- .../kernels/ragged/ragged_gather_reduce_v2.py | 41 ++++++---- src/maxtext/kernels/ragged/ragged_sort.py | 74 +++++++++++++++++-- src/maxtext/layers/moe.py | 15 +++- .../unit/pallas_mosaic_tpu_v2_kernel_test.py | 45 ----------- 4 files changed, 107 insertions(+), 68 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py index c6a1a14f9e..0594676b74 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py @@ -18,7 +18,6 @@ import dataclasses import functools -import math from typing import Any import jax @@ -176,10 +175,6 @@ def _fallback_implementation( reduce_group_size: int, ) -> jax.Array: """Fallback implementation using JAX ops for non-SparseCore TPU or small inputs.""" - if hasattr(x, "scale") and hasattr(x, "qvalue"): - import qwix # pylint: disable=import-outside-toplevel - - x = qwix.dequantize(x) out = x[indices] * topk_weights[:, None].astype(jnp.float32) out = jnp.where(valid_rows_mask[:, None], out, 0) out = out.reshape(-1, reduce_group_size, out.shape[-1]) @@ -201,17 +196,27 @@ def _calculate_num_column_partitions( preferred_num_stages = 4 num_column_partitions = 1 while ( - (num_cores == 2 or num_cores % (num_column_partitions * 2) == 0) + num_cores % (num_column_partitions * 2) == 0 and hidden_size % (num_lanes * num_column_partitions * 2) == 0 and hidden_size // (num_column_partitions * 2 * num_lanes) >= preferred_num_stages ): next_candidate = num_column_partitions * 2 + next_row_partitions = num_cores // next_candidate + + # Calculate exactly how many pipeline invocations (outer loop) + _, row_chunk_size = _calculate_row_tiling(input_size, num_simd_lanes, next_row_partitions) + num_iterations = input_size // (row_chunk_size * next_row_partitions) # Ensure we satisfy the hardware constraint (num_row_partitions <= num_simd_lanes) first. if num_cores // num_column_partitions > num_simd_lanes: num_column_partitions = next_candidate continue + # Too many iterations cause high cumulative pipeline overhead. Set the + # limit based on empirical data. + if num_iterations > _CostModelConstants.MAX_ITERATIONS: + break + num_column_partitions = next_candidate return num_column_partitions @@ -223,8 +228,8 @@ def _calculate_row_tiling( num_row_partitions: int, ) -> tuple[int, int]: """Calculates the number of row subchunks and row chunk size.""" - base_block_size = max(1, num_simd_lanes * num_row_partitions) - num_row_subchunks = max(1, min(2, math.ceil(int(input_size) / base_block_size))) + base_block_size = num_simd_lanes * num_row_partitions + num_row_subchunks = max(1, min(4, pl.cdiv(input_size, base_block_size))) row_chunk_size = num_simd_lanes * num_row_subchunks return num_row_subchunks, row_chunk_size @@ -620,7 +625,14 @@ def col_loop(col_compute_offset): @functools.partial( - jax.jit, static_argnames=("reduce_group_size", "enforce_fallback", "flops_override", "bytes_accessed_override") + jax.jit, + static_argnames=( + "reduce_group_size", + "enforce_fallback", + "flops_override", + "bytes_accessed_override", + "use_single_sparsecore", + ), ) def ragged_gather_reduce( x: jax.Array, @@ -631,6 +643,7 @@ def ragged_gather_reduce( enforce_fallback: bool = False, flops_override: int = -1, bytes_accessed_override: int = -1, + use_single_sparsecore: bool = False, ) -> jax.Array: """Gathers ``x`` by ``indices``, weights and masks, then reduces by group. @@ -668,12 +681,12 @@ def ragged_gather_reduce( input_size = indices.size num_simd_lanes = sc_info.num_lanes num_lanes = pltpu.get_tpu_info().num_lanes - num_cores = sc_info.num_cores * sc_info.num_subcores + num_sc_cores = 1 if use_single_sparsecore else sc_info.num_cores + num_cores = num_sc_cores * sc_info.num_subcores num_column_partitions = _calculate_num_column_partitions(hidden_size, input_size, num_cores, num_lanes, num_simd_lanes) num_row_partitions = num_cores // num_column_partitions - if num_row_partitions > num_simd_lanes: - return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size) + assert num_row_partitions <= num_simd_lanes, f"{num_row_partitions=} must be <= {num_simd_lanes=}" num_row_subchunks, row_chunk_size = _calculate_row_tiling(input_size, num_simd_lanes, num_row_partitions) aligned_hidden_size = _align_to(hidden_size, 128 * num_column_partitions) @@ -707,7 +720,7 @@ def ragged_gather_reduce( # Step 4: Launch the SparseCore kernel. vector_mesh = plsc.VectorSubcoreMesh( - num_cores=sc_info.num_cores, + num_cores=num_sc_cores, num_subcores=sc_info.num_subcores, core_axis_name="core", subcore_axis_name="subcore", @@ -747,7 +760,7 @@ def ragged_gather_reduce( flops_override=flops_override, bytes_accessed_override=bytes_accessed_override, ), - scratch_types=( + scratch_types=( # pyrefly: ignore[bad-argument-type] _Scratch( num_rows_per_row_partition_vmem=pltpu.VMEM((num_simd_lanes,), jnp.int32), prev_iter_last_row_vmem=pltpu.VMEM((col_size // col_chunk_size, col_chunk_size), jnp.float32), diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index 62b8e3ea5a..4d5940b0b3 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -35,6 +35,7 @@ def ring_ragged_sort( gather_reduce_flops_override=-1, gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, + use_single_sparsecore=False, ): """Ragged-gather variant for AG-RS Expert Parallelism token routing. @@ -113,6 +114,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) x_scale = ragged_gather( hidden_states_local.scale, @@ -122,6 +124,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) x = qpl.QArray(qvalue=x_qval, scale=x_scale) else: @@ -133,6 +136,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) else: local_buffer_size = buffer_size @@ -157,6 +161,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) x_scale = ragged_gather( hidden_states_local.scale, @@ -166,6 +171,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) x = qpl.QArray(qvalue=x_qval, scale=x_scale) else: @@ -177,6 +183,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) out = (x, group_sizes_local, topk_argsort_revert_indices) @@ -231,6 +238,7 @@ def _ring_ragged_sort_bwd(res, g_out): enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) else: # Buffering: g_x has size `local_buffer_size` (packed). @@ -256,6 +264,7 @@ def _ring_ragged_sort_bwd(res, g_out): enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) return grad_hidden_states, None @@ -278,6 +287,7 @@ def ring_ragged_unsort( gather_reduce_flops_override=-1, gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, + use_single_sparsecore=False, ): """Dual of :func:`ring_ragged_sort`. @@ -309,14 +319,27 @@ def ring_ragged_unsort( """ @jax.custom_vjp - def _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat): + def _ring_ragged_unsort( + sorted_tokens_local, + group_sizes_local, + topk_argsort_revert_indices, + topk_weights_flat, + ): """Unsort and scatter activations.""" return _ring_ragged_unsort_fwd( - sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat + sorted_tokens_local, + group_sizes_local, + topk_argsort_revert_indices, + topk_weights_flat, )[0] @jax.named_scope("ragged-unsort-fwd") - def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat): + def _ring_ragged_unsort_fwd( + sorted_tokens_local, + group_sizes_local, + topk_argsort_revert_indices, + topk_weights_flat, + ): """Executes unsorting sending tokens back.""" group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) @@ -352,6 +375,7 @@ def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) else: # Shift indices so they map to the packed local buffer [0, local_num_tokens). @@ -370,6 +394,7 @@ def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort enforce_fallback=enforce_gather_reduce_fallback, flops_override=gather_reduce_flops_override, bytes_accessed_override=gather_reduce_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) res = ( @@ -428,6 +453,7 @@ def _ring_ragged_unsort_bwd(res, g_out): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) # Mask out gradients that correspond to elements outside the valid shard # output range. @@ -451,6 +477,7 @@ def _ring_ragged_unsort_bwd(res, g_out): enforce_fallback=enforce_gather_fallback, flops_override=gather_flops_override, bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, ) # Mask out gradients for elements beyond the valid limit of the local buffer. limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) @@ -463,10 +490,22 @@ def _ring_ragged_unsort_bwd(res, g_out): # Build the flat weights array from the routing weights. topk_weights_flat = topk_weights.astype(jnp.float32) - return _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat) + return _ring_ragged_unsort( + sorted_tokens_local, + group_sizes_local, + topk_argsort_revert_indices, + topk_weights_flat, + ) -def a2a_ragged_sort(inputs, sort_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False): +def a2a_ragged_sort( + inputs, + sort_indices, + valid_end, + enforce_gather_fallback=False, + enforce_gather_reduce_fallback=False, + use_single_sparsecore=False, +): """Ragged-gather variant for ``local_permute``. Unlike :func:`ring_ragged_sort`, the rows valid for this shard live in @@ -506,7 +545,13 @@ def _a2a_ragged_sort(inputs, sort_indices, valid_end): def _a2a_ragged_sort_fwd(inputs, sort_indices, valid_end): start = jnp.int32(0) end = valid_end.astype(jnp.int32) if hasattr(valid_end, "astype") else jnp.int32(valid_end) - out = ragged_gather(inputs, sort_indices, start[None], end[None]) + out = ragged_gather( + inputs, + sort_indices, + start[None], + end[None], + use_single_sparsecore=use_single_sparsecore, + ) n = sort_indices.shape[0] valid_mask = jnp.arange(n) < end out = jnp.where(valid_mask[:, None], out, 0.0) @@ -530,6 +575,7 @@ def _a2a_ragged_sort_bwd(res, g_out): valid_rows_mask=valid_rows_mask[idx_inv], reduce_group_size=1, enforce_fallback=enforce_gather_reduce_fallback, + use_single_sparsecore=use_single_sparsecore, ) # custom_vjp must return one gradient per primal arg; valid_end is integer # and non-differentiable, so we return None for it. @@ -540,7 +586,12 @@ def _a2a_ragged_sort_bwd(res, g_out): def a2a_ragged_unsort( - sorted_tokens, revert_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False + sorted_tokens, + revert_indices, + valid_end, + enforce_gather_fallback=False, + enforce_gather_reduce_fallback=False, + use_single_sparsecore=False, ): """Dual of :func:`a2a_ragged_sort`. @@ -583,6 +634,7 @@ def _a2a_ragged_unsort_fwd(sorted_tokens, revert_indices, valid_end): valid_rows_mask=valid_rows_mask, reduce_group_size=1, enforce_fallback=enforce_gather_reduce_fallback, + use_single_sparsecore=use_single_sparsecore, ) res = (revert_indices, end, sorted_tokens.shape, start) return out, res @@ -594,7 +646,13 @@ def _a2a_ragged_unsort_bwd(res, g_out): # Because revert_indices is a permutation, build the inverse and use # ragged_gather to pull the per-row gradients to the right positions. idx_inv = jnp.argsort(revert_indices) - grad_sorted = ragged_gather(g_out, idx_inv, start[None], end[None]) + grad_sorted = ragged_gather( + g_out, + idx_inv, + start[None], + end[None], + use_single_sparsecore=use_single_sparsecore, + ) num_rows = sorted_tokens_shape[0] pos = jnp.arange(num_rows) valid = pos < end diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 3c03bbd2be..f8c481aa66 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -944,6 +944,7 @@ def permute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) if inputs_2d.scale.shape[0] == inputs_2d.qvalue.shape[0]: sorted_scale, _, _ = ring_ragged_sort( @@ -960,6 +961,7 @@ def permute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: sorted_scale = inputs_2d.scale @@ -979,6 +981,7 @@ def permute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -1077,6 +1080,7 @@ def unpermute( gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: unsort_intermediate = _sort_activations( @@ -1141,6 +1145,7 @@ def local_permute( use_custom_sort_vjp=True, use_ragged_sort=False, ragged_buffer_factor=-1.0, + use_single_sparsecore=False, ): """Permutes tokens locally within an expert shard. @@ -1224,7 +1229,12 @@ def local_permute( # the worst-case ragged buffer. Restricting the gather to that prefix # makes both forward and backward proportional to the routed token count. valid_end = jnp.sum(local_group_size).astype(jnp.int32) - sorted_inputs = a2a_ragged_sort(inputs, sorted_indices, valid_end) + sorted_inputs = a2a_ragged_sort( + inputs, + sorted_indices, + valid_end, + use_single_sparsecore=use_single_sparsecore, + ) else: sorted_inputs = _sort_activations(inputs, sorted_indices, use_custom_sort_vjp) sorted_experts_ids = expert_indices[sorted_indices] @@ -1845,6 +1855,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: x, local_sorted_indices, group_sizes, selected_experts = RoutedMoE.local_permute( @@ -1857,6 +1868,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in use_custom_sort_vjp=self.config.use_custom_sort_vjp, use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) return ( @@ -2042,6 +2054,7 @@ def unsort_output_and_ra2a( intermediate_output, jnp.argsort(route_metadata.local_sorted_indices), # pylint: disable=undefined-variable valid_end, + use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) else: local_output = _sort_activations( diff --git a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py index 586a9f8e4b..9c16084e24 100644 --- a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py +++ b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py @@ -1068,51 +1068,6 @@ def test_gmm_fused_activation( chex.assert_trees_all_close(actual, expected, atol=atol, rtol=rtol) - def test_gmm_prequantized_activation(self): - """Test GMM v2 with pre-quantized activation vs in-kernel quantization.""" - batch_size = 128 - in_size = 512 - out_size = 512 - num_groups = 4 - group_offset = 0 - - num_local_groups = num_groups - group_offset - key = jax.random.key(0) - k0, k1 = jax.random.split(key, 2) - - lhs_float = jax.random.normal(k0, (batch_size, in_size), dtype=jnp.bfloat16) - rhs_float = jax.random.normal(k1, (num_local_groups, in_size, out_size), dtype=jnp.bfloat16) - - group_sizes = get_group_sizes(batch_size, num_groups) - group_offset_arr = jnp.array([group_offset], dtype=jnp.int32) - - rhs_q, rhs_scale = quantize_tensor(rhs_float, jnp.int8, axis=1, block_size=256) - rhs_scale = jnp.expand_dims(rhs_scale, axis=2) - - # 1. In-kernel dynamic activation quantization - actual_inkernel = gmm_backend.gmm_v2( - lhs_float, - rhs_q, - group_sizes, - rhs_scale=rhs_scale, - group_offset=group_offset_arr, - maybe_quantize_lhs=True, - ) - - # 2. Pre-quantized activation - lhs_q, lhs_scale = quantize_tensor(lhs_float, jnp.int8, axis=1, block_size=in_size) - actual_prequantized = gmm_backend.gmm_v2( - lhs_q, - rhs_q, - group_sizes, - rhs_scale=rhs_scale, - lhs_scale=lhs_scale, - group_offset=group_offset_arr, - maybe_quantize_lhs=False, - ) - - assert_arrays_all_close(actual_prequantized, actual_inkernel, atol=1e-2, rtol=1e-2) - if __name__ == "__main__": absltest.main() From 595d36f29c80f5ddc4ac761a5733cd71e213b137 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 00:33:28 +0000 Subject: [PATCH 03/11] Restore PR #4735's static-LHS-scaling mechanism, layered alongside ours 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 --- src/maxtext/kernels/megablox/ops.py | 41 +++- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 180 ++++++++++++++---- 2 files changed, 188 insertions(+), 33 deletions(-) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index c295e58e1c..a002d0c6ac 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -211,7 +211,7 @@ def _gmm_fwd( out = _fwd_run_tokamax_v1(lhs, rhs, group_sizes, preferred_element_type, transpose_rhs, use_manual_quantization) elif use_tokamax_backend and use_gmm_v2: out = _fwd_run_tokamax_v2( - lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs + lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset, partial_sum, transpose_rhs, quantization_rule ) else: out = _fwd_run_megablox( @@ -335,6 +335,36 @@ def _fwd_prepare_rhs_scale(rhs: qpl.QArray, transpose_rhs: bool = False) -> jnp. return jnp.broadcast_to(rhs_scale, (G, num_quant_blocks, 1, N)) +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. + + If a static scale is used, GMM v2 requires it to be from a symmetric fixed-range + calibration (e.g., 'fixed,-max,max' or 'fixed,max'). If no static scale is + provided, the kernel will compute a dynamic scale on the fly. + + Enforces a default (1, 1) shape for per-tensor quantization kernels. + + Args: + quantization_rule: The Qwix quantization rule from which to extract the scale. + + Returns: + The extracted static scale array, or None if not using purely fixed calibration. + """ + if quantization_rule is None: + return None + + method = quantization_rule.act_calibration_method + qtype = quantization_rule.act_qtype + + # Use dynamic quantization, gmm_v2 calculates dynamic scale internally + if method is None or qtype is None or not method.lower().startswith("fixed"): + return None + + scale_val = quantizations.get_static_scale(qtype, method) + + return jnp.full((1, 1), scale_val, jnp.float32) + + def _fwd_run_tokamax_v2( lhs: jnp.ndarray | qpl.QArray, rhs: jnp.ndarray | qpl.QArray, @@ -344,6 +374,7 @@ def _fwd_run_tokamax_v2( group_offset: jnp.ndarray | None, partial_sum: jnp.ndarray | None, transpose_rhs: bool, + quantization_rule: qwix.QtRule | None = None, ) -> jnp.ndarray: """Executes the Tokamax GMM V2 backend for forward pass OUT = LHS @ RHS.""" # if transpose_rhs=False, rhs is [g, k, n], remain unchanged @@ -361,6 +392,13 @@ def _fwd_run_tokamax_v2( jnp.issubdtype(lhs_operand.dtype, jnp.integer) or jnp.issubdtype(lhs_operand.dtype, jnp.float8_e4m3fn) ) + # When lhs is not pre-quantized (ahead-of-time by the caller), fall back to + # a static scale from the quantization rule's fixed calibration if one is + # configured; the kernel uses it directly instead of computing a dynamic + # per-block absmax. If lhs is already a QArray, its scale is applied + # externally below, so no in-kernel scale is needed. + lhs_scale = None if isinstance(lhs, qpl.QArray) else _fwd_prepare_lhs_scale(quantization_rule) + custom_fwd_tiling = gmm_v2.TileSizes( tile_m=tiling[0], tile_k=tiling[1], @@ -383,6 +421,7 @@ def _fwd_run_tokamax_v2( partial_sum=partial_sum, group_offset=group_offset, maybe_quantize_lhs=maybe_quantize_lhs, + lhs_scale=lhs_scale, ) if isinstance(lhs, qpl.QArray): diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 24567d6272..49b7845c0b 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== # 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 """GMM kernel implemented using Pallas.""" from abc import ABC, abstractmethod @@ -141,6 +141,29 @@ def get_bias(self) -> jax.Array: return jnp.concatenate([b_gate, b_up], axis=-1) +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class LhsRef: + """Dataclass for the lhs value and its optional quantization scale. + + 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 + per-channel `[M, 1]` and sub-channel `[M, num_blocks]`). + """ + + value: Any + scale: Any | None + + def get_value(self) -> jax.Array: + return self.value[...] + + def get_scale(self) -> jax.Array: + assert self.scale is not None + return self.scale[...] + + @jax.tree_util.register_dataclass @dataclasses.dataclass(frozen=True) class MetadataRef: @@ -173,8 +196,21 @@ class InputConfigs: quant_block_size: int | None dtype: jnp.dtype has_bias: bool = False + # 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). has_scale: bool = False + @property + def should_use_external_scale(self) -> bool: + # A scale is present but the input is not yet quantized + # (dtype != quant_dtype). The kernel uses it to quantize the input online + # and multiply the result by the scale after. This differs from an already + # quantized input (dtype == quant_dtype), whose scale only dequantizes after + # the matmul. + return self.has_scale and self.quant_dtype is not None and self.dtype != self.quant_dtype + @property def should_bitcast(self) -> bool: bits = jax.dtypes.itemsize_bits(self.dtype) @@ -239,6 +275,14 @@ def lhs_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): return (pl.ds(row_start, row_size), 0, k_id) + def lhs_scale_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array): + # Per-tensor scale: a single [1, 1] value shared across every tile, so the + # block always reads index 0. Extension point: when the scale is per-channel + # or sub-channel, tile the row axis like `lhs_index_map` (using gm_id) and + # index the K-block axis from `k_id`. + del gm_id, k_id + return (0, 0) + def rhs_weight_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array): group_id = self.metadata_ref.gm_id_to_group_id[gm_id] return (group_id, k_id, n_id) @@ -283,16 +327,23 @@ def ps_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array): def generate_block_specs( metadata_ref: MetadataRef, cfgs: GmmConfigs -) -> Tuple[Tuple[pl.BlockSpec, WeightsRef, pl.BlockSpec | None], pl.BlockSpec]: +) -> Tuple[Tuple[LhsRef, WeightsRef, pl.BlockSpec | None], pl.BlockSpec]: """Generates block specs for the given lhs, rhs, and out refs.""" index_map = IndexMaps(metadata_ref, cfgs) bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane) - lhs_block_spec = pl.BlockSpec( + lhs_value_spec = pl.BlockSpec( (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k), index_map.lhs_index_map, ) + lhs_scale_spec = None + if cfgs.lhs_cfgs.has_scale: + lhs_scale_spec = pl.BlockSpec( + (1, 1), + index_map.lhs_scale_index_map, + ) + lhs_block_spec = LhsRef(value=lhs_value_spec, scale=lhs_scale_spec) tile_k_rhs = cfgs.tiles.tile_k if cfgs.rhs_cfgs.should_bitcast: @@ -341,7 +392,7 @@ def generate_block_specs( def inner_kernel( # In - tiled_lhs_ref: jax.Array, + tiled_lhs_ref: LhsRef, # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k] tiled_rhs_ref: RhsRef, # [tile_k, tile_n] # Partial Sum @@ -382,7 +433,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): mxu_size = tpu_info.mxu_column_size # Step 1: Input pre-processing. - tiled_lhs = tiled_lhs_ref.reshape(-1, cfgs.tiles.tile_k)[...] + tiled_lhs = tiled_lhs_ref.get_value().reshape(-1, cfgs.tiles.tile_k)[...] tiled_rhs = tiled_rhs_ref.get_weight() # When rhs is packed (quantized dtype packed into uint32), unpack it # back to the original dtype using pltpu.bitcast which operates on K @@ -446,6 +497,14 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): dtype_max = float(jnp.iinfo(lhs_q_dtype).max) preferred_element_type = jnp.int32 + # When the caller supplies a quantization scale, use it directly instead + # of computing a dynamic per-block absmax. + lhs_scale = lhs_scale_inv = None + should_use_external_scale = cfgs.lhs_cfgs.should_use_external_scale + if should_use_external_scale: + lhs_scale = tiled_lhs_ref.get_scale().astype(acc_ref.dtype) + lhs_scale_inv = 1.0 / lhs_scale + # Without n outer loop, result of quantized matmul becomes available only # at the last iteration of the loop. This means [tile_m, tile_n] value # needs to be stored until the last iteration. By adding n outer loop, @@ -467,8 +526,16 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): # same computation will be performed tiles_n//mxu_size times. # But we can let compiler perform CSE and avoid recomputation. if jnp.issubdtype(tiled_lhs.dtype, jnp.integer) or jnp.issubdtype(tiled_lhs.dtype, jnp.float8_e4m3fn): + # lhs block already quantized upstream (pre-quantized ahead-of-time + # by the caller); the real dequant scale is applied externally, so + # just pass the block through with an identity scale here. block_lhs_q = block_lhs block_scale = jnp.array(1.0, dtype=acc_ref.dtype) + elif should_use_external_scale: + assert lhs_scale is not None + assert lhs_scale_inv is not None + block_lhs_q = jnp.clip(block_lhs * lhs_scale_inv, -dtype_max, dtype_max).astype(lhs_q_dtype) + block_scale = lhs_scale # [1, 1] else: block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True) block_scale = block_abs_max / dtype_max @@ -770,7 +837,7 @@ def kernel_main( lhs_group_sizes_ref: jax.Array, # int32[size_lhs_group] group_offset_ref: jax.Array, # int32[1] # In - lhs_ref: jax.Array, # [size_m, size_k] + lhs_ref: LhsRef, # value: [size_m, size_k] rhs_ref: WeightsRef, # [size_group, size_k, size_n] partial_sum_ref: jax.Array, # [size_m, size_n] # Out @@ -859,8 +926,10 @@ def kernel_main( ) # Bounded slice requires second last dim to be aligned to the sublane size. - # rhs_ref uses static tiling thus reshape is not needed. - lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.shape[-1]) + # rhs_ref uses static tiling thus reshape is not needed. The lhs quant scale + # (when present) is small and statically tiled, so it is passed through as-is. + lhs_value_in = lhs_ref.value.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.value.shape[-1]) + lhs_in = LhsRef(value=lhs_value_in, scale=lhs_ref.scale) ps_in = None if cfgs.has_partial_sum: ps_in = partial_sum_ref.reshape(-1, cfgs.dims.size_lhs_sublane, partial_sum_ref.shape[-1]) @@ -993,6 +1062,8 @@ def validate_inputs( group_sizes: jax.Array, group_offset: jax.Array, fuse_act: str | None = None, + maybe_quantize_lhs: bool = True, + lhs_scale: jax.Array | None = None, ) -> Dimensions: """Validates the inputs for the GMM kernel.""" @@ -1011,9 +1082,18 @@ def validate_inputs( assert partial_sum.shape[0] <= size_m if rhs_scale is not None: num_quant_blocks = rhs_scale.shape[1] - assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n) + assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n), ( + f"rhs_scale shape {rhs_scale.shape}. Expecting ({size_group}," f" {num_quant_blocks}, 1, {size_n})" + ) assert size_k % num_quant_blocks == 0 + if lhs_scale is not None: + assert maybe_quantize_lhs, "lhs_scale requires maybe_quantize_lhs=True." + # Only per-tensor scales are supported for now. The current implementation generalizes to per-channel [M, 1] and + # sub-channel [M, num_k_blocks]; extend the validation and the block spec / + # index map together when adding those. + assert lhs_scale.shape == (1, 1), "Only per-tensor lhs_scale of shape (1, 1) is supported, got " f"{lhs_scale.shape}." + assert group_offset.shape == (1,) size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype) @@ -1096,10 +1176,22 @@ def make_gmm_configs( maybe_quantize_lhs: bool, zero_initialize: bool, fuse_act: str | None = None, + lhs_scale: jax.Array | None = None, ): """Fills the GMM config for the GMM kernel.""" - dims = validate_inputs(lhs, rhs, rhs_scale, rhs_bias, partial_sum, group_sizes, group_offset, fuse_act) + dims = validate_inputs( + lhs, + rhs, + rhs_scale, + rhs_bias, + partial_sum, + group_sizes, + group_offset, + fuse_act, + maybe_quantize_lhs, + lhs_scale, + ) if rhs_scale is not None: has_scale = True @@ -1121,6 +1213,8 @@ def make_gmm_configs( lhs_q_dtype = None if jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn): + # lhs arrives already quantized (e.g. pre-quantized ahead-of-time by the + # caller): use its dtype as-is, no in-kernel quantization/scale needed. lhs_q_dtype = lhs.dtype elif maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul: # Choose lhs quantization dtype based on TPU hardware support. @@ -1138,6 +1232,14 @@ def make_gmm_configs( if not is_rhs_float: lhs_q_dtype = jnp.int8.dtype + if lhs_scale is not None: + assert lhs_q_dtype is not None, ( + "lhs_scale requires lhs quantization to engage, but no lhs quant " + "dtype was selected. Ensure rhs is quantized and the hardware supports " + "fp8/int8 matmul." + ) + has_lhs_scale = lhs_scale is not None and lhs_q_dtype is not None + lhs_cfgs = InputConfigs( quant_dtype=lhs_q_dtype, # Input quantization involves reading all elements in a block to compute @@ -1146,9 +1248,14 @@ def make_gmm_configs( # enough to minimize compute overhead of quantization. quant_block_size=512, dtype=lhs.dtype, + has_scale=has_lhs_scale, ) if out_dtype is None or jnp.issubdtype(out_dtype, jnp.float8_e4m3fn): + # The raw quantized-domain matmul output isn't yet rescaled -- writing it + # directly as fp8 would lose precision before the scale multiply happens + # (either inside this kernel via lhs_scale/block_scale, or externally by + # the caller for a pre-quantized lhs). Floor to bf16 as a safe intermediate. out_dtype = jnp.bfloat16.dtype if acc_dtype is None: @@ -1209,6 +1316,7 @@ def gmm_v2( rhs_bias: jax.Array | None = None, # [size_group, 1, out_size] partial_sum: jax.Array | None = None, # [size_m, size_n] group_offset: jax.Array | None = None, # int32[1] + lhs_scale: jax.Array | None = None, # [1, 1] (per-tensor) *, tile_info: TileSizes | TileFn = calculate_tiling, # pyrefly: ignore[bad-function-definition] vmem_limit_bytes: int | None = None, @@ -1233,6 +1341,12 @@ def gmm_v2( rhs_bias: The rhs bias of shape [size_group, 1, out_size]. partial_sum: Optional. Per-token partial sums of shape [size_m, size_n]. group_offset: Optional. The group offset of shape [1,]. + lhs_scale: Optional scale used to quantize the (unquantized) lhs + inside the kernel and the result is multiplied back by `scale`. The shape + encodes granularity; currently only per-tensor `[1, 1]` is supported. When + None, a quantized lhs uses the default dynamic per-block absmax + calibration. Only takes effect when maybe_quantize_lhs is True and rhs is + quantized. tile_info: The tile sizes or tile function to use. vmem_limit_bytes: Optional vmem limit in bytes. precision: Unused. Exists for compatibility reasons. @@ -1272,11 +1386,20 @@ def gmm_v2( maybe_quantize_lhs=maybe_quantize_lhs, zero_initialize=zero_initialize, fuse_act=fuse_act, + lhs_scale=lhs_scale, ) dims = cfgs.dims tiles = cfgs.tiles # Prepare block specs. + lhs_scale_spec = None + if cfgs.lhs_cfgs.has_scale: + assert lhs_scale is not None + lhs_scale = lhs_scale.astype(jnp.float32) + lhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM) + else: + lhs_scale = None + rhs_scale_spec = rhs_bias_spec = None if rhs_scale is not None: rhs_scale = rhs_scale.astype(jnp.float32) @@ -1328,33 +1451,15 @@ def gmm_v2( aligned_n = align_to(cfgs.out_size_n, num_lanes) out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n), cfgs.out_dtype) + lhs_in = LhsRef(value=lhs, scale=lhs_scale) rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias) - in_specs = [ - pl.BlockSpec(memory_space=pltpu.HBM), - WeightsRef( - weight=pl.BlockSpec(memory_space=pltpu.HBM), - scale=rhs_scale_spec, - bias=rhs_bias_spec, - ), - ] - partial_sum_spec = None if partial_sum is not None: - in_specs.append(pl.BlockSpec(memory_space=pltpu.HBM)) partial_sum_spec = pl.BlockSpec(memory_space=pltpu.HBM) - in_specs = [ - pl.BlockSpec(memory_space=pltpu.HBM), # lhs - WeightsRef( - weight=pl.BlockSpec(memory_space=pltpu.HBM), - scale=rhs_scale_spec, - bias=rhs_bias_spec, - ), # rhs_weights - partial_sum_spec, # partial_sum - ] input_output_aliases = {} if partial_sum is not None: - flat_args_preceding = (group_sizes, group_offset, lhs, rhs_weights) + flat_args_preceding = (group_sizes, group_offset, lhs_in, rhs_weights) leaves = jax.tree_util.tree_leaves(flat_args_preceding) partial_sum_idx = sum(1 for x in leaves if x is not None) input_output_aliases = {partial_sum_idx: 0} @@ -1364,7 +1469,18 @@ def gmm_v2( out_shape=out_init, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, - in_specs=in_specs, + in_specs=[ + LhsRef( + value=pl.BlockSpec(memory_space=pltpu.HBM), + scale=lhs_scale_spec, + ), + WeightsRef( + weight=pl.BlockSpec(memory_space=pltpu.HBM), + scale=rhs_scale_spec, + bias=rhs_bias_spec, + ), + partial_sum_spec, + ], out_specs=pl.BlockSpec(memory_space=pltpu.HBM), scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), @@ -1376,4 +1492,4 @@ def gmm_v2( cost_estimate=get_cost_estimate(cfgs), metadata=get_metadata(cfgs), input_output_aliases=input_output_aliases, - )(group_sizes, group_offset, lhs, rhs_weights, partial_sum)[:, : cfgs.out_size_n] + )(group_sizes, group_offset, lhs_in, rhs_weights, partial_sum)[:, : cfgs.out_size_n] From 59a7acb2dd3b1b81ae0f71fd8f5674a36ecc3564 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 00:49:09 +0000 Subject: [PATCH 04/11] Use qwix's should_quantize() for the "already quantized" check in ops.py Replaces the ops.py-level jnp.issubdtype(int) or jnp.issubdtype(e4m3fn) enumeration with qwix._src.core.numerics.should_quantize(dtype), which generalizes correctly to any quantized dtype (e5m2, int4, etc.) instead of only recognizing float8_e4m3fn specifically -- consistent with the call_with_generic_broadcast private-qwix-import precedent already in this file. Did NOT apply the same change to pallas_mosaic_tpu_v2_gmm_kernel.py: that file's make_gmm_configs/inner_kernel are shared by the backward DLHS/DRHS gmm_v2 calls, where gradients are quantized to bwd_qtype=float8_e5m2 (per the fp8_full recipe) -- a different dtype than the forward act_qtype (float8_e4m3fn). The old, narrower check only recognized e4m3fn as "already quantized", so e5m2 gradients fell through to the dynamic in-kernel quantization branch (safe). Generalizing it there too causes the kernel to treat e5m2 lhs as already-quantized and hit an unsupported Mosaic cast (float8_e4m3fn -> float8_e5m2) when pairing it with e4m3fn rhs -- confirmed by reproducing and reverting just those two call sites. Left as the original enumeration there, verified via a full AOT compile check (train_compile_671b_v6e64.sh) both ways. Co-Authored-By: Claude Sonnet 5 --- src/maxtext/kernels/megablox/ops.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index a002d0c6ac..d6d3905178 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -27,6 +27,7 @@ from maxtext.layers import quantizations import qwix import qwix.pallas as qpl +from qwix._src.core import numerics as qwix_numerics from qwix._src.core.qarray import call_with_generic_broadcast import tokamax @@ -239,11 +240,7 @@ def _fwd_quantize_activation_and_weight( transpose_rhs: bool, ) -> tuple[jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray]: """Handles act and weight quantization for GMM forward inputs.""" - if ( - quantization_rule.act_qtype - and not isinstance(lhs, qpl.QArray) - and not (jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn)) - ): + if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs.dtype): lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, @@ -388,9 +385,7 @@ def _fwd_run_tokamax_v2( rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs) lhs_operand = lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs - maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and not ( - jnp.issubdtype(lhs_operand.dtype, jnp.integer) or jnp.issubdtype(lhs_operand.dtype, jnp.float8_e4m3fn) - ) + maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs_operand.dtype) # When lhs is not pre-quantized (ahead-of-time by the caller), fall back to # a static scale from the quantization rule's fixed calibration if one is @@ -405,11 +400,7 @@ def _fwd_run_tokamax_v2( tile_n=tiling[2], ) - eff_pref_dtype = ( - jnp.bfloat16 - if (jnp.issubdtype(lhs_operand.dtype, jnp.integer) or jnp.issubdtype(lhs_operand.dtype, jnp.float8_e4m3fn)) - else preferred_element_type - ) + eff_pref_dtype = preferred_element_type if qwix_numerics.should_quantize(lhs_operand.dtype) else jnp.bfloat16 out = gmm_v2.gmm_v2( lhs=lhs_operand, # pyrefly: ignore[bad-argument-type] @@ -593,7 +584,7 @@ def _bwd_prepare_inputs( quantization_rule and quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray) - and not (jnp.issubdtype(lhs.dtype, jnp.integer) or jnp.issubdtype(lhs.dtype, jnp.float8_e4m3fn)) + and qwix_numerics.should_quantize(lhs.dtype) ): lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, From f438e644ab0737e51ead8df3c5baca1ed8ab59dd Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 00:55:09 +0000 Subject: [PATCH 05/11] moe.py: dedup ring_ragged_sort call sites, drop redundant/unsafe quantize block 1. Factor the 3 near-identical ring_ragged_sort call sites in permute() (qvalue, scale, non-QArray) into a local _permute_ring_ragged_sort() closure over the shared config args (topk_indices_2d, buffer_size, etc). 2. Remove permute()'s own "pre-quantize before ragged gather" block. It was redundant for roe_ag_and_route (ring-of-experts): that caller already quantizes x into a QArray before its EP all-gather, so by the time permute() ran this block, `not isinstance(inputs_2d, qpl.QArray)` was always False and it never actually fired. Worse, it was live but unsafe for ra2a_and_route (the non-ring-of-experts caller), which passes x in raw: this block would quantize it into a QArray that then flows into either jax.lax.ragged_all_to_all (a raw XLA primitive that doesn't understand a QArray pytree) or local_permute()/a2a_ragged_sort (which, unlike ring_ragged_sort, has no QArray dual-dispatch at all). Since the feature is specifically "quantize before EP all-gather" for the ring-of-experts path, and roe_ag_and_route already owns that correctly, removing this block closes the gap rather than leaving latent-broken code in the unrelated ra2a path. Verified via a full AOT compile check (train_compile_671b_v6e64.sh, deepseek3-671b fp8, v6e-64, ring-of-experts) -- same memory footprint as before, confirming roe_ag_and_route's own quantization is sufficient on its own. Co-Authored-By: Claude Sonnet 5 --- src/maxtext/layers/moe.py | 53 ++++++--------------------------------- 1 file changed, 8 insertions(+), 45 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index f8c481aa66..fa6cf2682e 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -895,16 +895,6 @@ def permute( num_expert_parallelism = self.get_expert_parallelism_size() - # Pre-quantize activations with Qwix before permute / ragged gather if quantization is configured - quantization_rule = qpl.get_current_rule("gmm") - if quantization_rule and quantization_rule.act_qtype and not isinstance(inputs_2d, qpl.QArray): - inputs_2d = qpl.quantize( - inputs_2d, - quantization_rule.act_qtype, - channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], - calibration_method=quantization_rule.act_calibration_method, - ) - # The ragged-kernel path inside permute()/unpermute() is only correct for # the ring-of-experts strategy: each shard's output is masked to its own # [start, end) range within a globally-sorted layout. When ring of experts @@ -929,9 +919,9 @@ def permute( else: buffer_size = None - if isinstance(inputs_2d, qpl.QArray): - sorted_qvalue, group_size, sorted_selected_experts = ring_ragged_sort( - inputs_2d.qvalue, + def _permute_ring_ragged_sort(tensor): + return ring_ragged_sort( + tensor, topk_indices_2d, self.config.num_experts, self.num_experts_per_tok, @@ -946,43 +936,16 @@ def permute( gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, ) + + if isinstance(inputs_2d, qpl.QArray): + sorted_qvalue, group_size, sorted_selected_experts = _permute_ring_ragged_sort(inputs_2d.qvalue) if inputs_2d.scale.shape[0] == inputs_2d.qvalue.shape[0]: - sorted_scale, _, _ = ring_ragged_sort( - inputs_2d.scale, - topk_indices_2d, - self.config.num_experts, - self.num_experts_per_tok, - self._expert_parallelism_name, - num_expert_parallelism, - buffer_size=buffer_size, - enforce_gather_fallback=self.config.ragged_gather_fallback, - enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, - gather_flops_override=self.config.ragged_gather_cost_estimate_flops, - gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, - gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, - gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, - ) + sorted_scale, _, _ = _permute_ring_ragged_sort(inputs_2d.scale) else: sorted_scale = inputs_2d.scale sorted_inputs = qpl.QArray(qvalue=sorted_qvalue, scale=sorted_scale) else: - sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort( - inputs_2d, - topk_indices_2d, - self.config.num_experts, - self.num_experts_per_tok, - self._expert_parallelism_name, - num_expert_parallelism, - buffer_size=buffer_size, - enforce_gather_fallback=self.config.ragged_gather_fallback, - enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback, - gather_flops_override=self.config.ragged_gather_cost_estimate_flops, - gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops, - gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, - gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, - use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, - ) + sorted_inputs, group_size, sorted_selected_experts = _permute_ring_ragged_sort(inputs_2d) else: flatten_selected_experts = jnp.ravel(selected_experts) From 09dc56e6509faba7d4b454fe0594875cece2a767 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 18:19:11 +0000 Subject: [PATCH 06/11] Add quantize_before_ep_all_gather flag; use qpl.get_current_rule in roe_ag_and_route - New config flag `quantize_before_ep_all_gather` (default True) gates roe_ag_and_route's ahead-of-time activation quantization, so it can be disabled to quantize inside the gmm call instead (e.g. for an A/B numerical-equivalence test comparing the two, or as a rollback lever). - roe_ag_and_route now derives its quantization rule via qpl.get_current_rule("gmm") instead of manually re-deriving it from static config via quantizations.get_quantization_rule(self.config) + list/QtRule unwrapping -- simpler, and consistent with how the rest of this file (and ops.py) already fetches the active rule. Did NOT apply a similar qpl.get_current_rule()-based rewrite to permute()'s scale-sort condition (attempted, then reverted): disable_channelwise_axes does not reliably predict whether a QArray's materialized scale ends up per-row. Confirmed via debug prints on the real fp8_full/absmax recipe -- scale came out shape (8, 1) (looks like one value per expert-parallel shard) while disable_channelwise_axes=False, so a rule-based check would have ring_ragged_sort'd an 8-row tensor against indices built for a 32768-row space, producing a garbage-shaped result and a downstream "input tokens != expert assignments" compile failure. The existing shape-based check (scale.shape[0] == qvalue.shape[0]) verifies the actual invariant directly and is more robust; left unchanged. Verified via a full AOT compile check (train_compile_671b_v6e64.sh, deepseek3-671b fp8, v6e-64) -- same memory footprint as before. Co-Authored-By: Claude Sonnet 5 --- src/maxtext/configs/base.yml | 2 ++ src/maxtext/configs/types.py | 7 +++++++ src/maxtext/layers/moe.py | 7 ++----- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 08d22bae24..6bb0765530 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -226,6 +226,8 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss use_random_routing: false # whether to use random routing for debug/test purpose use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism +quantize_before_ep_all_gather: true # whether to quantize activations before the ring-of-experts EP all-gather + # (fp8 collective + ragged sort) vs. quantizing later inside the gmm call num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe). diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 68684cb38a..f0465579ba 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -867,6 +867,13 @@ class MoEGeneral(BaseModel): False, description="Whether to use Ring of Experts for sparse matmul expert parallelism.", ) + quantize_before_ep_all_gather: bool = Field( + True, + description=( + "Whether to quantize activations before the Ring of Experts EP all-gather (so the " + "collective and ragged sort move fp8, not bf16), vs. quantizing later inside the gmm call." + ), + ) moe_dispatch_no_expert_sharding: bool = Field( False, description=( diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index fa6cf2682e..2d08e6d860 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1691,12 +1691,9 @@ def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, # expert shards, and then routes within each shard. # Duplicate inputs to all expert shards - rule = None - if self.config.quantization and self.config.use_qwix_quantization: - rules = quantizations.get_quantization_rule(self.config) - rule = rules[0] if isinstance(rules, list) and rules else (rules if isinstance(rules, qwix.QtRule) else None) + rule = qpl.get_current_rule("gmm") - if rule and rule.act_qtype and not isinstance(x, qpl.QArray): + if self.config.quantize_before_ep_all_gather and rule and rule.act_qtype and not isinstance(x, qpl.QArray): x_q = qpl.quantize( x, rule.act_qtype, From 207b9434a149d07d648b73997ca41c40933f6806 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 19:04:57 +0000 Subject: [PATCH 07/11] Add test_quantize_before_ep_all_gather_equivalence Compares quantizing activations ahead of the ring-of-experts EP all-gather (this branch's default) against quantizing them dynamically inside the gmm call (quantize_before_ep_all_gather=False), with everything else identical (same weights/inputs/EP degree/routing). Verifies the ahead-of-time relocation is a pure optimization, not a semantic change. Confirmed on real TPU hardware: forward output and every weight gradient are bit-exact (relative_norm_diff=0) between the two configs, as expected since channelwise/absmax scale is computed per-token and doesn't depend on when quantization happens. Uses chex.assert_trees_all_close (absolute + relative tolerance) rather than the existing compare_tree helper's pure relative-norm metric: the hidden-state gradient's reference norm sits near the float32 noise floor for this loss formulation, so a relative-only metric reported a spurious 400%+ "mismatch" on what was actually ~1e-9 absolute floating-point reordering noise (verified via a debug script that dumped the full per-leaf diff before landing on the tolerance choice). Co-Authored-By: Claude Sonnet 5 --- tests/unit/moe_test.py | 131 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 7781f4b797..500eff63a1 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -15,6 +15,7 @@ import unittest from absl.testing import parameterized +import chex import pytest from flax import nnx @@ -1623,6 +1624,136 @@ def loss_fn(params, x): diff_summary = compare_tree(tree_ref, tree_tgt, relative_norm_diff_threshold) max_logging.log("\n" + diff_summary) + @pytest.mark.skip_on_tpu7x + @pytest.mark.tpu_only + def test_quantize_before_ep_all_gather_equivalence(self): + """Quantizing activations before the ring-of-experts EP all-gather should be + numerically equivalent to quantizing them later inside the gmm call: both + sides quantize the same per-token values (channelwise/absmax scale), just + at a different point in the pipeline. Unlike test_gmm_grad_equivalence + (which compares quantized vs. unquantized and expects real quantization + noise), any meaningful gap here is a bug in the ahead-of-time relocation, + not expected noise -- so this uses a much tighter tolerance. + """ + + def _build_cfg(quantize_before_ep_all_gather): + return pyconfig.initialize( + [None, get_test_config_path()], + run_name="quantize_before_ep_all_gather_equivalence_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + weight_dtype="float32", + dtype="bfloat16", + per_device_batch_size=2, + max_target_length=256, + float32_gate_logits=True, + ici_expert_parallelism=4, + sparse_matmul=True, + megablox=True, + use_tokamax_gmm=True, + use_gmm_v2=True, + use_ring_of_experts=True, + use_ragged_sort=True, + quantization="fp8_full", + use_qwix_quantization=True, + weight_quantization_calibration_method="absmax", + act_quantization_calibration_method="absmax", + bwd_quantization_calibration_method="absmax", + quantize_before_ep_all_gather=quantize_before_ep_all_gather, + wi_tile_fwd_batch_seq=128, + wi_tile_dlhs_batch_seq=128, + wi_tile_dlhs_embed_dim=256, + wi_tile_drhs_batch_seq=128, + wo_tile_fwd_batch_seq=128, + wo_tile_fwd_embed_dim=256, + wo_tile_dlhs_batch_seq=128, + wo_tile_dlhs_mlp_dim=256, + wo_tile_drhs_batch_seq=128, + ) + + def _build_model(cfg, mesh): + model = moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + + # Similar to `quantizations.get_fp8_full_qwix_rule_w_sparsity`. + def get_fp8_full_qwix_rule_for_test(config): + return [ + qwix.QtRule( + module_path=".*", + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + bwd_qtype=jnp.float8_e5m2, + weight_calibration_method=config.weight_quantization_calibration_method, + act_calibration_method=config.act_quantization_calibration_method, + bwd_calibration_method=config.bwd_quantization_calibration_method, + op_names=("gmm", "ragged_dot"), + ), + ] + + quantization_provider = qwix.QtProvider(get_fp8_full_qwix_rule_for_test(cfg)) + model = qwix.quantize_model(model, quantization_provider) + return model + + def _loss_and_grad(model, variables, hidden_states): + def loss_fn(params, x): + out, lb_loss, _ = model.apply({"params": params}, x) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + return jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(variables["params"], hidden_states) + + rng = jax.random.PRNGKey(4567) + rng_model, rng_hidden_states = jax.random.split(rng) + device_count = jax.device_count() + + # Reference run: quantize dynamically inside the gmm call (activations stay + # bf16 through the EP all-gather and ragged sort). + cfg_ref = _build_cfg(quantize_before_ep_all_gather=False) + # Normal distribution for realistic variance/negative values, so the + # quantization scale != 1.0 and scale-dropping bugs are actually caught. + hidden_states = jax.random.normal( + rng_hidden_states, + (int(cfg_ref.per_device_batch_size) * device_count, cfg_ref.max_target_length, cfg_ref.base_emb_dim), + dtype=cfg_ref.dtype, + ) + devices_array_ref = maxtext_utils.create_device_mesh(cfg_ref) + mesh_ref = Mesh(devices_array_ref, cfg_ref.mesh_axes) + model_ref = _build_model(cfg_ref, mesh_ref) + with jax.set_mesh(mesh_ref), nn_partitioning.axis_rules(cfg_ref.logical_axis_rules): + variables = model_ref.init({"params": rng_model, "dropout": rng_model}, hidden_states) + (_, output_ref), (grads_ref, x_grad_ref) = _loss_and_grad(model_ref, variables, hidden_states) + + # Target run: quantize ahead of the EP all-gather (this branch's default). + cfg_tgt = _build_cfg(quantize_before_ep_all_gather=True) + devices_array_tgt = maxtext_utils.create_device_mesh(cfg_tgt) + mesh_tgt = Mesh(devices_array_tgt, cfg_tgt.mesh_axes) + model_tgt = _build_model(cfg_tgt, mesh_tgt) + with jax.set_mesh(mesh_tgt), nn_partitioning.axis_rules(cfg_tgt.logical_axis_rules): + # Re-initialize for the target mesh, but with the same RNG so the + # initial weights match the reference run's. + variables_tgt = model_tgt.init({"params": rng_model, "dropout": rng_model}, hidden_states) + (_, output_tgt), (grads_tgt, x_grad_tgt) = _loss_and_grad(model_tgt, variables_tgt, hidden_states) + + tree_ref = {"output": output_ref, "state_grad": x_grad_ref, "var_grad": grads_ref} + tree_tgt = {"output": output_tgt, "state_grad": x_grad_tgt, "var_grad": grads_tgt} + # Use an absolute+relative tolerance (not compare_tree's pure relative-norm + # metric): the hidden-state gradient's reference norm is near the float32 + # noise floor for this loss, so a relative-only metric blows up on + # ordinary floating-point reordering noise even when both sides agree to + # ~1e-9 in absolute terms. + chex.assert_trees_all_close(tree_tgt, tree_ref, atol=1e-6, rtol=1e-3) + def make_moe(cfg, mesh): return moe.RoutedMoE( From 76f209778fbd2787525dd55641f7f25c76e66cfc Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 20:32:09 +0000 Subject: [PATCH 08/11] update test --- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 7 +- tests/unit/moe_test.py | 193 +++++++++--------- 2 files changed, 95 insertions(+), 105 deletions(-) diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 49b7845c0b..cdd89179c2 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -525,10 +525,9 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): # Perform lhs quantization. Note that for every block_lhs, # same computation will be performed tiles_n//mxu_size times. # But we can let compiler perform CSE and avoid recomputation. - if jnp.issubdtype(tiled_lhs.dtype, jnp.integer) or jnp.issubdtype(tiled_lhs.dtype, jnp.float8_e4m3fn): - # lhs block already quantized upstream (pre-quantized ahead-of-time - # by the caller); the real dequant scale is applied externally, so - # just pass the block through with an identity scale here. + if tiled_lhs.dtype == lhs_q_dtype: + # lhs block already arrives quantized, the real dequant + # scale is applied externally, so just pass an identity scale here. Comparing against lhs_q_dtype block_lhs_q = block_lhs block_scale = jnp.array(1.0, dtype=acc_ref.dtype) elif should_use_external_scale: diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 500eff63a1..49b99851d5 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -507,6 +507,46 @@ def get_moe_output(self, variables, hidden_states, cfg, mesh): output = jax.jit(model.apply)(moe_variables, hidden_states) # pylint: disable=not-callable return output + def get_quantized_moe_model(self, cfg, mesh): + """Builds a RoutedMoE wrapped with the fp8_full qwix quantization rule.""" + model = moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + # Similar to `quantizations.get_fp8_full_qwix_rule_w_sparsity`. + quantization_rule = [ + qwix.QtRule( + module_path=".*", + weight_qtype=jnp.float8_e4m3fn, + act_qtype=jnp.float8_e4m3fn, + bwd_qtype=jnp.float8_e5m2, + weight_calibration_method=cfg.weight_quantization_calibration_method, + act_calibration_method=cfg.act_quantization_calibration_method, + bwd_calibration_method=cfg.bwd_quantization_calibration_method, + op_names=("gmm", "ragged_dot"), + ), + ] + return qwix.quantize_model(model, qwix.QtProvider(quantization_rule)) + + def get_moe_loss_and_grad(self, model, variables, hidden_states): + """Computes (loss, output) and grads w.r.t. params and the input hidden states.""" + + def loss_fn(params, x): + out, lb_loss, _ = model.apply({"params": params}, x) + loss = jnp.mean(out.astype(jnp.float32) ** 2) + if lb_loss is not None: + loss = loss + lb_loss.astype(jnp.float32) + return loss, out + + return jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(variables["params"], hidden_states) + @pytest.mark.tpu_only def test_megablox(self): cfg = pyconfig.initialize( @@ -1624,6 +1664,49 @@ def loss_fn(params, x): diff_summary = compare_tree(tree_ref, tree_tgt, relative_norm_diff_threshold) max_logging.log("\n" + diff_summary) + def _build_ep_all_gather_test_cfg(self, quantize_before_ep_all_gather): + return pyconfig.initialize( + [None, get_test_config_path()], + run_name="quantize_before_ep_all_gather_equivalence_test", + enable_checkpointing=False, + model_name="mixtral-8x7b", + weight_dtype="float32", + dtype="bfloat16", + per_device_batch_size=2, + max_target_length=256, + float32_gate_logits=True, + ici_expert_parallelism=4, + sparse_matmul=True, + megablox=True, + use_tokamax_gmm=True, + use_gmm_v2=True, + use_ring_of_experts=True, + use_ragged_sort=True, + quantization="fp8_full", + use_qwix_quantization=True, + weight_quantization_calibration_method="absmax", + act_quantization_calibration_method="absmax", + bwd_quantization_calibration_method="absmax", + quantize_before_ep_all_gather=quantize_before_ep_all_gather, + wi_tile_fwd_batch_seq=128, + wi_tile_dlhs_batch_seq=128, + wi_tile_dlhs_embed_dim=256, + wi_tile_drhs_batch_seq=128, + wo_tile_fwd_batch_seq=128, + wo_tile_fwd_embed_dim=256, + wo_tile_dlhs_batch_seq=128, + wo_tile_dlhs_mlp_dim=256, + wo_tile_drhs_batch_seq=128, + ) + + def _run_ep_all_gather_test(self, quantize_before_ep_all_gather, rng_model, hidden_states): + cfg = self._build_ep_all_gather_test_cfg(quantize_before_ep_all_gather) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + model = self.get_quantized_moe_model(cfg, mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + variables = model.init({"params": rng_model, "dropout": rng_model}, hidden_states) + return self.get_moe_loss_and_grad(model, variables, hidden_states) + @pytest.mark.skip_on_tpu7x @pytest.mark.tpu_only def test_quantize_before_ep_all_gather_equivalence(self): @@ -1635,115 +1718,23 @@ def test_quantize_before_ep_all_gather_equivalence(self): noise), any meaningful gap here is a bug in the ahead-of-time relocation, not expected noise -- so this uses a much tighter tolerance. """ - - def _build_cfg(quantize_before_ep_all_gather): - return pyconfig.initialize( - [None, get_test_config_path()], - run_name="quantize_before_ep_all_gather_equivalence_test", - enable_checkpointing=False, - model_name="mixtral-8x7b", - weight_dtype="float32", - dtype="bfloat16", - per_device_batch_size=2, - max_target_length=256, - float32_gate_logits=True, - ici_expert_parallelism=4, - sparse_matmul=True, - megablox=True, - use_tokamax_gmm=True, - use_gmm_v2=True, - use_ring_of_experts=True, - use_ragged_sort=True, - quantization="fp8_full", - use_qwix_quantization=True, - weight_quantization_calibration_method="absmax", - act_quantization_calibration_method="absmax", - bwd_quantization_calibration_method="absmax", - quantize_before_ep_all_gather=quantize_before_ep_all_gather, - wi_tile_fwd_batch_seq=128, - wi_tile_dlhs_batch_seq=128, - wi_tile_dlhs_embed_dim=256, - wi_tile_drhs_batch_seq=128, - wo_tile_fwd_batch_seq=128, - wo_tile_fwd_embed_dim=256, - wo_tile_dlhs_batch_seq=128, - wo_tile_dlhs_mlp_dim=256, - wo_tile_drhs_batch_seq=128, - ) - - def _build_model(cfg, mesh): - model = moe.get_routed_moe( - name="MoeBlock", - config=cfg, - num_experts=cfg.num_experts, - num_experts_per_tok=cfg.num_experts_per_tok, - mesh=mesh, - kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), - kernel_axes=("embed", "mlp"), - intermediate_dim=cfg.mlp_dim, - dtype=cfg.dtype, - ) - - # Similar to `quantizations.get_fp8_full_qwix_rule_w_sparsity`. - def get_fp8_full_qwix_rule_for_test(config): - return [ - qwix.QtRule( - module_path=".*", - weight_qtype=jnp.float8_e4m3fn, - act_qtype=jnp.float8_e4m3fn, - bwd_qtype=jnp.float8_e5m2, - weight_calibration_method=config.weight_quantization_calibration_method, - act_calibration_method=config.act_quantization_calibration_method, - bwd_calibration_method=config.bwd_quantization_calibration_method, - op_names=("gmm", "ragged_dot"), - ), - ] - - quantization_provider = qwix.QtProvider(get_fp8_full_qwix_rule_for_test(cfg)) - model = qwix.quantize_model(model, quantization_provider) - return model - - def _loss_and_grad(model, variables, hidden_states): - def loss_fn(params, x): - out, lb_loss, _ = model.apply({"params": params}, x) - loss = jnp.mean(out.astype(jnp.float32) ** 2) - if lb_loss is not None: - loss = loss + lb_loss.astype(jnp.float32) - return loss, out - - return jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(variables["params"], hidden_states) - rng = jax.random.PRNGKey(4567) rng_model, rng_hidden_states = jax.random.split(rng) - device_count = jax.device_count() - - # Reference run: quantize dynamically inside the gmm call (activations stay - # bf16 through the EP all-gather and ragged sort). - cfg_ref = _build_cfg(quantize_before_ep_all_gather=False) + cfg = self._build_ep_all_gather_test_cfg(quantize_before_ep_all_gather=False) # Normal distribution for realistic variance/negative values, so the # quantization scale != 1.0 and scale-dropping bugs are actually caught. hidden_states = jax.random.normal( rng_hidden_states, - (int(cfg_ref.per_device_batch_size) * device_count, cfg_ref.max_target_length, cfg_ref.base_emb_dim), - dtype=cfg_ref.dtype, + (int(cfg.per_device_batch_size) * jax.device_count(), cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, ) - devices_array_ref = maxtext_utils.create_device_mesh(cfg_ref) - mesh_ref = Mesh(devices_array_ref, cfg_ref.mesh_axes) - model_ref = _build_model(cfg_ref, mesh_ref) - with jax.set_mesh(mesh_ref), nn_partitioning.axis_rules(cfg_ref.logical_axis_rules): - variables = model_ref.init({"params": rng_model, "dropout": rng_model}, hidden_states) - (_, output_ref), (grads_ref, x_grad_ref) = _loss_and_grad(model_ref, variables, hidden_states) - # Target run: quantize ahead of the EP all-gather (this branch's default). - cfg_tgt = _build_cfg(quantize_before_ep_all_gather=True) - devices_array_tgt = maxtext_utils.create_device_mesh(cfg_tgt) - mesh_tgt = Mesh(devices_array_tgt, cfg_tgt.mesh_axes) - model_tgt = _build_model(cfg_tgt, mesh_tgt) - with jax.set_mesh(mesh_tgt), nn_partitioning.axis_rules(cfg_tgt.logical_axis_rules): - # Re-initialize for the target mesh, but with the same RNG so the - # initial weights match the reference run's. - variables_tgt = model_tgt.init({"params": rng_model, "dropout": rng_model}, hidden_states) - (_, output_tgt), (grads_tgt, x_grad_tgt) = _loss_and_grad(model_tgt, variables_tgt, hidden_states) + # Reference: quantize dynamically inside the gmm call (activations stay + # bf16 through the EP all-gather and ragged sort). Target: quantize ahead + # of the EP all-gather (this branch's default). Same RNG for both, so the + # initial weights match. + (_, output_ref), (grads_ref, x_grad_ref) = self._run_ep_all_gather_test(False, rng_model, hidden_states) + (_, output_tgt), (grads_tgt, x_grad_tgt) = self._run_ep_all_gather_test(True, rng_model, hidden_states) tree_ref = {"output": output_ref, "state_grad": x_grad_ref, "var_grad": grads_ref} tree_tgt = {"output": output_tgt, "state_grad": x_grad_tgt, "var_grad": grads_tgt} From 5728eac1f627f958344b6ccf8d673229cd0fc1f7 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 20:45:22 +0000 Subject: [PATCH 09/11] validate --- src/maxtext/configs/types.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f0465579ba..a9bb27a089 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2848,6 +2848,16 @@ def validate_ragged_buffer_factor(self): " 2. Ragged sort with ring of experts (use_ring_of_experts=True AND use_ragged_sort=True)" ) + def _validate_quantize_before_ep_all_gather(self): + """Validates quantize_before_ep_all_gather is used with supported settings.""" + if self.quantize_before_ep_all_gather and not ( + self.use_ring_of_experts and self.use_qwix_quantization and self.use_gmm_v2 + ): + raise ValueError( + "quantize_before_ep_all_gather=True is only supported with use_ring_of_experts=True and " + "qwix quantization, and gmm v2 kernel" + ) + def _validate_use_te_comm_gemm_overlap(self): """Validates that use_te_comm_gemm_overlap is used with supported settings to enable TE Collective GEMM ops.""" te_has_distributed_env = jax.local_device_count() == 1 and jax.distributed.is_initialized() @@ -3955,6 +3965,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if self.use_batch_split_schedule: raise ValueError("GMM v2 is not supported with a batch split schedule.") + self._validate_quantize_before_ep_all_gather() + for val in self.compress_ratios: if val != 0 and val < 4: raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}") From b1ed0f0a5823224138971f54690f186f5e1bf7ae Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Tue, 11 Aug 2026 22:02:01 +0000 Subject: [PATCH 10/11] simplify --- src/maxtext/kernels/ragged/ragged_sort.py | 114 +++++----------------- src/maxtext/layers/moe.py | 31 +++--- 2 files changed, 42 insertions(+), 103 deletions(-) diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index 4d5940b0b3..b0ced6f607 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -18,7 +18,6 @@ import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce -import qwix.pallas as qpl def ring_ragged_sort( @@ -103,41 +102,21 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): shard_output_start = group_offsets[experts_start] shard_output_end = group_offsets[experts_end] + def _gather(indices, start, end): + return ragged_gather( + hidden_states_local, + indices, + start, + end, + enforce_fallback=enforce_gather_fallback, + flops_override=gather_flops_override, + bytes_accessed_override=gather_bytes_accessed_override, + use_single_sparsecore=use_single_sparsecore, + ) + if buffer_size is None or buffer_size >= num_tokens_local * topk: local_buffer_size = num_tokens_local * topk - if isinstance(hidden_states_local, qpl.QArray): - x_qval = ragged_gather( - hidden_states_local.qvalue, - token_indices_sorted, - shard_output_start[None], - shard_output_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - x_scale = ragged_gather( - hidden_states_local.scale, - token_indices_sorted, - shard_output_start[None], - shard_output_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - x = qpl.QArray(qvalue=x_qval, scale=x_scale) - else: - x = ragged_gather( - hidden_states_local, - token_indices_sorted, - shard_output_start[None], - shard_output_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + x = _gather(token_indices_sorted, shard_output_start[None], shard_output_end[None]) else: local_buffer_size = buffer_size # We only gather up to the available buffer size or the actual number of @@ -152,39 +131,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): local_buffer_size, axis=0, ) - if isinstance(hidden_states_local, qpl.QArray): - x_qval = ragged_gather( - hidden_states_local.qvalue, - sliced_indices, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - x_scale = ragged_gather( - hidden_states_local.scale, - sliced_indices, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) - x = qpl.QArray(qvalue=x_qval, scale=x_scale) - else: - x = ragged_gather( - hidden_states_local, - sliced_indices, - jnp.int32(0)[None], - gather_end[None], - enforce_fallback=enforce_gather_fallback, - flops_override=gather_flops_override, - bytes_accessed_override=gather_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + x = _gather(sliced_indices, jnp.int32(0)[None], gather_end[None]) out = (x, group_sizes_local, topk_argsort_revert_indices) @@ -221,17 +168,11 @@ def _ring_ragged_sort_bwd(res, g_out): # rather than materializing a (mostly-zero) dense buffer ourselves. n = topk_argsort_revert_indices.shape[0] - if local_buffer_size >= n: - valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( - topk_argsort_revert_indices < shard_output_end - ) - # The forward scatter-add over `token_indices_sorted` is equivalent to a - # gather-reduce: each input token has exactly `topk` contributions located - # at sorted positions `topk_argsort_revert_indices[t*topk:(t+1)*topk]`. - # `topk_weights` is set to ones because this op has no per-row weighting. - grad_hidden_states = ragged_gather_reduce( + def _gather_reduce(indices, valid_rows_mask): + """`topk_weights` is set to ones because this op has no per-row weighting.""" + return ragged_gather_reduce( g_x, - topk_argsort_revert_indices, + indices, topk_weights=jnp.ones((n,), dtype=jnp.float32), valid_rows_mask=valid_rows_mask, reduce_group_size=topk, @@ -240,6 +181,12 @@ def _ring_ragged_sort_bwd(res, g_out): bytes_accessed_override=gather_reduce_bytes_accessed_override, use_single_sparsecore=use_single_sparsecore, ) + + if local_buffer_size >= n: + valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) + grad_hidden_states = _gather_reduce(topk_argsort_revert_indices, valid_rows_mask) else: # Buffering: g_x has size `local_buffer_size` (packed). # The revert indices are global [0, n), but they must map to the local @@ -254,18 +201,7 @@ def _ring_ragged_sort_bwd(res, g_out): # Clamp invalid indices to 0 to prevent compile-time/run-time out-of-bounds # in JAX. These clamped values will be ignored due to `valid_rows_mask`. safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) - - grad_hidden_states = ragged_gather_reduce( - g_x, - safe_indices, - topk_weights=jnp.ones((n,), dtype=jnp.float32), - valid_rows_mask=valid_rows_mask, - reduce_group_size=topk, - enforce_fallback=enforce_gather_reduce_fallback, - flops_override=gather_reduce_flops_override, - bytes_accessed_override=gather_reduce_bytes_accessed_override, - use_single_sparsecore=use_single_sparsecore, - ) + grad_hidden_states = _gather_reduce(safe_indices, valid_rows_mask) return grad_hidden_states, None _ring_ragged_sort.defvjp(_ring_ragged_sort_fwd, _ring_ragged_sort_bwd) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 2d08e6d860..cf451d57ec 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -112,6 +112,17 @@ def _truncate_matrix(all_shards_group_sizes: jax.Array, buffer_size: int) -> jax return jnp.diff(clamped_cumsum_extended, axis=0) +def _map_qarray(fn, value: jax.Array | qpl.QArray) -> jax.Array | qpl.QArray: + """Applies `fn` to `value`. + + If `value` is a QArray, applies `fn` to `qvalue` and `scale` independently + and rewraps the result, instead of requiring the caller to unpack/repack it. + """ + if isinstance(value, qpl.QArray): + return qpl.QArray(qvalue=fn(value.qvalue), scale=fn(value.scale)) + return fn(value) + + def _sort_activations( inputs: jax.Array, sort_indices: jax.Array, @@ -953,20 +964,12 @@ def _permute_ring_ragged_sort(tensor): flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts sorted_selected_experts = jnp.argsort(flatten_selected_experts) # sort inputs for number of selected experts - if isinstance(inputs_2d, qpl.QArray): - replicated_inputs_2d = qpl.QArray( - qvalue=jnp.repeat(inputs_2d.qvalue, self.num_experts_per_tok, axis=0), - scale=jnp.repeat(inputs_2d.scale, self.num_experts_per_tok, axis=0), - ) - sorted_inputs = qpl.QArray( - qvalue=_sort_activations(replicated_inputs_2d.qvalue, sorted_selected_experts, use_custom_sort_vjp), - scale=_sort_activations(replicated_inputs_2d.scale, sorted_selected_experts, use_custom_sort_vjp), - ) - else: - replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0) - sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype( - self.dtype - ) + replicated_inputs_2d = _map_qarray(lambda t: jnp.repeat(t, self.num_experts_per_tok, axis=0), inputs_2d) + sorted_inputs = _map_qarray( + lambda t: _sort_activations(t, sorted_selected_experts, use_custom_sort_vjp), replicated_inputs_2d + ) + if not isinstance(sorted_inputs, qpl.QArray): + sorted_inputs = sorted_inputs.astype(self.dtype) group_size = jnp.bincount(flatten_selected_experts, length=self.num_experts) num_tokens = bsz_times_seq_len * self.num_experts_per_tok From 1ed8bb58970f6165cc06f0edc23bbdf88a9c72c0 Mon Sep 17 00:00:00 2001 From: Shuwen-Fang Date: Wed, 12 Aug 2026 00:27:45 +0000 Subject: [PATCH 11/11] fix --- run_qwen3_next_80b_xpk.sh | 152 ++++++++++++++++++++++++++++ src/maxtext/kernels/megablox/ops.py | 7 +- 2 files changed, 153 insertions(+), 6 deletions(-) create mode 100755 run_qwen3_next_80b_xpk.sh diff --git a/run_qwen3_next_80b_xpk.sh b/run_qwen3_next_80b_xpk.sh new file mode 100755 index 0000000000..22a2711862 --- /dev/null +++ b/run_qwen3_next_80b_xpk.sh @@ -0,0 +1,152 @@ +#!/bin/bash +set -e + +# Activate Python virtual environment +source /home/shuwenf_google_com/venv-maxtext/bin/activate + +# --- Environment Variables --- +export PROJECT_ID="tpu-prod-env-one-vm" +export CLUSTER_NAME="bodaborg-v6e-256-lcscld-c" +export ZONE="southamerica-west1-a" + +# --- Configuration & Automated Image Build --- +TIMESTAMP=$(date +%m%d%H%M%S) +export WORKLOAD_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:shuwenf_${TIMESTAMP}" +export WORKLOAD_NAME="shuwenf-qn80b-${TIMESTAMP}" +export DEVICE_TYPE="v6e-256" +export NUM_SLICES=1 +export PRIORITY="very-high" +export NUM_STEPS=15 +export MAX_RESTARTS=${MAX_RESTARTS:-0} +export MODEL_NAME="qwen3-next-80b-a3b" +export BASE_OUTPUT_DIR="gs://shuwenf-hlo-dumps/qwen3-next-80b-profiles/run-${TIMESTAMP}" + +echo "========================================================================" +echo "Building and uploading Docker runner image for full 15-step execution..." +echo "Target Image: ${WORKLOAD_IMAGE}" +echo "========================================================================" + +( + cd /home/shuwenf_google_com/maxtext && sudo CLOUD_IMAGE_NAME="${WORKLOAD_IMAGE}" BASE_IMAGE="gcr.io/tpu-prod-env-one-vm/param3_21jul:latest" bash src/dependencies/scripts/docker_upload_runner.sh +) + +echo "Docker image upload complete: ${WORKLOAD_IMAGE}" + +# --- XLA Flags --- +XLA_FLAGS_ARRAY=( + "--xla_tpu_scheduler_percent_shared_memory_limit=35" + "--xla_msa_enable_sync_slice_replacement=false" + "--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true" + "--xla_msa_enable_sync_copy_replacement=false" + "--xla_tpu_scoped_vmem_limit_kib=81000" + "--xla_tpu_enable_sparse_core_collective_offload_all_gather=true" + "--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true" + "--xla_tpu_offload_gather_to_sparsecore=true" + "--xla_tpu_dvfs_p_state=7" + "--xla_tpu_disable_sparse_core_collective_offload_remover=true" + "--xla_tpu_enable_async_collective_fusion=true" + "--xla_tpu_overlap_compute_collective_tc=true" + "--xla_tpu_enable_async_collective_fusion_multiple_steps=true" + "--xla_tpu_enable_latency_hiding_scheduler=true" + "--xla_latency_hiding_scheduler_rerun=10" + "--xla_tpu_all_gather_collective_matmul_mode=post_spmd_conservative" + "--xla_tpu_reduce_scatter_collective_matmul_mode=post_spmd_conservative" + "--xla_latency_hiding_scheduler_enable_selective_resources=true" + "--xla_tpu_enable_ilp_latency_hiding_scheduler=true" + "--xla_tpu_enable_all_experimental_scheduler_features=true" + "--xla_tpu_enable_scheduler_memory_pressure_tracking=true" + "--xla_tpu_host_transfer_overlap_limit=24" + "--xla_tpu_aggressive_opt_barrier_removal=ENABLED" + "--xla_lhs_prioritize_async_depth_over_stall=DISABLED" + "--xla_tpu_enable_ag_backward_pipelining=true" + "--xla_should_allow_loop_variant_parameter_in_chain=ENABLED" + "--xla_should_add_loop_invariant_op_in_chain=ENABLED" + "--xla_max_concurrent_host_send_recv=100" +) +export XLA_FLAGS="${XLA_FLAGS_ARRAY[*]}" + +# --- MaxText Workload Overrides --- +MAXTEXT_ARGS_ARRAY=( + "model_name=${MODEL_NAME}" + "base_output_directory=${BASE_OUTPUT_DIR}" + "run_name=param-3" + "dataset_type=synthetic" + "dataset_name=synthetic" + "dtype=bfloat16" + "allow_split_physical_axes=True" + "ici_expert_parallelism=4" + "use_ring_of_experts=True" + "custom_mesh=hybrid_ring_64x4" + "use_ragged_sort=True" + "use_random_routing=True" + "per_device_batch_size=4" + "opt_type=muon" + "muon_consistent_rms=0.2" + "muon_weight_decay=0.1" + "learning_rate=1e-5" + "max_target_length=2048" + "ragged_buffer_factor=1.5" + "remat_policy=full" + "reuse_example_batch=1" + "decoder_layer_input=offload" + "context=device" + "ici_fsdp_parallelism=-1" + "steps=15" + "shard_exp_on_fsdp=True" + "sharding_tolerance=0.5" + "sa_q_layout=SEQ_MINOR" + "sa_k_layout=HEAD_DIM_MINOR" + "sa_v_layout=HEAD_DIM_MINOR" + "sa_block_q=2048" + "sa_block_kv=2048" + "sa_block_kv_compute=1024" + "sa_block_q_dkv=2048" + "sa_block_kv_dkv=2048" + "sa_block_kv_dkv_compute=1024" + "hardware=tpu" + "skip_jax_distributed_system=False" + "attention=flash" + "use_tokamax_splash=True" + "sa_use_fused_bwd_kernel=True" + "use_tokamax_gmm=True" + "use_gmm_v2=True" + "sparse_matmul=True" + "megablox=True" + "optimizer_memory_host_offload=True" + "parameter_memory_host_offload=False" + "enable_checkpointing=False" + "async_checkpointing=False" + "tokenizer_type=tiktoken" + "tokenizer_path=tokenizer_74B/" + "override_model_config=true" + "mhc_expansion_rate=4" + "profiler=xplane" + "profiler_steps=5" + "skip_first_n_steps_for_profiler=2" + "enable_tpu_profiling_options=True" + "upload_all_profiler_results=true" +) +MAXTEXT_ARGS="${MAXTEXT_ARGS_ARRAY[*]}" + +# Clean container temporary log setup safely +RUN_COMMAND="set -e && rm -rf /tmp/tpu_logs/* 2>/dev/null || true; mkdir -p /tmp/tpu_logs && export LIBTPU_INIT_ARGS=\"${XLA_FLAGS}\" && export JAX_PLATFORMS='tpu,cpu' && export ENABLE_PJRT_COMPATIBILITY='true' && export JAX_DISTRIBUTED_INITIALIZE_TIMEOUT=1800 && export PYTHONPATH=/deps:/deps/src:/deps/src/maxtext/src && python3 src/maxtext/trainers/pre_train/train.py src/maxtext/configs/base.yml ${MAXTEXT_ARGS}" + +# --- XPK Workload Creation --- +echo "Creating XPK workload: ${WORKLOAD_NAME} on cluster: ${CLUSTER_NAME}" + +python3 -m xpk.main workload create --cluster="${CLUSTER_NAME}" --project="${PROJECT_ID}" --zone="${ZONE}" --priority="${PRIORITY}" --max-restarts="${MAX_RESTARTS}" --device-type="${DEVICE_TYPE}" --num-slices="${NUM_SLICES}" --docker-image="${WORKLOAD_IMAGE}" --enable-debug-logs --workload="${WORKLOAD_NAME}" --command="${RUN_COMMAND}" + +LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" +GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" +TB_URL="https://tensorboard.corp.google.com/?logdir=${BASE_OUTPUT_DIR}/param-3/tensorboard" + +echo "========================================================================" +echo "πŸ“‹ Pantheon Cloud Logging (Worker 0 Logs):" +echo "${LOGS_URL}" +echo "" +echo "☸️ GKE Workload Details:" +echo "${GKE_URL}" +echo "" +echo "πŸ“Š GCS TensorBoard Link:" +echo "${TB_URL}" +echo "========================================================================" diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index d6d3905178..e890c90db2 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -387,12 +387,7 @@ def _fwd_run_tokamax_v2( lhs_operand = lhs.qvalue if isinstance(lhs, qpl.QArray) else lhs maybe_quantize_lhs = not isinstance(lhs, qpl.QArray) and qwix_numerics.should_quantize(lhs_operand.dtype) - # When lhs is not pre-quantized (ahead-of-time by the caller), fall back to - # a static scale from the quantization rule's fixed calibration if one is - # configured; the kernel uses it directly instead of computing a dynamic - # per-block absmax. If lhs is already a QArray, its scale is applied - # externally below, so no in-kernel scale is needed. - lhs_scale = None if isinstance(lhs, qpl.QArray) else _fwd_prepare_lhs_scale(quantization_rule) + lhs_scale = _fwd_prepare_lhs_scale(quantization_rule) if maybe_quantize_lhs else None custom_fwd_tiling = gmm_v2.TileSizes( tile_m=tiling[0],