From 5fd507a3b1e19e53756ad39d1962af5812e3303f Mon Sep 17 00:00:00 2001 From: Suraj Kolla Date: Mon, 10 Aug 2026 16:11:59 +0000 Subject: [PATCH] add static LHS scaling support for GMM v2 --- src/maxtext/kernels/megablox/ops.py | 44 +++- .../pallas_mosaic_tpu_v2_gmm_kernel.py | 190 ++++++++++++++---- src/maxtext/layers/quantizations.py | 58 +++--- tests/unit/quantizations_test.py | 58 ++++++ 4 files changed, 284 insertions(+), 66 deletions(-) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index 979d219475..01b132f508 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -200,7 +200,15 @@ 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( @@ -319,6 +327,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, @@ -328,6 +366,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 @@ -354,6 +393,7 @@ def _fwd_run_tokamax_v2( preferred_element_type=preferred_element_type, partial_sum=partial_sum, group_offset=group_offset, + lhs_scale=_fwd_prepare_lhs_scale(quantization_rule), ) @@ -498,7 +538,7 @@ def _bwd_prepare_inputs( dlhs_dout = grad drhs_dout = grad - # Apply rhs.scale to dlhs_dout, dlhs_dout[m, n] @ rhs_tranpose[g, n, k] = dlhs[m, k] + # Apply rhs.scale to dlhs_dout, dlhs_dout[m, n] @ rhs_transpose[g, n, k] = dlhs[m, k] # Assume channelwise scale on rhs n. # Apply rhs.scale to dlhs_dout to avoid dequantizing or requantizing rhs. # We cannot apply the scale to dlhs because axis n will disappear there. 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 46fcfbeafb..4e5e33abde 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, @@ -466,15 +525,21 @@ 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. - block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True) - block_scale = block_abs_max / dtype_max - - # If block_scale=0, it will cause division by zero and return either - # NaN or Inf. Since this can cause numeric issue when downcasting to - # quantized value, we convert them into 0. - block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale) - # Convert lhs into quantized dtype. - block_lhs_q = (block_lhs * block_scale_inv).astype(lhs_q_dtype) + if should_use_external_scale: + 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 + + # If block_scale=0, it will cause division by zero and return either + # NaN or Inf. Since this can cause numeric issue when downcasting to + # quantized value, we convert them into 0. + block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale) + # Convert lhs into quantized dtype. + block_lhs_q = (block_lhs * block_scale_inv).astype(lhs_q_dtype) # Unlike unquantized path, compiler may not perform implicit type # conversion due to numeric concerns. As this can cause unsupported @@ -766,7 +831,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 @@ -855,8 +920,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]) @@ -987,6 +1054,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.""" @@ -1005,9 +1074,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) @@ -1090,10 +1168,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 @@ -1130,6 +1220,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 @@ -1138,6 +1236,7 @@ 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: @@ -1201,6 +1300,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, @@ -1225,6 +1325,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. @@ -1264,11 +1370,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) @@ -1320,33 +1435,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} @@ -1356,7 +1453,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] ), @@ -1368,4 +1476,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] diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index c2c3aba722..55b9559825 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -28,6 +28,7 @@ from aqt.jax.v2 import calibration import qwix +from qwix._src.core import numerics from qwix._src.core import dot_general_qt from qwix._src.core import sparsity @@ -905,11 +906,37 @@ def _make_scale_tensor(scale, arr): return _cast_reduced_from(scale_tensor, arr) -def _get_max_min(target_dtype): - if target_dtype in (jnp.int4, jnp.int8): - return jnp.iinfo(target_dtype).max, jnp.iinfo(target_dtype).min - else: - return jnp.finfo(target_dtype).max.astype(jnp.bfloat16), jnp.finfo(target_dtype).min.astype(jnp.bfloat16) +def get_static_scale(qtype: jax.typing.DTypeLike, calibration_method: str) -> float: + """Extracts the static scale. + Currently, only symmetric fixed range calibration is supported. + For symmetric calibration, the calibration_method must be in the format 'fixed,-max,max' or 'fixed,max'. + + Args: + qtype: The dtype to quantize to. + calibration_method: A string specifying the calibration method. + + Returns: + The extracted static scale value. + """ + if calibration_method is None or not calibration_method.lower().startswith("fixed"): + raise ValueError(f"Only static scale quantization is supported, got {calibration_method}") + + args = [float(a) for a in calibration_method.split(",")[1:]] + if len(args) == 1: + args = [-args[0], args[0]] + + if len(args) != 2 or args[0] + args[1] != 0 or args[1] <= 0: + raise ValueError(f"Expected format: 'fixed,max' or 'fixed,-max,max'. Got: {calibration_method}") + + qmax = numerics.get_symmetric_bound(qtype) + scale_val = args[1] / qmax + + # Prevent scale from being 0 + tiny_sqrt = jnp.finfo(jnp.float32).tiny ** 0.5 + if scale_val < tiny_sqrt: + scale_val = 1.0 + + return scale_val def manual_quantize(tensor: jax.Array, dtype: jax.typing.DTypeLike, calibration_method: str) -> qwix.QArray: @@ -927,24 +954,9 @@ def manual_quantize(tensor: jax.Array, dtype: jax.typing.DTypeLike, calibration_ Raises: ValueError: If calibration_method is None or has an unexpected format. """ - # validate calibration method and parse - calib_method = calibration_method - if calib_method is None: - raise ValueError("calibration_method cannot be None for manual quantization") - if not calib_method.startswith("fixed"): - # we can use static scale for weight/activation, but grad usually needs dynamic - raise ValueError(f"Only static scale quantization is supported, but got {calib_method}") - parts = calib_method.split(",") - if len(parts) != 3: - raise ValueError(f"Unexpected format for calibration method: {calib_method}") - - dtype_max, dtype_min = _get_max_min(dtype) - min_val, max_val = float(parts[1]), float(parts[2]) - if max_val <= 0 or min_val != -max_val: - raise ValueError(f"Unexpected format for calibration method: {calib_method}") - - scale = max_val / dtype_max - scale = jnp.where(scale == 0, 1.0, scale) + scale = get_static_scale(dtype, calibration_method) + dtype_max = numerics.get_symmetric_bound(dtype) + dtype_min = -dtype_max # scale must be converted to a tensor because grad has reduced axes. scale_tensor = _make_scale_tensor(scale, tensor) min_bound = _make_scale_tensor(dtype_min, tensor) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index c11f5ea3b7..0f9bd681b3 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -30,6 +30,7 @@ from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR +from maxtext.kernels.megablox import ops from maxtext.kernels.megablox import gmm from maxtext.layers import nnx_wrappers, quantizations from maxtext.utils import maxtext_utils @@ -37,6 +38,7 @@ from tests.utils.test_helpers import get_test_config_path import numpy as np import pytest +import qwix _QUERY_REGEX = ".*/query" _VALUE_REGEX = ".*/value" @@ -716,5 +718,61 @@ def test_nnx_abstract_state_has_no_intermediates(self): self.assertNotIn("intermediates", state_dict) +class StaticScaleTest(unittest.TestCase): + """Tests for static scale extraction.""" + + def test_get_static_scale_fixed_symmetric(self): + scale = quantizations.get_static_scale(jnp.float8_e4m3fn, "fixed,224.0") + self.assertIsNotNone(scale) + np.testing.assert_allclose(scale, 0.5, rtol=1e-5) + + def test_get_static_scale_invalid_format(self): + with self.assertRaises(ValueError): + quantizations.get_static_scale(jnp.float8_e4m3fn, "fixed,1,2,3") + + with self.assertRaises(ValueError): + quantizations.get_static_scale(jnp.float8_e4m3fn, "fixed,-200.0,224.0") + + with self.assertRaises(ValueError): + quantizations.get_static_scale(jnp.float8_e4m3fn, "fixed") + + with self.assertRaises(ValueError): + quantizations.get_static_scale(jnp.float8_e4m3fn, "absmax") + + +class LhsScaleTest(unittest.TestCase): + """Tests for LHS scale extraction in GMM v2 forward.""" + + def test_fwd_prepare_lhs_scale_fixed_symmetric(self): + + rule = qwix.QtRule( + act_qtype=jnp.float8_e4m3fn, + act_calibration_method="fixed,224.0", + ) + scale = ops._fwd_prepare_lhs_scale(rule) # pylint: disable=protected-access + self.assertIsNotNone(scale) + self.assertEqual(scale.shape, (1, 1)) + self.assertEqual(scale.dtype, jnp.float32) + np.testing.assert_allclose(scale, 0.5, rtol=1e-5) + + def test_fwd_prepare_lhs_scale_dynamic_returns_none(self): + + rule = qwix.QtRule( + act_qtype=jnp.float8_e4m3fn, + act_calibration_method="absmax", + ) + scale = ops._fwd_prepare_lhs_scale(rule) # pylint: disable=protected-access + self.assertIsNone(scale) + + def test_fwd_prepare_lhs_scale_no_rule_returns_none(self): + + rule = qwix.QtRule( + act_qtype=None, + act_calibration_method=None, + ) + scale = ops._fwd_prepare_lhs_scale(rule) # pylint: disable=protected-access + self.assertIsNone(scale) + + if __name__ == "__main__": unittest.main()