diff --git a/qwix/_src/core/dot_general_qt.py b/qwix/_src/core/dot_general_qt.py index bf7a2b71..61f3aa8a 100644 --- a/qwix/_src/core/dot_general_qt.py +++ b/qwix/_src/core/dot_general_qt.py @@ -22,6 +22,7 @@ import numpy as np from qwix._src import interception from qwix._src.core import dot_general +from qwix._src.core import multipass from qwix._src.core import numerics from qwix._src.core import qarray from qwix._src.core import sparsity @@ -43,6 +44,7 @@ class DotGeneralQtConfig: rhs_collect_quant_stat: Callable[[Any], Any] | None = None lhs_disable_channelwise_axes: bool = False rhs_disable_channelwise_axes: bool = False + multipass_mode: multipass.MultiPassMode | None = None # Backward pass (dlhs). dlhs_grad_qtype: jax.typing.DTypeLike | None = None # incoming gradient @@ -50,6 +52,7 @@ class DotGeneralQtConfig: dlhs_tile_size: int | float | None = None dlhs_stochastic_rounding_noise_fn: stochastic_rounding.NoiseFn | None = None dlhs_grad_disable_channelwise_axes: bool = False + dlhs_multipass_mode: multipass.MultiPassMode | None = None # Backward pass (drhs). drhs_grad_qtype: jax.typing.DTypeLike | None = None # incoming gradient @@ -57,6 +60,7 @@ class DotGeneralQtConfig: drhs_tile_size: int | float | None = None drhs_stochastic_rounding_noise_fn: stochastic_rounding.NoiseFn | None = None drhs_grad_disable_channelwise_axes: bool = False + drhs_multipass_mode: multipass.MultiPassMode | None = None # Whether not to clip the gradients to the calibration ranges of the quantized # inputs. Enabling this improves the performance but may decrease the @@ -168,6 +172,58 @@ def _apply_rhs_scale_to_lhs(lhs, rhs_scale, dnums): return qarray.call_with_generic_broadcast(jnp.multiply, lhs, lhs_scale) +_MICROSCALED_QTYPES = frozenset({ + 'mxfp8', + 'mxfp8_16', + 'mxfp4', + 'nvfp4', + 'mxint8', + 'mxint4', +}) + + +def _is_block_scaled(operand: qarray.MaybeQArray) -> bool: + """Returns True if the operand uses block-scaled or tiled quantization. + + Args: + operand: The potentially quantized operand. + + Returns: + True if the operand uses block-scaled or tiled quantization. + """ + if not isinstance(operand, qarray.QArray): + return False + return ( + bool(qarray.get_tiled_axes(operand)) + or operand.qtype in _MICROSCALED_QTYPES + ) + + +def _requires_unquantized_residual( + config: DotGeneralQtConfig, operand_qt: qarray.MaybeQArray +) -> bool: + """Returns True if the backward contraction cannot reuse the forward operand. + + Block-scaled residuals cannot be reused because quantization scales defined + for the forward contraction axis do not align with the new contraction axis + in the backward pass. Multi-pass modes also require original residuals. + + Args: + config: The quantization configuration. + operand_qt: The potentially quantized operand. + + Returns: + True if the backward contraction cannot reuse the forward operand. + """ + return ( + config.use_original_residuals + or config.multipass_mode is not None + or config.dlhs_multipass_mode is not None + or config.drhs_multipass_mode is not None + or _is_block_scaled(operand_qt) + ) + + def _get_residual_for_backward( config: DotGeneralQtConfig, operand_in: jax.Array | None, @@ -184,15 +240,11 @@ def _get_residual_for_backward( config: The quantization configuration. operand_in: The original, unquantized operand, or None if not retained. operand_qt: The potentially quantized operand. + + Returns: + The residual to be used in the backward pass. """ - if config.use_original_residuals or ( - isinstance(operand_qt, qarray.QArray) - and ( - qarray.get_tiled_axes(operand_qt) - or isinstance(operand_qt.qtype, str) - and operand_qt.qtype in ('mxfp8', 'mxfp8_16', 'mxfp4', 'nvfp4') - ) - ): + if _requires_unquantized_residual(config, operand_qt): assert operand_in is not None return operand_in return operand_qt @@ -204,7 +256,17 @@ def _needs_original_residual( calibration: dict[str, jax.Array] | None, calibration_method: str, ) -> bool: - """Returns True if the unquantized operand must be retained in residuals.""" + """Returns True if the unquantized operand must be retained in residuals. + + Args: + config: The quantization configuration. + operand_qt: The potentially quantized operand. + calibration: The calibration dictionary or None. + calibration_method: The calibration method string. + + Returns: + True if the unquantized operand must be retained in residuals. + """ # 1. Contraction residual: # The unquantized operand must be retained if: # a) The user explicitly opted into original residuals @@ -213,16 +275,7 @@ def _needs_original_residual( # tiled scales). # c) The operand is an MXFP/microscaling type (block-level scales also # require original operands). - if config.use_original_residuals or ( - isinstance(operand_qt, qarray.QArray) - and ( - qarray.get_tiled_axes(operand_qt) - or ( - isinstance(operand_qt.qtype, str) - and operand_qt.qtype in ('mxfp8', 'mxfp8_16', 'mxfp4', 'nvfp4') - ) - ) - ): + if _requires_unquantized_residual(config, operand_qt): return True # 2. Gradient clipping: @@ -248,6 +301,46 @@ def dot_general_qt_fwd( ): """Forward pass for dot_general_qt custom VJP.""" lhs_in, rhs_in = lhs, rhs + if config.multipass_mode is not None: + lhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dimension_numbers, + ndims=(lhs.ndim, rhs.ndim), + for_lhs=True, + qtype=config.lhs_qtype, + tile_size=config.tile_size, + calibration_method=config.lhs_calibration_method, + ) + if config.lhs_disable_channelwise_axes: + lhs_how = dataclasses.replace(lhs_how, channelwise_axes=[]) + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dimension_numbers, + ndims=(lhs.ndim, rhs.ndim), + for_lhs=False, + qtype=config.rhs_qtype, + tile_size=config.tile_size, + calibration_method=config.rhs_calibration_method, + ) + if config.rhs_disable_channelwise_axes: + rhs_how = dataclasses.replace(rhs_how, channelwise_axes=[]) + + out = multipass.multipass_dot( + lhs, + rhs, + dimension_numbers=dimension_numbers, + mode=config.multipass_mode, + lhs_how=lhs_how, + rhs_how=rhs_how, + ) + residuals = ( + lhs_in, + rhs_in, + lhs_in, + rhs_in, + None, + None, + config, + ) + return out, residuals if lhs_calibration is not None: scale, zero_point = qarray.compute_scale_zero_point( lhs_calibration, config.lhs_qtype # pyrefly: ignore[bad-argument-type] @@ -335,6 +428,53 @@ def _compute_gradient_for_operand(g: jax.Array, *, for_dlhs: bool): y_calibration_method = config.drhs_residual_calibration_method y_disable_channelwise_axes = config.drhs_residual_disable_channelwise_axes + multipass_mode = ( + config.dlhs_multipass_mode if for_dlhs else config.drhs_multipass_mode + ) + if multipass_mode is not None: + if isinstance(y, qarray.QArray): + y = qarray.dequantize(y) + resolved_g_qtype = ( + g_qtype + or (config.lhs_qtype if for_dlhs else config.rhs_qtype) + or 'mxfp8_16' + ) + resolved_y_qtype = ( + y_qtype + or (config.rhs_qtype if for_dlhs else config.lhs_qtype) + or 'mxfp8_16' + ) + resolved_tile_size = g_tile_size or config.tile_size + g_how = dot_general.get_how_to_quantize( + dimension_numbers=bwd_dnums, + ndims=(g.ndim, y.ndim), + for_lhs=True, + qtype=resolved_g_qtype, + tile_size=resolved_tile_size, + calibration_method=g_calibration_method, + ) + if g_disable_channelwise_axes: + g_how = dataclasses.replace(g_how, channelwise_axes=[]) + y_how = dot_general.get_how_to_quantize( + dimension_numbers=bwd_dnums, + ndims=(g.ndim, y.ndim), + for_lhs=False, + qtype=resolved_y_qtype, + tile_size=resolved_tile_size, + calibration_method=y_calibration_method, + ) + if y_disable_channelwise_axes: + y_how = dataclasses.replace(y_how, channelwise_axes=[]) + grad_res = multipass.multipass_dot( + g, + y, + dimension_numbers=bwd_dnums, + mode=multipass_mode, + lhs_how=g_how, + rhs_how=y_how, + ) + return jax.lax.transpose(grad_res, transpose_axes) + if g_qtype and numerics.should_quantize(g.dtype): if isinstance(y, qarray.QArray): # Scale shifting for quantized residuals (use_original_residuals=False) @@ -420,6 +560,11 @@ def dot_general_qt( config: DotGeneralQtConfig, ) -> jax.Array: """Quantized dot_general with backpropagation support.""" + if config.multipass_mode is not None: + return dot_general_qt_fwd_bwd( + lhs, rhs, None, None, dimension_numbers, config + ) + lhs_calibration = None rhs_calibration = None diff --git a/qwix/_src/core/multipass.py b/qwix/_src/core/multipass.py new file mode 100644 index 00000000..597dbca9 --- /dev/null +++ b/qwix/_src/core/multipass.py @@ -0,0 +1,260 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-pass emulation for microscaled floating-point and integer formats. + +This module provides multi-pass residual quantization emulation: +- 2-pass residual decomposition: + Pass 0: X_0 = quantize(X), R_1 = X - dequantize(X_0) + Pass 1: X_1 = quantize(R_1) +- Multi-pass matrix multiplication modes: + 1. 'triangular' (3 passes): A_0 B_0 + A_0 B_1 + A_1 B_0 + Drops second-order cross-residual A_1 B_1 (scaled by ~2^-14). + 2. 'full_cross' (4 passes): A_0 B_0 + A_0 B_1 + A_1 B_0 + A_1 B_1 + Evaluates all 4 cross-products. + 3. 'lhs_high_precision' (2 passes): A_0 B_0 + A_1 B_0 + LHS uses 2 residual passes, RHS uses 1 pass. + 4. 'rhs_high_precision' (2 passes): A_0 B_0 + A_0 B_1 + LHS uses 1 pass, RHS uses 2 residual passes. + +Scaling factor design choice for residual passes: +While an alternative formulation could independently compute fresh scaling +factors for each residual pass (via block-level max-reduction searches over +R_1), + +we instead derive residual scales by shifting the initial pass's scale factor by +a fixed power-of-2 exponent offset (e.g. 2^-4 for FP8 E4M3, 2^-2 for FP4 E2M1). +This design choice is made because: + +1. Hardware efficiency: It eliminates the expensive second block-level reduction + tree across elements, replacing dynamic exponent search with a single ALU + integer shift (E_1 = E_0 - shift_bits). +2. Memory bandwidth and storage: The residual scale is implicit rather than a + separate scale tensor that must be stored and loaded from memory, cutting + scale metadata traffic by 50%. +3. Fixed accumulator alignment: Constant power-of-2 scale offsets allow matrix + units to align cross-term products with simple arithmetic bit-shifts prior + to register accumulation, rather than performing variable per-block floating- + point scale multiplications. +4. Numerical closeness: Because the maximum rounding residual of round-to- + nearest is bounded by the top bin step size, shifting by these exact bit + yields virtually identical SQNR to independent scale choices. +""" + +from typing import Literal, TypeAlias +import jax +import jax.numpy as jnp +from qwix._src.core import dot_general +from qwix._src.core import qarray + +MultiPassMode: TypeAlias = Literal[ + 'triangular', + 'full_cross', + 'lhs_high_precision', + 'rhs_high_precision', +] + + +def get_residual_scale_shift_bits(qtype: jax.typing.DTypeLike) -> int | None: + """Returns the power-of-2 scale shift bits for residual quantization. + + For microscaled floating-point formats, the residual quantization scale + can be derived directly from the preceding pass's scale by shifting its + exponent, eliminating redundant block-level reduction passes: + - mxfp8 / mxfp8_16 (E4M3): 16 / 448 = 1/28 -> shift by 4 bits (2^-4 = 1/16). + - mxfp4 / nvfp4 (E2M1): 1 / 6 = 1/6 -> shift by 2 bits (2^-2 = 1/4). + - float8_e5m2 (E5M2): 8192 / 57344 = 1/7 -> shift by 2 bits (2^-2 = 1/4). + + For integer formats (e.g. int4, int8, mxint8), returns None to indicate + exact algebraic decomposition or independent calibration. + + Args: + qtype: Quantization format or dtype string. + + Returns: + The integer exponent shift in bits, or None. + """ + match qtype: + case 'mxfp8' | 'mxfp8_16': + return 4 + case 'mxfp4' | 'nvfp4': + return 2 + case _: + if qtype == jnp.float8_e4m3fn or qtype == 'float8_e4m3fn': + return 4 + elif qtype == jnp.float4_e2m1fn or qtype == 'float4_e2m1fn': + return 2 + elif qtype == jnp.float8_e5m2 or qtype == 'float8_e5m2': + return 2 + return None + + +def residual_decompose( + x: jax.Array, + how: qarray.HowToQuantize, + n_passes: int = 2, +) -> tuple[qarray.QArray, ...]: + """Decomposes tensor x into residual quantized passes. + + For floating-point formats (e.g., mxfp8, mxfp4), residual passes reuse the + initial pass's scale factor shifted by the format's bit-precision (e.g. 2^-4 + for FP8, 2^-2 for FP4), eliminating redundant block-level reduction passes. + For integer formats, passes are calibrated independently. + + Args: + x: Input array. + how: HowToQuantize configuration specifying format, tile size, scaling, etc. + n_passes: Number of residual passes (default 2). + + Returns: + Tuple of quantized QArray objects (q_0, q_1, ...). + """ + passes = [] + current = x + shift_bits = get_residual_scale_shift_bits(how.qtype) + + # Pass 0 + q0 = qarray.quantize(current, how) + passes.append(q0) + deq0 = qarray.dequantize(q0) + current = current - deq0 + + base_scale = q0.scale + for p in range(1, n_passes): + if shift_bits is not None: + scale_p = (base_scale * (2.0 ** (-shift_bits * p))).astype( + base_scale.dtype + ) + q = qarray.quantize_with_scale_zero_point( + current, + how.qtype, + scale=scale_p, + zero_point=q0.zero_point, + noise_fn=how.noise_fn, + ) + else: + q = qarray.quantize(current, how) + passes.append(q) + deq = qarray.dequantize(q) + current = current - deq + + return tuple(passes) + + +def multipass_dot( + lhs: jax.Array, + rhs: jax.Array, + dimension_numbers: jax.lax.DotDimensionNumbers = (((1,), (0,)), ((), ())), + mode: MultiPassMode = 'triangular', + *, + lhs_how: qarray.HowToQuantize | None = None, + rhs_how: qarray.HowToQuantize | None = None, + lhs_qtype: str = 'mxfp8_16', + rhs_qtype: str = 'mxfp8_16', + tile_size: int = 16, + preferred_element_type: jax.typing.DTypeLike | None = None, +) -> jax.Array: + """Executes multi-pass emulated matrix multiplication. + + Args: + lhs: Left-hand side tensor. + rhs: Right-hand side tensor. + dimension_numbers: Standard JAX dot_general dimension specification. + mode: Multi-pass mode: - 'triangular': Computes 3 GEMMs: A_0 B_0 + A_0 B_1 + + A_1 B_0. - 'full_cross': Computes 4 GEMMs: A_0 B_0 + A_0 B_1 + A_1 B_0 + + A_1 B_1. - 'lhs_high_precision': Computes 2 GEMMs: A_0 B_0 + A_1 B_0 (2 + passes on LHS, 1 on RHS). - 'rhs_high_precision': Computes 2 GEMMs: A_0 + B_0 + A_0 B_1 (1 pass on LHS, 2 on RHS). + lhs_how: Explicit HowToQuantize for LHS (overrides lhs_qtype, tile_size, + etc.). + rhs_how: Explicit HowToQuantize for RHS (overrides rhs_qtype, tile_size, + etc.). + lhs_qtype: QType for LHS if lhs_how is None (default 'mxfp8_16'). + rhs_qtype: QType for RHS if rhs_how is None (default 'mxfp8_16'). + tile_size: Microscaling tile size (default 16 for mxfp8_16). + preferred_element_type: Output/accumulator dtype. + + Returns: + The resulting matrix product array. + """ + if lhs_how is None: + lhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dimension_numbers, + ndims=(lhs.ndim, rhs.ndim), + for_lhs=True, + qtype=lhs_qtype, + tile_size=tile_size, + ) + if rhs_how is None: + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dimension_numbers, + ndims=(lhs.ndim, rhs.ndim), + for_lhs=False, + qtype=rhs_qtype, + tile_size=tile_size, + ) + + if mode == 'lhs_high_precision': + a_passes = residual_decompose(lhs, lhs_how, n_passes=2) + a0, a1 = a_passes[0], a_passes[1] + b0 = qarray.quantize(rhs, rhs_how) + c0 = dot_general.dot_general( + a0, b0, dimension_numbers, preferred_element_type=preferred_element_type + ) + c1 = dot_general.dot_general( + a1, b0, dimension_numbers, preferred_element_type=preferred_element_type + ) + return c0 + c1 + + elif mode == 'rhs_high_precision': + a0 = qarray.quantize(lhs, lhs_how) + b_passes = residual_decompose(rhs, rhs_how, n_passes=2) + b0, b1 = b_passes[0], b_passes[1] + c0 = dot_general.dot_general( + a0, b0, dimension_numbers, preferred_element_type=preferred_element_type + ) + c1 = dot_general.dot_general( + a0, b1, dimension_numbers, preferred_element_type=preferred_element_type + ) + return c0 + c1 + + elif mode in ('triangular', 'full_cross'): + a_passes = residual_decompose(lhs, lhs_how, n_passes=2) + a0, a1 = a_passes[0], a_passes[1] + b_passes = residual_decompose(rhs, rhs_how, n_passes=2) + b0, b1 = b_passes[0], b_passes[1] + c00 = dot_general.dot_general( + a0, b0, dimension_numbers, preferred_element_type=preferred_element_type + ) + c01 = dot_general.dot_general( + a0, b1, dimension_numbers, preferred_element_type=preferred_element_type + ) + c10 = dot_general.dot_general( + a1, b0, dimension_numbers, preferred_element_type=preferred_element_type + ) + res = c00 + c01 + c10 + if mode == 'full_cross': + c11 = dot_general.dot_general( + a1, + b1, + dimension_numbers, + preferred_element_type=preferred_element_type, + ) + res = res + c11 + return res + + else: + raise ValueError( + f"Unknown multipass mode: {mode!r}. Expected one of: 'triangular'," + " 'full_cross', 'lhs_high_precision', 'rhs_high_precision'." + ) diff --git a/qwix/_src/core/numerics.py b/qwix/_src/core/numerics.py index 9a087700..a1c9cc89 100644 --- a/qwix/_src/core/numerics.py +++ b/qwix/_src/core/numerics.py @@ -117,6 +117,10 @@ def get_symmetric_bound(qtype: jax.typing.DTypeLike) -> float: qtype = jnp.float8_e4m3fn case 'mxfp4' | 'nvfp4': qtype = jnp.float4_e2m1fn + case 'mxint8': + qtype = jnp.int8 + case 'mxint4': + qtype = jnp.int4 # Prevent common misconfigurations, e.g., use bf16 as qtype. if jnp.dtype(qtype).itemsize > 1 and jnp.dtype(qtype) != jnp.int16: @@ -174,6 +178,10 @@ def convert_to( qtype = jnp.float8_e4m3fn case 'mxfp4' | 'nvfp4': qtype = jnp.float4_e2m1fn + case 'mxint8': + qtype = jnp.int8 + case 'mxint4': + qtype = jnp.int4 # Handles builtin qtypes. try: diff --git a/qwix/_src/core/qarray.py b/qwix/_src/core/qarray.py index b7a4b2fe..2ae69e91 100644 --- a/qwix/_src/core/qarray.py +++ b/qwix/_src/core/qarray.py @@ -309,8 +309,12 @@ def __post_init__(self): 'mxfp8_16', 'mxfp4', 'nvfp4', + 'mxint8', + 'mxint4', ): - resolved_tile_size = 32 if self.qtype in ('mxfp8', 'mxfp4') else 16 + resolved_tile_size = ( + 32 if self.qtype in ('mxfp8', 'mxfp4', 'mxint8', 'mxint4') else 16 + ) if not self.tiled_axes: raise ValueError( @@ -581,6 +585,21 @@ def compute_scale_zero_point( .view(jnp.bfloat16) .astype(scale.dtype) ) + elif qtype in ('mxint8', 'mxint4'): + # Efficient bit manipulation for 2 ** ceil(log2(scale)) without + # transcendentals: + # In IEEE-754 float32, adding the mantissa mask (0x007FFFFF) carries into + # the exponent if and only if the mantissa > 0 (scale is not already an + # exact power of 2). Masking with 0x7F800000 clears the mantissa, yielding + # the exact ceil power-of-2. + scale_f32 = scale.astype(jnp.float32) + scale_bits = scale_f32.view(jnp.int32) + scale_pow2 = ( + ((scale_bits + 0x007FFFFF) & 0x7F800000) + .view(jnp.float32) + .astype(scale.dtype) + ) + scale = jnp.where(scale > 0, scale_pow2, scale) elif qtype == 'nvfp4': scale = numerics.convert_to(scale, jnp.float8_e4m3fn).astype(scale.dtype) return scale, zero_point diff --git a/tests/_src/core/dot_general_qt_test.py b/tests/_src/core/dot_general_qt_test.py index 0f8fa1ab..db01e928 100644 --- a/tests/_src/core/dot_general_qt_test.py +++ b/tests/_src/core/dot_general_qt_test.py @@ -673,6 +673,146 @@ def quant_loss(l, r): ' 25 dB (indicates multi-axis contraction corruption)', ) + def test_mxint8_sqnr(self): + """Verifies that mxint8 dot_general_qt achieves expected SQNR.""" + key = jax.random.PRNGKey(42) + k1, k2, k3 = jax.random.split(key, 3) + + lhs = jax.random.normal(k1, (2, 32, 64), dtype=jnp.float32) + rhs = jax.random.normal(k2, (64, 128), dtype=jnp.float32) + dout = jax.random.normal(k3, (2, 32, 128), dtype=jnp.float32) + + dnums = (((2,), (0,)), ((), ())) + ref_fwd = jax.lax.dot_general(lhs, rhs, dnums) + + def ref_loss(l, r): + return jnp.sum(jax.lax.dot_general(l, r, dnums) * dout) + + ref_grad_lhs, ref_grad_rhs = jax.grad(ref_loss, argnums=(0, 1))(lhs, rhs) + + config = dot_general_qt.DotGeneralQtConfig( + lhs_qtype='mxint8', + rhs_qtype='mxint8', + dlhs_grad_qtype='mxint8', + drhs_grad_qtype='mxint8', + dlhs_residual_qtype='mxint8', + drhs_residual_qtype='mxint8', + tile_size=32, + dlhs_tile_size=32, + drhs_tile_size=32, + use_original_residuals=False, + ) + test_fwd = dot_general_qt.dot_general_qt(lhs, rhs, dnums, config=config) + fwd_sqnr = float( + metrax.SNR.from_model_output( + predictions=test_fwd, targets=ref_fwd + ).compute() + ) + self.assertGreater( + fwd_sqnr, + 30.0, + f'mxint8 forward pass SQNR {fwd_sqnr:.2f} dB is below 30 dB', + ) + + def quant_loss(l, r): + return jnp.sum( + dot_general_qt.dot_general_qt(l, r, dnums, config=config) * dout + ) + + grad_lhs, grad_rhs = jax.grad(quant_loss, argnums=(0, 1))(lhs, rhs) + + dlhs_sqnr = float( + metrax.SNR.from_model_output( + predictions=grad_lhs, targets=ref_grad_lhs + ).compute() + ) + drhs_sqnr = float( + metrax.SNR.from_model_output( + predictions=grad_rhs, targets=ref_grad_rhs + ).compute() + ) + + self.assertGreater( + dlhs_sqnr, + 30.0, + f'mxint8 dlhs gradient SQNR {dlhs_sqnr:.2f} dB is below 30 dB', + ) + self.assertGreater( + drhs_sqnr, + 30.0, + f'mxint8 drhs weight gradient SQNR {drhs_sqnr:.2f} dB is below 30 dB', + ) + + def test_mxint4_sqnr(self): + """Verifies that mxint4 dot_general_qt achieves expected SQNR.""" + key = jax.random.PRNGKey(42) + k1, k2, k3 = jax.random.split(key, 3) + + lhs = jax.random.normal(k1, (2, 32, 64), dtype=jnp.float32) + rhs = jax.random.normal(k2, (64, 128), dtype=jnp.float32) + dout = jax.random.normal(k3, (2, 32, 128), dtype=jnp.float32) + + dnums = (((2,), (0,)), ((), ())) + ref_fwd = jax.lax.dot_general(lhs, rhs, dnums) + + def ref_loss(l, r): + return jnp.sum(jax.lax.dot_general(l, r, dnums) * dout) + + ref_grad_lhs, ref_grad_rhs = jax.grad(ref_loss, argnums=(0, 1))(lhs, rhs) + + config = dot_general_qt.DotGeneralQtConfig( + lhs_qtype='mxint4', + rhs_qtype='mxint4', + dlhs_grad_qtype='mxint4', + drhs_grad_qtype='mxint4', + dlhs_residual_qtype='mxint4', + drhs_residual_qtype='mxint4', + tile_size=32, + dlhs_tile_size=32, + drhs_tile_size=32, + use_original_residuals=False, + ) + test_fwd = dot_general_qt.dot_general_qt(lhs, rhs, dnums, config=config) + fwd_sqnr = float( + metrax.SNR.from_model_output( + predictions=test_fwd, targets=ref_fwd + ).compute() + ) + self.assertGreater( + fwd_sqnr, + 12.0, + f'mxint4 forward pass SQNR {fwd_sqnr:.2f} dB is below 12 dB', + ) + + def quant_loss(l, r): + return jnp.sum( + dot_general_qt.dot_general_qt(l, r, dnums, config=config) * dout + ) + + grad_lhs, grad_rhs = jax.grad(quant_loss, argnums=(0, 1))(lhs, rhs) + + dlhs_sqnr = float( + metrax.SNR.from_model_output( + predictions=grad_lhs, targets=ref_grad_lhs + ).compute() + ) + drhs_sqnr = float( + metrax.SNR.from_model_output( + predictions=grad_rhs, targets=ref_grad_rhs + ).compute() + ) + + self.assertGreater( + dlhs_sqnr, + 12.0, + f'mxint4 dlhs gradient SQNR {dlhs_sqnr:.2f} dB is below 12 dB', + ) + self.assertGreater( + drhs_sqnr, + 12.0, + f'mxint4 drhs weight gradient SQNR {drhs_sqnr:.2f} dB is below 12 dB', + ) + if __name__ == '__main__': absltest.main() diff --git a/tests/_src/core/multipass_test.py b/tests/_src/core/multipass_test.py new file mode 100644 index 00000000..50cb5fb2 --- /dev/null +++ b/tests/_src/core/multipass_test.py @@ -0,0 +1,762 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import lax +import jax.numpy as jnp +import metrax +import numpy as np +from qwix._src.core import dot_general +from qwix._src.core import dot_general_qt +from qwix._src.core import multipass +from qwix._src.core import qarray + + +def compute_snr_db(true_val: jax.Array, approx_val: jax.Array) -> float: + """Computes Signal-to-Noise Ratio (SNR) in decibels (dB) via metrax.SNR.""" + return float( + metrax.SNR.from_model_output( + predictions=approx_val.astype(jnp.float32), + targets=true_val.astype(jnp.float32), + ).compute() + ) + + +def compute_relative_error(true_val: jax.Array, approx_val: jax.Array) -> float: + """Computes relative L2 error ||true - approx|| / ||true||.""" + true_f = true_val.astype(jnp.float32) + approx_f = approx_val.astype(jnp.float32) + norm_diff = jnp.linalg.norm(true_f - approx_f) + norm_true = jnp.linalg.norm(true_f) + return float(norm_diff / jnp.maximum(norm_true, 1e-12)) + + +class MultiPassTest(parameterized.TestCase): + + def setUp(self): + super().setUp() + self.rng = jax.random.PRNGKey(42) + + def test_residual_decompose(self): + key = self.rng + x = jax.random.normal(key, (4, 32), dtype=jnp.float32) + how = dot_general.get_how_to_quantize( + dimension_numbers=(((1,), (0,)), ((), ())), + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + passes = multipass.residual_decompose(x, how, n_passes=2) + self.assertLen(passes, 2) + q0, q1 = passes + self.assertIsInstance(q0, qarray.QArray) + self.assertIsInstance(q1, qarray.QArray) + + # Dequantized values + x0 = qarray.dequantize(q0) + x1 = qarray.dequantize(q1) + reconstructed = x0 + x1 + + err_single = compute_relative_error(x, x0) + err_double = compute_relative_error(x, reconstructed) + self.assertLess(err_double, err_single) + + snr_single = compute_snr_db(x, x0) + snr_double = compute_snr_db(x, reconstructed) + self.assertGreater(snr_double, snr_single) + + def test_residual_decomposition_n_passes_progression(self): + """Verifies that each residual pass reduces error and increases SNR.""" + key = self.rng + x = jax.random.normal(key, (8, 64), dtype=jnp.float32) + how = dot_general.get_how_to_quantize( + dimension_numbers=(((1,), (0,)), ((), ())), + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + n_passes = 4 + passes = multipass.residual_decompose(x, how, n_passes=n_passes) + self.assertLen(passes, n_passes) + + reconstructed = jnp.zeros_like(x) + prev_err = float('inf') + prev_snr = -float('inf') + prev_scale_mean = float('inf') + + for p in range(n_passes): + deq_p = qarray.dequantize(passes[p]) + reconstructed = reconstructed + deq_p + + err_p = compute_relative_error(x, reconstructed) + snr_p = compute_snr_db(x, reconstructed) + scale_mean = float(jnp.mean(passes[p].scale)) + + # Strict error decrease and SNR increase with each added pass + self.assertLess(err_p, prev_err) + self.assertGreater(snr_p, prev_snr) + # Residual scale must drop substantially with each pass (~16x for FP8) + self.assertLess(scale_mean, prev_scale_mean / 4.0) + + prev_err = err_p + prev_snr = snr_p + prev_scale_mean = scale_mean + + # 4 passes of FP8 should achieve near-lossless reconstruction + # (> 85 dB SNR with float32 eps). + self.assertGreater(prev_snr, 85.0) + + @parameterized.parameters( + ('mxfp8_16', 16), + ) + def test_exact_quantized_input_zero_residual(self, qtype, tile_size): + """Verifies that on-grid inputs produce zero residual in pass 1.""" + k = self.rng + raw = jax.random.normal(k, (8, 64), dtype=jnp.float32) + how = dot_general.get_how_to_quantize( + dimension_numbers=(((1,), (0,)), ((), ())), + ndims=(2, 2), + for_lhs=True, + qtype=qtype, + tile_size=tile_size, + ) + # Project raw values onto the exact quantization grid + exact_vals = qarray.dequantize(qarray.quantize(raw, how)) + + # Decompose the already-quantized values + passes = multipass.residual_decompose(exact_vals, how, n_passes=2) + self.assertLen(passes, 2) + q0, q1 = passes + + # Pass 0 must exactly equal exact_vals + x0 = qarray.dequantize(q0) + np.testing.assert_allclose(x0, exact_vals, rtol=1e-5, atol=1e-5) + + # Pass 1 must have identically zero qvalue and dequantize to 0.0 + np.testing.assert_array_equal(q1.qvalue, 0) + x1 = qarray.dequantize(q1) + np.testing.assert_allclose(x1, 0.0, atol=1e-7) + + @parameterized.parameters( + 'triangular', + 'full_cross', + 'lhs_high_precision', + 'rhs_high_precision', + ) + def test_exact_algebraic_component_equivalence(self, mode): + """Verifies that multipass_dot matches the exact sum of component GEMMs.""" + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, (8, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 8), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + lhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype='mxfp8_16', + tile_size=16, + ) + + actual = multipass.multipass_dot( + lhs, + rhs, + dimension_numbers=dnums, + mode=mode, + lhs_how=lhs_how, + rhs_how=rhs_how, + ) + + # Compute manual reference terms from dequantized component passes + a_passes = multipass.residual_decompose(lhs, lhs_how, n_passes=2) + b_passes = multipass.residual_decompose(rhs, rhs_how, n_passes=2) + a0 = qarray.dequantize(a_passes[0]) + a1 = qarray.dequantize(a_passes[1]) + b0 = qarray.dequantize(b_passes[0]) + b1 = qarray.dequantize(b_passes[1]) + + c00 = lax.dot_general(a0, b0, dnums) + c01 = lax.dot_general(a0, b1, dnums) + c10 = lax.dot_general(a1, b0, dnums) + c11 = lax.dot_general(a1, b1, dnums) + + if mode == 'triangular': + expected = c00 + c01 + c10 + elif mode == 'full_cross': + expected = c00 + c01 + c10 + c11 + elif mode == 'lhs_high_precision': + expected = c00 + c10 + elif mode == 'rhs_high_precision': + expected = c00 + c01 + else: + raise ValueError(f'Unexpected mode: {mode}') + + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + def test_exact_dropped_cross_residual(self): + """Verifies full_cross - triangular is identically equal to A_1 @ B_1.""" + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, (8, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 8), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + lhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype='mxfp8_16', + tile_size=16, + ) + + res_tri = multipass.multipass_dot( + lhs, rhs, dnums, mode='triangular', lhs_how=lhs_how, rhs_how=rhs_how + ) + res_full = multipass.multipass_dot( + lhs, rhs, dnums, mode='full_cross', lhs_how=lhs_how, rhs_how=rhs_how + ) + + a_passes = multipass.residual_decompose(lhs, lhs_how, n_passes=2) + b_passes = multipass.residual_decompose(rhs, rhs_how, n_passes=2) + a1 = qarray.dequantize(a_passes[1]) + b1 = qarray.dequantize(b_passes[1]) + a1_b1 = lax.dot_general(a1, b1, dnums) + + # Difference must match A1 @ B1 to floating point precision + diff = res_full - res_tri + np.testing.assert_allclose(diff, a1_b1, rtol=1e-5, atol=1e-5) + + # Second order term magnitude must be ~10^-4 or smaller relative to full + rel_second_order = jnp.linalg.norm(a1_b1) / jnp.linalg.norm(res_full) + self.assertLess(float(rel_second_order), 1e-3) + + @parameterized.parameters( + 'gaussian', + 'uniform', + 'heavy_tailed', + 'outliers', + ) + def test_multipass_snr_hierarchy_across_distributions(self, dist): + """Verifies strict SNR monotonic improvement across distributions.""" + k1, k2, k3, k4 = jax.random.split(self.rng, 4) + shape_lhs = (16, 64) + shape_rhs = (64, 16) + dnums = (((1,), (0,)), ((), ())) + + if dist == 'gaussian': + lhs = jax.random.normal(k1, shape_lhs, dtype=jnp.float32) + rhs = jax.random.normal(k2, shape_rhs, dtype=jnp.float32) + elif dist == 'uniform': + lhs = jax.random.uniform(k1, shape_lhs, minval=-3.0, maxval=3.0) + rhs = jax.random.uniform(k2, shape_rhs, minval=-3.0, maxval=3.0) + elif dist == 'heavy_tailed': + lhs = jax.random.laplace(k1, shape_lhs, dtype=jnp.float32) + rhs = jax.random.laplace(k2, shape_rhs, dtype=jnp.float32) + elif dist == 'outliers': + lhs = jax.random.normal(k1, shape_lhs, dtype=jnp.float32) + rhs = jax.random.normal(k2, shape_rhs, dtype=jnp.float32) + # Add 2% large outliers (30x scale) + mask_l = jax.random.bernoulli(k3, p=0.02, shape=shape_lhs) + mask_r = jax.random.bernoulli(k4, p=0.02, shape=shape_rhs) + lhs = jnp.where(mask_l, lhs * 30.0, lhs) + rhs = jnp.where(mask_r, rhs * 30.0, rhs) + else: + raise ValueError(f'Unknown distribution: {dist}') + + true_res = lax.dot_general(lhs, rhs, dnums) + + # 1-pass baseline + how_l = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + how_r = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype='mxfp8_16', + tile_size=16, + ) + res_1p = lax.dot_general( + qarray.dequantize(qarray.quantize(lhs, how_l)), + qarray.dequantize(qarray.quantize(rhs, how_r)), + dnums, + ) + snr_1p = compute_snr_db(true_res, res_1p) + + # 2-pass (lhs_high_precision) + res_2p = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='lhs_high_precision', + lhs_how=how_l, + rhs_how=how_r, + ) + snr_2p = compute_snr_db(true_res, res_2p) + + # 3-pass (triangular) + res_3p = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='triangular', + lhs_how=how_l, + rhs_how=how_r, + ) + snr_3p = compute_snr_db(true_res, res_3p) + + # 4-pass (full_cross) + res_4p = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='full_cross', + lhs_how=how_l, + rhs_how=how_r, + ) + snr_4p = compute_snr_db(true_res, res_4p) + + # Strict SNR progression: 1p < 2p < 3p <= 4p + self.assertGreater(snr_2p, snr_1p) + self.assertGreater(snr_3p, snr_2p) + self.assertGreaterEqual(snr_4p, snr_3p - 0.05) + # Triangular should deliver substantial gain over 1-pass (> 20 dB gain) + self.assertGreater(snr_3p - snr_1p, 20.0) + + @parameterized.parameters( + ((2, 8, 32), (2, 32, 16), (((2,), (1,)), ((0,), (0,))), (2, 8, 16)), + ( + (2, 4, 8, 16), + (2, 4, 16, 8), + (((3,), (2,)), ((0, 1), (0, 1))), + (2, 4, 8, 8), + ), + ((4, 8, 16), (8, 16, 4), (((1, 2), (0, 1)), ((), ())), (4, 4)), + ) + def test_batched_and_higher_rank_matmuls( + self, lhs_shape, rhs_shape, dnums, expected_shape + ): + """Verifies multipass_dot works on batched and multi-axis contractions.""" + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, lhs_shape, dtype=jnp.float32) + rhs = jax.random.normal(k2, rhs_shape, dtype=jnp.float32) + + true_res = lax.dot_general(lhs, rhs, dnums) + + res = multipass.multipass_dot( + lhs, + rhs, + dimension_numbers=dnums, + mode='triangular', + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + ) + + self.assertEqual(res.shape, expected_shape) + self.assertFalse(jnp.isnan(res).any()) + + snr = compute_snr_db(true_res, res) + self.assertGreater(snr, 45.0) + + @parameterized.parameters( + 'triangular', + 'full_cross', + 'lhs_high_precision', + 'rhs_high_precision', + ) + def test_dot_general_qt_multipass_fwd(self, mode): + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, (4, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 4), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + config = dot_general_qt.DotGeneralQtConfig( + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + multipass_mode=mode, + ) + res = dot_general_qt.dot_general_qt(lhs, rhs, dnums, config) + self.assertEqual(res.shape, (4, 4)) + self.assertFalse(jnp.isnan(res).any()) + + def test_gradient_accuracy_and_convergence_vs_float32(self): + """Verifies that multi-pass backward gradients converge to float32.""" + k1, k2, k3 = jax.random.split(self.rng, 3) + lhs = jax.random.normal(k1, (16, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 16), dtype=jnp.float32) + target = jax.random.normal(k3, (16, 16), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + # Loss function with continuous non-trivial incoming gradient + def true_loss(a, b): + out = lax.dot_general(a, b, dnums) + return 0.5 * jnp.sum((out - target) ** 2) + + true_ga, true_gb = jax.grad(true_loss, argnums=(0, 1))(lhs, rhs) + + # 1. Single pass fwd & bwd config (quantizing grad and residual in bwd) + cfg_1p = dot_general_qt.DotGeneralQtConfig( + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + dlhs_grad_qtype='mxfp8_16', + drhs_grad_qtype='mxfp8_16', + dlhs_residual_qtype='mxfp8_16', + drhs_residual_qtype='mxfp8_16', + dlhs_tile_size=16, + drhs_tile_size=16, + ) + + def loss_1p(a, b): + out = dot_general_qt.dot_general_qt(a, b, dnums, cfg_1p) + return 0.5 * jnp.sum((out - target) ** 2) + + ga_1p, gb_1p = jax.grad(loss_1p, argnums=(0, 1))(lhs, rhs) + snr_ga_1p = compute_snr_db(true_ga, ga_1p) + snr_gb_1p = compute_snr_db(true_gb, gb_1p) + + # 2. Multi-pass triangular fwd & bwd config + cfg_tri = dot_general_qt.DotGeneralQtConfig( + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + multipass_mode='triangular', + dlhs_grad_qtype='mxfp8_16', + drhs_grad_qtype='mxfp8_16', + dlhs_residual_qtype='mxfp8_16', + drhs_residual_qtype='mxfp8_16', + dlhs_tile_size=16, + drhs_tile_size=16, + dlhs_multipass_mode='triangular', + drhs_multipass_mode='triangular', + ) + + def loss_tri(a, b): + out = dot_general_qt.dot_general_qt(a, b, dnums, cfg_tri) + return 0.5 * jnp.sum((out - target) ** 2) + + ga_tri, gb_tri = jax.grad(loss_tri, argnums=(0, 1))(lhs, rhs) + snr_ga_tri = compute_snr_db(true_ga, ga_tri) + snr_gb_tri = compute_snr_db(true_gb, gb_tri) + + self.assertFalse(jnp.isnan(ga_tri).any()) + self.assertFalse(jnp.isnan(gb_tri).any()) + # Multi-pass gradient should achieve high fidelity (> 45 dB SNR) + self.assertGreater(snr_ga_tri, 45.0) + self.assertGreater(snr_gb_tri, 45.0) + # Multi-pass gradient should have substantially higher SNR than single pass + self.assertGreater(snr_ga_tri, snr_ga_1p + 15.0) + self.assertGreater(snr_gb_tri, snr_gb_1p + 15.0) + + @parameterized.parameters( + 'triangular', + 'full_cross', + ) + def test_numerical_edge_cases(self, mode): + """Verifies behavior on zero matrices, asymmetric zeros, and wide dynamic range.""" + dnums = (((1,), (0,)), ((), ())) + + # Zero matrices + zeros_l = jnp.zeros((4, 16), dtype=jnp.float32) + zeros_r = jnp.zeros((16, 4), dtype=jnp.float32) + res_zero = multipass.multipass_dot( + zeros_l, zeros_r, dnums, mode=mode, tile_size=16 + ) + np.testing.assert_array_equal(res_zero, 0.0) + self.assertFalse(jnp.isnan(res_zero).any()) + + # One side zero + k1 = self.rng + normal_l = jax.random.normal(k1, (4, 16), dtype=jnp.float32) + res_one_zero = multipass.multipass_dot( + normal_l, zeros_r, dnums, mode=mode, tile_size=16 + ) + np.testing.assert_array_equal(res_one_zero, 0.0) + self.assertFalse(jnp.isnan(res_one_zero).any()) + + # Extreme scaling (1e-5 to 1e5) + k2, k3 = jax.random.split(self.rng) + small_l = jax.random.normal(k2, (4, 16), dtype=jnp.float32) * 1e-4 + large_r = jax.random.normal(k3, (16, 4), dtype=jnp.float32) * 1e4 + res_scaled = multipass.multipass_dot( + small_l, large_r, dnums, mode=mode, tile_size=16 + ) + self.assertFalse(jnp.isnan(res_scaled).any()) + self.assertFalse(jnp.isinf(res_scaled).any()) + true_scaled = lax.dot_general(small_l, large_r, dnums) + snr = compute_snr_db(true_scaled, res_scaled) + self.assertGreater(snr, 40.0) + + def test_bilinear_scaling_and_transposition_symmetry(self): + """Verifies transposition symmetry and exact bilinear scaling of multi-pass.""" + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, (16, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 16), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + res = multipass.multipass_dot( + lhs, rhs, dnums, mode='triangular', tile_size=16 + ) + + # Transposition symmetry: (A @ B).T == B.T @ A.T + res_t = multipass.multipass_dot( + rhs.T, lhs.T, dnums, mode='triangular', tile_size=16 + ) + np.testing.assert_allclose(res.T, res_t, rtol=1e-5, atol=1e-5) + + # Power-of-2 scaling linearity: multipass(2 * A, B) == 2 * multipass(A, B) + res_scaled = multipass.multipass_dot( + lhs * 2.0, rhs, dnums, mode='triangular', tile_size=16 + ) + np.testing.assert_allclose(res * 2.0, res_scaled, rtol=1e-5, atol=1e-5) + + def test_multipass_jit_compatibility(self): + """Verifies that multipass_dot compiles under jax.jit and matches eager.""" + k1, k2 = jax.random.split(self.rng) + lhs = jax.random.normal(k1, (8, 32), dtype=jnp.float32) + rhs = jax.random.normal(k2, (32, 8), dtype=jnp.float32) + dnums = (((1,), (0,)), ((), ())) + + eager_res = multipass.multipass_dot( + lhs, rhs, dnums, mode='triangular', tile_size=16 + ) + + jitted_dot = jax.jit( + multipass.multipass_dot, + static_argnames=( + 'dimension_numbers', + 'mode', + 'lhs_qtype', + 'rhs_qtype', + 'tile_size', + ), + ) + jit_res = jitted_dot(lhs, rhs, dnums, mode='triangular', tile_size=16) + + np.testing.assert_array_equal(eager_res, jit_res) + + def test_transformer_attention_multipass_accuracy(self): + """Verifies multi-pass accuracy on 4D Transformer attention operations.""" + # (batch=2, num_heads=4, seq_len=16, head_dim=32) + k1, k2, k3 = jax.random.split(self.rng, 3) + query = jax.random.normal(k1, (2, 4, 16, 32), dtype=jnp.float32) + key = jax.random.normal(k2, (2, 4, 16, 32), dtype=jnp.float32) + value = jax.random.normal(k3, (2, 4, 16, 32), dtype=jnp.float32) + + # Q @ K.T: contracting head_dim (axis 3) + qk_dnums = (((3,), (3,)), ((0, 1), (0, 1))) + true_logits = lax.dot_general(query, key, qk_dnums) + + # 1-pass baseline + how_q = dot_general.get_how_to_quantize( + dimension_numbers=qk_dnums, + ndims=(4, 4), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + how_k = dot_general.get_how_to_quantize( + dimension_numbers=qk_dnums, + ndims=(4, 4), + for_lhs=False, + qtype='mxfp8_16', + tile_size=16, + ) + logits_1p = lax.dot_general( + qarray.dequantize(qarray.quantize(query, how_q)), + qarray.dequantize(qarray.quantize(key, how_k)), + qk_dnums, + ) + snr_1p = compute_snr_db(true_logits, logits_1p) + + # Multi-pass triangular (3-pass) + logits_tri = multipass.multipass_dot( + query, + key, + dimension_numbers=qk_dnums, + mode='triangular', + lhs_how=how_q, + rhs_how=how_k, + ) + snr_tri = compute_snr_db(true_logits, logits_tri) + + self.assertEqual(logits_tri.shape, (2, 4, 16, 16)) + self.assertGreater(snr_tri, 50.0) + self.assertGreater(snr_tri - snr_1p, 20.0) + + # Attn_weights @ V: contracting seq_len (axis 3 of weights, axis 2 of V) + weights = jax.nn.softmax(logits_tri, axis=-1) + av_dnums = (((3,), (2,)), ((0, 1), (0, 1))) + true_context = lax.dot_general(weights, value, av_dnums) + + context_tri = multipass.multipass_dot( + weights, + value, + dimension_numbers=av_dnums, + mode='triangular', + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + ) + self.assertEqual(context_tri.shape, (2, 4, 16, 32)) + snr_context = compute_snr_db(true_context, context_tri) + self.assertGreater(snr_context, 50.0) + + def test_normal_distribution_sqnr_benchmark(self): + """Verifies SQNR and relative error on N(0, 1) inputs against reference table. + + Cross-checks the QWIX multi-pass and asymmetric implementation against the + published benchmarks in zfc_emulation_utils/README.md Table 2.1. + """ + k1, k2 = jax.random.split(self.rng) + shape_l = (256, 512) + shape_r = (512, 256) + dnums = (((1,), (0,)), ((), ())) + + # Continuous standard normal Gaussian inputs N(0, 1) + lhs = jax.random.normal(k1, shape_l, dtype=jnp.float32) + rhs = jax.random.normal(k2, shape_r, dtype=jnp.float32) + ref_f32 = lax.dot_general(lhs, rhs, dnums) + + # 1. Single-Pass FP8 (1 pass, block size 16) - target ~28.41 dB + how_l_fp8 = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=True, + qtype='mxfp8_16', + tile_size=16, + ) + how_r_fp8 = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype='mxfp8_16', + tile_size=16, + ) + res_fp8_1p = lax.dot_general( + qarray.dequantize(qarray.quantize(lhs, how_l_fp8)), + qarray.dequantize(qarray.quantize(rhs, how_r_fp8)), + dnums, + ) + snr_fp8_1p = float(compute_snr_db(ref_f32, res_fp8_1p)) + err_fp8_1p = float(compute_relative_error(ref_f32, res_fp8_1p)) + + # 2. Asymmetric FP8 2-Pass (2 passes, block size 16) - target ~31.56 dB + res_fp8_2p = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='lhs_high_precision', + lhs_how=how_l_fp8, + rhs_how=how_r_fp8, + ) + snr_fp8_2p = float(compute_snr_db(ref_f32, res_fp8_2p)) + err_fp8_2p = float(compute_relative_error(ref_f32, res_fp8_2p)) + + # 3. Multi-Pass FP8 Triangular (3 passes, block size 16) - target ~54.84 dB + res_fp8_tri = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='triangular', + lhs_how=how_l_fp8, + rhs_how=how_r_fp8, + ) + snr_fp8_tri = float(compute_snr_db(ref_f32, res_fp8_tri)) + err_fp8_tri = float(compute_relative_error(ref_f32, res_fp8_tri)) + + # 4. Multi-Pass FP8 Full Cross (4 passes, block size 16) - target ~55.59 dB + res_fp8_full = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='full_cross', + lhs_how=how_l_fp8, + rhs_how=how_r_fp8, + ) + snr_fp8_full = float(compute_snr_db(ref_f32, res_fp8_full)) + err_fp8_full = float(compute_relative_error(ref_f32, res_fp8_full)) + + # Also evaluate with bfloat16 accumulation (matching zfc_emulation_utils) + res_fp8_tri_bf16 = multipass.multipass_dot( + lhs, + rhs, + dnums, + mode='triangular', + lhs_how=how_l_fp8, + rhs_how=how_r_fp8, + preferred_element_type=jnp.bfloat16, + ) + snr_fp8_tri_bf16 = float(compute_snr_db(ref_f32, res_fp8_tri_bf16)) + err_fp8_tri_bf16 = float(compute_relative_error(ref_f32, res_fp8_tri_bf16)) + + print('=== SQNR BENCHMARKS ===') + print(f'FP8 1-pass: SNR={snr_fp8_1p:.2f} dB, err={err_fp8_1p:.4f}') + print(f'FP8 2-pass: SNR={snr_fp8_2p:.2f} dB, err={err_fp8_2p:.4f}') + print(f'FP8 3-pass (tri): SNR={snr_fp8_tri:.2f} dB, err={err_fp8_tri:.4f}') + print(f'FP8 4-pass (full): SNR={snr_fp8_full:.2f} dB') + print(f'FP8 3-pass (bf16): SNR={snr_fp8_tri_bf16:.2f} dB') + + # Reference values aligned with zfc_emulation_utils/README.md + # 1-pass FP8: ~28.4 dB, relative error ~0.038 + self.assertAlmostEqual(snr_fp8_1p, 28.41, delta=1.5) + self.assertAlmostEqual(err_fp8_1p, 0.0380, delta=0.01) + + # 2-pass Asymmetric FP8: ~31.5 dB (+3.1 dB over 1-pass) + self.assertAlmostEqual(snr_fp8_2p, 31.56, delta=1.5) + self.assertAlmostEqual(err_fp8_2p, 0.0264, delta=0.01) + self.assertGreater(snr_fp8_2p, snr_fp8_1p + 2.0) + + # 3-pass Triangular FP8: > 54.0 dB (fp32), > 49.0 dB (bf16) + self.assertGreater(snr_fp8_tri, 54.0) + self.assertGreater(snr_fp8_tri_bf16, 49.0) + self.assertLess(err_fp8_tri, 0.002) + self.assertLess(err_fp8_tri_bf16, 0.005) + + # 4-pass Full Cross FP8: > 55.0 dB in float32 + self.assertGreater(snr_fp8_full, 55.0) + self.assertLess(err_fp8_full, 0.002) + self.assertGreaterEqual(snr_fp8_full, snr_fp8_tri - 0.2) + + # Strict SNR progression hierarchy across passes + self.assertGreater(snr_fp8_2p, snr_fp8_1p) + self.assertGreater(snr_fp8_tri, snr_fp8_2p) + self.assertGreaterEqual(snr_fp8_full, snr_fp8_tri - 0.2) + + +if __name__ == '__main__': + absltest.main() diff --git a/tests/_src/core/numerics_test.py b/tests/_src/core/numerics_test.py index 5b42a2e7..fa454af5 100644 --- a/tests/_src/core/numerics_test.py +++ b/tests/_src/core/numerics_test.py @@ -202,6 +202,39 @@ def test_mxfp(self): out_nvfp4 = numerics.convert_from(in_nvfp4, "nvfp4") self.assertIs(in_nvfp4, out_nvfp4) + def test_mxint(self): + with self.subTest("mxint8_bound"): + bound = numerics.get_symmetric_bound("mxint8") + self.assertEqual(bound, numerics.get_symmetric_bound(jnp.int8)) + + with self.subTest("mxint8_conversion"): + in_array = jnp.array([10.0, 150.0, -200.0], dtype=jnp.float32) + converted = numerics.convert_to(in_array, "mxint8") + self.assertEqual(converted.dtype, jnp.int8) + expected = jnp.array([10, 127, -128], dtype=jnp.int8) + self._assert_equal(converted, expected) + + with self.subTest("mxint8_convert_from"): + in_mxint8 = jnp.array([1, -2], dtype=jnp.int8) + out_mxint8 = numerics.convert_from(in_mxint8, "mxint8") + self.assertIs(in_mxint8, out_mxint8) + + with self.subTest("mxint4_bound"): + bound = numerics.get_symmetric_bound("mxint4") + self.assertEqual(bound, 7.5) + + with self.subTest("mxint4_conversion"): + in_array = jnp.array([3.0, 15.0, -20.0], dtype=jnp.float32) + converted = numerics.convert_to(in_array, "mxint4") + self.assertEqual(converted.dtype, jnp.int4) + expected = jnp.array([3, 7, -8], dtype=jnp.int4) + self._assert_equal(converted, expected) + + with self.subTest("mxint4_convert_from"): + in_mxint4 = jnp.array([1, -2], dtype=jnp.int4) + out_mxint4 = numerics.convert_from(in_mxint4, "mxint4") + self.assertIs(in_mxint4, out_mxint4) + if __name__ == "__main__": absltest.main() diff --git a/tests/_src/core/qarray_test.py b/tests/_src/core/qarray_test.py index 954da736..e49a4148 100644 --- a/tests/_src/core/qarray_test.py +++ b/tests/_src/core/qarray_test.py @@ -516,6 +516,34 @@ def test_mxfp_tile_size_validation(self): tiled_axes={1: 32}, ) + with self.assertRaisesRegex( + ValueError, 'Format mxint8 requires `tiled_axes` to be specified.' + ): + qarray.HowToQuantize(qtype='mxint8') + + with self.assertRaisesRegex( + ValueError, + 'Format mxint8 requires a tile size of 32, but axis 1 got 16', + ): + qarray.HowToQuantize( + qtype='mxint8', + tiled_axes={1: 16}, + ) + + with self.assertRaisesRegex( + ValueError, 'Format mxint4 requires `tiled_axes` to be specified.' + ): + qarray.HowToQuantize(qtype='mxint4') + + with self.assertRaisesRegex( + ValueError, + 'Format mxint4 requires a tile size of 32, but axis 1 got 16', + ): + qarray.HowToQuantize( + qtype='mxint4', + tiled_axes={1: 16}, + ) + @parameterized.named_parameters( dict( testcase_name='mxfp4', @@ -572,6 +600,101 @@ def test_compute_scale_zero_point_mxfp_oas_bias( jnp.array_equal(scale, jnp.array(expected_scales, dtype=scale.dtype)) ) + def test_compute_scale_zero_point_mxint_power_of_2(self): + calibration_i8 = { + 'absmax': jnp.array([63.75, 64.0, 127.5, 127.6, 255.0, 256.0]) + } + scale, zero_point = qarray.compute_scale_zero_point( + calibration_i8, 'mxint8' + ) + self.assertIsNone(zero_point) + expected_scales = jnp.array( + [0.5, 1.0, 1.0, 2.0, 2.0, 4.0], dtype=scale.dtype + ) + self.assertTrue(jnp.array_equal(scale, expected_scales)) + + calibration_i4 = {'absmax': jnp.array([3.75, 3.8, 7.5, 7.6, 15.0, 15.1])} + scale_i4, zero_point_i4 = qarray.compute_scale_zero_point( + calibration_i4, 'mxint4' + ) + self.assertIsNone(zero_point_i4) + expected_scales_i4 = jnp.array( + [0.5, 1.0, 1.0, 2.0, 2.0, 4.0], dtype=scale_i4.dtype + ) + self.assertTrue(jnp.array_equal(scale_i4, expected_scales_i4)) + + def test_mxint8_quantize_dequantize(self): + x = jnp.array( + [[10.0, -20.0, 30.0, -40.0] * 8, [50.0, -60.0, 70.0, -80.0] * 8], + dtype=jnp.float32, + ) # shape (2, 32) + how = qarray.HowToQuantize( + qtype='mxint8', channelwise_axes=[0], tiled_axes={1: 32} + ) + q = qarray.quantize(x, how) + self.assertEqual(q.qvalue.dtype, jnp.int8) + self.assertEqual(q.scale.shape, (2, 1)) + self.assertIsNone(q.zero_point) + # Scale must be a power of 2 + log2_scale = jnp.log2(q.scale) + self.assertTrue(jnp.all(jnp.equal(log2_scale, jnp.round(log2_scale)))) + # Check dequantize + deq = qarray.dequantize(q) + self.assertEqual(deq.shape, x.shape) + # Dequantized values should be close to original + self.assertTrue(jnp.allclose(deq, x, atol=2.0)) + + def test_mxint4_quantize_dequantize(self): + x = jnp.array( + [[1.0, -2.0, 3.0, -4.0] * 8, [5.0, -6.0, 7.0, -7.0] * 8], + dtype=jnp.float32, + ) # shape (2, 32) + how = qarray.HowToQuantize( + qtype='mxint4', channelwise_axes=[0], tiled_axes={1: 32} + ) + q = qarray.quantize(x, how) + self.assertEqual(q.qvalue.dtype, jnp.int4) + self.assertEqual(q.scale.shape, (2, 1)) + self.assertIsNone(q.zero_point) + # Scale must be a power of 2 + log2_scale = jnp.log2(q.scale) + self.assertTrue(jnp.all(jnp.equal(log2_scale, jnp.round(log2_scale)))) + # Check dequantize + deq = qarray.dequantize(q) + self.assertEqual(deq.shape, x.shape) + # Dequantized values should be close to original + self.assertTrue(jnp.allclose(deq, x, atol=1.0)) + + def test_mxint8_sqnr(self): + key = jax.random.PRNGKey(0) + x = jax.random.normal(key, (4, 128), dtype=jnp.float32) + how = qarray.HowToQuantize( + qtype='mxint8', channelwise_axes=[0], tiled_axes={1: 32} + ) + q = qarray.quantize(x, how) + deq = qarray.dequantize(q) + signal_power = jnp.mean(jnp.square(x)) + noise_power = jnp.mean(jnp.square(x - deq)) + sqnr = float(10.0 * jnp.log10(signal_power / noise_power)) + self.assertGreater( + sqnr, 35.0, f'mxint8 SQNR {sqnr:.2f} dB is below expected 35 dB' + ) + + def test_mxint4_sqnr(self): + key = jax.random.PRNGKey(0) + x = jax.random.normal(key, (4, 128), dtype=jnp.float32) + how = qarray.HowToQuantize( + qtype='mxint4', channelwise_axes=[0], tiled_axes={1: 32} + ) + q = qarray.quantize(x, how) + deq = qarray.dequantize(q) + signal_power = jnp.mean(jnp.square(x)) + noise_power = jnp.mean(jnp.square(x - deq)) + sqnr = float(10.0 * jnp.log10(signal_power / noise_power)) + self.assertGreater( + sqnr, 15.0, f'mxint4 SQNR {sqnr:.2f} dB is below expected 15 dB' + ) + if __name__ == '__main__': absltest.main()