From 1e79de589a426076a0637e408e407a8d6d728af6 Mon Sep 17 00:00:00 2001 From: Gil Tabak Date: Fri, 11 Sep 2026 14:31:00 -0700 Subject: [PATCH] Support multi-pass residual quantization in QWIX via Jax - Support multi-pass residual decomposition (residual_decompose) with format-based scale exponent shifting. - Support multi-pass matrix multiplication modes: 1-pass, 2-pass (lhs_high_precision, rhs_high_precision), 3-pass (triangular), and 4-pass (full_cross) in multipass_dot.py. - Support multipass_mode at the dot_general primitive level, allowing direct invocation for inference and PTQ. - Support forward and backward multipass_mode in DotGeneralQtConfig, preserving calibration statistics collection and standard quantized residual packing. - Add unit tests in multipass_dot_test verifying: - Algebraic component equivalence across all multi-pass modes. - SNR hierarchy across normal, uniform, and log-normal distributions. - Hardware FP8 accumulation path preservation (jnp.float8_e4m3fn in dot_general for channelwise and subchannel tile_size >= 128) vs scaled_matmul and dequantized bfloat16 paths for mxfp8_16. - Backward pass gradient accuracy and calibration statistics collection. - Multi-pass benchmarks across 1, 2, 3, and 4 passes. PiperOrigin-RevId: 980022272 --- qwix/_src/core/dot_general.py | 19 +- qwix/_src/core/dot_general_qt.py | 46 +- qwix/_src/core/multipass_dot.py | 309 ++++++++ tests/_src/core/multipass_dot_test.py | 1004 +++++++++++++++++++++++++ 4 files changed, 1374 insertions(+), 4 deletions(-) create mode 100644 qwix/_src/core/multipass_dot.py create mode 100644 tests/_src/core/multipass_dot_test.py diff --git a/qwix/_src/core/dot_general.py b/qwix/_src/core/dot_general.py index 6b390629..8707c552 100644 --- a/qwix/_src/core/dot_general.py +++ b/qwix/_src/core/dot_general.py @@ -429,7 +429,24 @@ def dot_general( Returns: a floating-point jax.Array. """ - # Try hardware-accelerated MXFP dot first + if kwargs.get('multipass_mode') is not None: + if isinstance(lhs, qarray.QArray) or isinstance(rhs, qarray.QArray): + raise ValueError( + 'Inputs to multipass dot_general must strictly be unquantized' + f' jax.Array, but got lhs={type(lhs)}, rhs={type(rhs)}.' + ) + from qwix._src.core import multipass_dot # pylint: disable=g-import-not-at-top + + return multipass_dot.multipass_dot_general( + lhs, + rhs, + dimension_numbers=dimension_numbers, + precision=precision, + preferred_element_type=preferred_element_type, + **kwargs, + ) + + # Try hardware-accelerated MXFP dot. mxfp_result = mxfp_dot.mxfp_dot_general( lhs, rhs, dimension_numbers, preferred_element_type ) diff --git a/qwix/_src/core/dot_general_qt.py b/qwix/_src/core/dot_general_qt.py index 0f86a48f..827e262c 100644 --- a/qwix/_src/core/dot_general_qt.py +++ b/qwix/_src/core/dot_general_qt.py @@ -43,6 +43,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: str | None = None # Backward pass (dlhs). dlhs_grad_qtype: jax.typing.DTypeLike | None = None # incoming gradient @@ -50,6 +51,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: str | None = None # Backward pass (drhs). drhs_grad_qtype: jax.typing.DTypeLike | None = None # incoming gradient @@ -57,6 +59,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: str | 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 @@ -202,7 +205,7 @@ def _requires_unquantized_residual( 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. + in the backward pass. Multi-pass modes also require original residuals. Args: config: The quantization configuration. @@ -211,7 +214,13 @@ def _requires_unquantized_residual( Returns: True if the backward contraction cannot reuse the forward operand. """ - return config.use_original_residuals or _is_block_scaled(operand_qt) + 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( @@ -333,7 +342,23 @@ def dot_general_qt_fwd( saved_rhs_calibration, config, ) - return dot_general.dot_general(lhs, rhs, dimension_numbers), residuals + if config.multipass_mode is not None: + # Multi-pass residual decomposition requires unquantized inputs (lhs_in, + # rhs_in) to compute successive residual passes (A_0 = quant(A), + # A_1 = quant(A - dequant(A_0))). Passing already-quantized QArrays (lhs, + # rhs) would cause residual decomposition to zero out subsequent passes. + out = dot_general.dot_general( + lhs_in, + rhs_in, + dimension_numbers, + multipass_mode=config.multipass_mode, + lhs_qtype=config.lhs_qtype or jnp.float8_e4m3fn, + rhs_qtype=config.rhs_qtype or jnp.float8_e4m3fn, + tile_size=config.tile_size, + ) + else: + out = dot_general.dot_general(lhs, rhs, dimension_numbers) + return out, residuals def dot_general_qt_bwd( @@ -378,6 +403,21 @@ 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: + grad_res = dot_general.dot_general( + g, + y, + dimension_numbers=bwd_dnums, + multipass_mode=multipass_mode, + lhs_qtype=g_qtype or jnp.float8_e4m3fn, + rhs_qtype=y_qtype or jnp.float8_e4m3fn, + tile_size=g_tile_size, + ) + 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) diff --git a/qwix/_src/core/multipass_dot.py b/qwix/_src/core/multipass_dot.py new file mode 100644 index 00000000..a07a8939 --- /dev/null +++ b/qwix/_src/core/multipass_dot.py @@ -0,0 +1,309 @@ +# 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 for fp8 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. + +All decomposed passes are evaluated on quantized operands directly through the +hardware FP8 path (via `_fast_dot_general` with hardware FP8 accumulation) or +via `jax.nn.scaled_matmul` on supported GPUs. + +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). +This design choice is made because: + +1. Hardware efficiency: It eliminates the expensive second 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. +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 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 bits + yields similar SQNR to independent scale choices. +""" + +from collections.abc import Mapping +from typing import Any, Literal, TypeAlias +import jax +import jax.numpy as jnp +from qwix._src.core import dot_general as dg +from qwix._src.core import qarray + +MultiPassMode: TypeAlias = Literal[ + 'triangular', + 'full_cross', + 'lhs_high_precision', + 'rhs_high_precision', +] + + +def get_how_to_quantize( + *, + dimension_numbers: jax.lax.DotDimensionNumbers, + ndims: tuple[int, int], + for_lhs: bool, + tile_size: Mapping[int, int | float] | int | float | None, + **kwargs: Any, +) -> qarray.HowToQuantize: + """Get how to quantize from dimension_numbers and remaining_dims.""" + if for_lhs: + ndim = ndims[0] + contracting_axes = dimension_numbers[0][0] + else: + ndim = ndims[1] + contracting_axes = dimension_numbers[0][1] + + if isinstance(tile_size, Mapping): + tiled_axes = tile_size + else: + tiled_axes = {} + if tile_size: + tiled_axes = {contracting_axes[-1]: tile_size} + + channelwise_axes = sorted( + set(range(ndim)) - set(contracting_axes) - set(tiled_axes.keys()) + ) + + return qarray.HowToQuantize( + channelwise_axes=channelwise_axes, + tiled_axes=tiled_axes, + **kwargs, + ) + + +def get_residual_scale_shift_bits(qtype: jax.typing.DTypeLike) -> int | None: + """Returns the power-of-2 scale shift bits for residual quantization. + + For floating-point and microscaled 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: + - fp8 / float8_e4m3fn / mxfp8 / mxfp8_16 (E4M3): shift by 4 bits (2^-4 = + 1/16). + - mxfp4 / nvfp4 / float4_e2m1fn (E2M1): shift by 2 bits (2^-2 = 1/4). + - float8_e5m2 (E5M2): 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. + """ + # Normalize synthetic string aliases to standard JAX dtypes. + match qtype: + case 'mxfp8' | 'mxfp8_16' | 'float8_e4m3' | 'fp8': + qtype = jnp.float8_e4m3fn + case 'mxfp4' | 'nvfp4': + qtype = jnp.float4_e2m1fn + + try: + dt = jnp.dtype(qtype) + except (TypeError, ValueError): + return None + + if dt == jnp.float8_e4m3fn: + return 4 + if dt in (jnp.float4_e2m1fn, jnp.float8_e5m2): + return 2 + return None + + +def residual_decompose( + x: qarray.MaybeQArray, + 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, float8_e4m3fn), 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. + + If x is already a QArray, it is used as the first pass, and subsequent passes + are zero. + + Args: + x: Input array or QArray. + 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, ...). + """ + if isinstance(x, qarray.QArray): + raise ValueError( + 'Input to residual_decompose must strictly be unquantized jax.Array,' + f' but got {type(x)}.' + ) + + 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_general( + lhs: jax.Array, + rhs: jax.Array, + dimension_numbers: jax.lax.DotDimensionNumbers = (((1,), (0,)), ((), ())), + precision: jax.lax.PrecisionLike = None, + preferred_element_type: jax.typing.DTypeLike | None = None, + **kwargs: Any, +) -> jax.Array: + """Executes multi-pass emulated matrix multiplication. + + Inputs must strictly be unquantized jax.Array to allow residual calculation. + + Args: + lhs: Left-hand side unquantized array. + rhs: Right-hand side unquantized array. + dimension_numbers: Standard JAX dot_general dimension specification. + precision: The precision for dot_general. + preferred_element_type: Output/accumulator dtype. + **kwargs: Additional arguments forwarded to dot_general, including optional + 'multipass_mode', 'lhs_how', 'rhs_how', 'lhs_qtype', 'rhs_qtype', + 'tile_size'. + + Returns: + The resulting matrix product array. + """ + if isinstance(lhs, qarray.QArray) or isinstance(rhs, qarray.QArray): + raise ValueError( + 'Inputs to multipass_dot_general must strictly be unquantized' + f' jax.Array, but got lhs={type(lhs)}, rhs={type(rhs)}.' + ) + + kwargs = dict(kwargs) + mode: MultiPassMode = kwargs.pop( + 'mode', kwargs.pop('multipass_mode', 'triangular') + ) + lhs_how = kwargs.pop('lhs_how', None) + rhs_how = kwargs.pop('rhs_how', None) + lhs_qtype = kwargs.pop('lhs_qtype', jnp.float8_e4m3fn) + rhs_qtype = kwargs.pop('rhs_qtype', jnp.float8_e4m3fn) + tile_size = kwargs.pop('tile_size', None) + kwargs.pop('dot_general_fn', None) + + if lhs_how is None: + lhs_how = 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 = get_how_to_quantize( + dimension_numbers=dimension_numbers, + ndims=(lhs.ndim, rhs.ndim), + for_lhs=False, + qtype=rhs_qtype, + tile_size=tile_size, + ) + + def _dot(a: qarray.QArray, b: qarray.QArray) -> jax.Array: + return dg.dot_general( + a, + b, + dimension_numbers=dimension_numbers, + precision=precision, + preferred_element_type=preferred_element_type, + **kwargs, + ) + + 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(a0, b0) + c1 = _dot(a1, b0) + 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(a0, b0) + c1 = _dot(a0, b1) + 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(a0, b0) + c01 = _dot(a0, b1) + c10 = _dot(a1, b0) + res = c00 + c01 + c10 + if mode == 'full_cross': + c11 = _dot(a1, b1) + 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/tests/_src/core/multipass_dot_test.py b/tests/_src/core/multipass_dot_test.py new file mode 100644 index 00000000..15937f57 --- /dev/null +++ b/tests/_src/core/multipass_dot_test.py @@ -0,0 +1,1004 @@ +# 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_dot +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 MultiPassDotTest(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=jnp.float8_e4m3fn, + tile_size=None, + ) + passes = multipass_dot.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_quantized_input_raises_error(self): + """Verifies that passing QArray to multipass raises ValueError.""" + 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=jnp.float8_e4m3fn, + tile_size=None, + ) + qx = qarray.quantize(x, how) + with self.assertRaisesRegex(ValueError, 'must strictly be unquantized'): + multipass_dot.residual_decompose(qx, how) + with self.assertRaisesRegex(ValueError, 'must strictly be unquantized'): + multipass_dot.multipass_dot_general(qx, x) + with self.assertRaisesRegex(ValueError, 'must strictly be unquantized'): + multipass_dot.multipass_dot_general(x, qx) + + 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=jnp.float8_e4m3fn, + tile_size=None, + ) + n_passes = 4 + passes = multipass_dot.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( + (jnp.float8_e4m3fn, None), + ('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_dot.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=jnp.float8_e4m3fn, + tile_size=None, + ) + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + + actual = multipass_dot.multipass_dot_general( + lhs, + rhs, + dimension_numbers=dnums, + mode=mode, + lhs_how=lhs_how, + rhs_how=rhs_how, + ) + + # Compute manual reference terms from component passes + a_passes = multipass_dot.residual_decompose(lhs, lhs_how, n_passes=2) + b_passes = multipass_dot.residual_decompose(rhs, rhs_how, n_passes=2) + + c00 = dot_general.dot_general(a_passes[0], b_passes[0], dnums) + c01 = dot_general.dot_general(a_passes[0], b_passes[1], dnums) + c10 = dot_general.dot_general(a_passes[1], b_passes[0], dnums) + c11 = dot_general.dot_general(a_passes[1], b_passes[1], 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=jnp.float8_e4m3fn, + tile_size=None, + ) + rhs_how = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + + res_tri = multipass_dot.multipass_dot_general( + lhs, rhs, dnums, mode='triangular', lhs_how=lhs_how, rhs_how=rhs_how + ) + res_full = multipass_dot.multipass_dot_general( + lhs, rhs, dnums, mode='full_cross', lhs_how=lhs_how, rhs_how=rhs_how + ) + + a_passes = multipass_dot.residual_decompose(lhs, lhs_how, n_passes=2) + b_passes = multipass_dot.residual_decompose(rhs, rhs_how, n_passes=2) + a1_b1 = dot_general.dot_general(a_passes[1], b_passes[1], 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.named_parameters( + ( + 'batched_matmul_3d', + (2, 8, 32), + (2, 32, 16), + (((2,), (1,)), ((0,), (0,))), + (2, 8, 16), + ), + ( + 'multi_axis_contraction_weight_grad', + (2, 4, 8), + (2, 4, 16), + (((0, 1), (0, 1)), ((), ())), + (8, 16), + ), + ( + 'high_rank_attention_proj_4d', + (2, 4, 8, 16), + (16, 32), + (((3,), (0,)), ((), ())), + (2, 4, 8, 32), + ), + ) + def test_multipass_dot_general_arbitrary_dimensions( + self, lhs_shape, rhs_shape, dnums, expected_shape + ): + """Verifies multipass_dot_general works with arbitrary dimension numbers.""" + 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) + + # 1. Direct call to multipass_dot_general + out = multipass_dot.multipass_dot_general( + lhs, rhs, dimension_numbers=dnums, mode='triangular' + ) + self.assertEqual(out.shape, expected_shape) + + # 2. Dispatch via dot_general.dot_general with multipass_mode + out_dg = dot_general.dot_general( + lhs, rhs, dimension_numbers=dnums, multipass_mode='triangular' + ) + self.assertEqual(out_dg.shape, expected_shape) + np.testing.assert_allclose(out, out_dg, rtol=1e-5, atol=1e-5) + + # 3. Verify high numerical accuracy against unquantized FP32 reference + ref = jax.lax.dot_general(lhs, rhs, dimension_numbers=dnums) + self.assertEqual(ref.shape, expected_shape) + snr = compute_snr_db(ref, out) + self.assertGreater(snr, 35.0) + + @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=jnp.float8_e4m3fn, + tile_size=None, + ) + how_r = dot_general.get_how_to_quantize( + dimension_numbers=dnums, + ndims=(2, 2), + for_lhs=False, + qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + lhs, + rhs, + dimension_numbers=dnums, + mode='triangular', + lhs_qtype=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + + 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=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + 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()) + + @parameterized.parameters( + 'triangular', + 'full_cross', + 'lhs_high_precision', + 'rhs_high_precision', + ) + def test_dot_general_multipass_mode_dispatch(self, mode): + """Verifies that dot_general.dot_general correctly dispatches multipass_mode.""" + 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,)), ((), ())) + + res_dot_general = dot_general.dot_general( + lhs, rhs, dnums, multipass_mode=mode + ) + res_multipass = multipass_dot.multipass_dot_general( + lhs, rhs, dnums, mode=mode + ) + + np.testing.assert_array_equal(res_dot_general, res_multipass) + + def test_fp8_vs_mxfp8_accumulation_path(self): + """Verifies standard FP8 preserves hardware FP8 dot_general, while MXFP8_16 dequantizes.""" + dnums = (((1,), (0,)), ((), ())) + x = jnp.ones((4, 256), dtype=jnp.bfloat16) + y = jnp.ones((256, 4), dtype=jnp.bfloat16) + + # 1. Multi-pass with standard FP8 (default: jnp.float8_e4m3fn, + # tile_size=None): Tracing with jax.make_jaxpr verifies the dot_general + # primitive receives float8_e4m3fn inputs directly. + jaxpr_mp_fp8 = jax.make_jaxpr( + lambda a, b: multipass_dot.multipass_dot_general( + a, + b, + dnums, + mode='triangular', + ) + )(x, y) + fp8_dot_eqns = [ + eqn for eqn in jaxpr_mp_fp8.eqns if eqn.primitive.name == 'dot_general' + ] + self.assertNotEmpty(fp8_dot_eqns) + for eqn in fp8_dot_eqns: + self.assertEqual(eqn.invars[0].aval.dtype, jnp.float8_e4m3fn) + self.assertEqual(eqn.invars[1].aval.dtype, jnp.float8_e4m3fn) + + # 2. Multi-pass with FP8 and subchannel tile_size=256 (>= 128 threshold): + # Also preserves hardware FP8 dot_general. + jaxpr_mp_fp8_subchan = jax.make_jaxpr( + lambda a, b: multipass_dot.multipass_dot_general( + a, + b, + dnums, + mode='triangular', + tile_size=256, + ) + )(x, y) + subchan_dot_eqns = [ + eqn + for eqn in jaxpr_mp_fp8_subchan.eqns + if eqn.primitive.name == 'dot_general' + ] + self.assertNotEmpty(subchan_dot_eqns) + for eqn in subchan_dot_eqns: + self.assertEqual(eqn.invars[0].aval.dtype, jnp.float8_e4m3fn) + self.assertEqual(eqn.invars[1].aval.dtype, jnp.float8_e4m3fn) + + # 3. Multi-pass with mxfp8_16 (block_size=16): + # Evaluates through scaled_matmul_wrapper (hardware GPU or JAX decomp). + jaxpr_mp_mxfp = jax.make_jaxpr( + lambda a, b: multipass_dot.multipass_dot_general( + a, + b, + dnums, + mode='triangular', + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + ) + )(x, y) + jit_eqns = [ + eqn for eqn in jaxpr_mp_mxfp.eqns if eqn.primitive.name == 'jit' + ] + scaled_matmul_calls = [ + je + for je in jit_eqns + if any( + e.primitive.name == 'scaled_matmul_wrapper' + for e in je.params['jaxpr'].eqns + ) + ] + # For 3-pass triangular mode, exactly 3 scaled_matmul passes (c00, c01, c10) + # are executed. + self.assertLen(scaled_matmul_calls, 3) + + # 4. Multi-axis contraction with mxfp8_16 and K > 16 (len(ca) > 1, e.g. + # backward pass G^T Y): scaled_matmul returns None and falls back to + # _slow_dot_general (tile_size=16 < 128), dequantizing to bfloat16. + x_multi = jnp.ones((2, 4, 32), dtype=jnp.bfloat16) + y_multi = jnp.ones((4, 32, 2), dtype=jnp.bfloat16) + dnums_multi = (((1, 2), (0, 1)), ((), ())) + jaxpr_mp_mxfp_multi = jax.make_jaxpr( + lambda a, b: multipass_dot.multipass_dot_general( + a, + b, + dnums_multi, + mode='triangular', + lhs_qtype='mxfp8_16', + rhs_qtype='mxfp8_16', + tile_size=16, + dot_general_fn=dot_general.dot_general, + ) + )(x_multi, y_multi) + dequant_dot_eqns = [ + eqn + for eqn in jaxpr_mp_mxfp_multi.eqns + if eqn.primitive.name == 'dot_general' + ] + self.assertNotEmpty(dequant_dot_eqns) + for eqn in dequant_dot_eqns: + self.assertEqual(eqn.invars[0].aval.dtype, jnp.bfloat16) + self.assertEqual(eqn.invars[1].aval.dtype, jnp.bfloat16) + + def test_dot_general_qt_multipass_calibrates_stats(self): + """Verifies that lhs/rhs_collect_quant_stat are invoked when multipass_mode is set.""" + 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,)), ((), ())) + + lhs_stat_called = [] + rhs_stat_called = [] + + def lhs_stat_hook(calib): + lhs_stat_called.append(calib) + return calib + + def rhs_stat_hook(calib): + rhs_stat_called.append(calib) + return calib + + config = dot_general_qt.DotGeneralQtConfig( + lhs_qtype=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + multipass_mode='triangular', + lhs_collect_quant_stat=lhs_stat_hook, + rhs_collect_quant_stat=rhs_stat_hook, + ) + res = dot_general_qt.dot_general_qt(lhs, rhs, dnums, config) + self.assertEqual(res.shape, (4, 4)) + self.assertLen(lhs_stat_called, 1) + self.assertLen(rhs_stat_called, 1) + + def test_dot_general_qt_fwd_multipass_bwd_singlepass_quantized_residual(self): + """Verifies backward pass correctly receives quantized QArray residual when dlhs_multipass_mode is None.""" + 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,)), ((), ())) + + # Fwd has multipass_mode='triangular', but backward has single pass + # (dlhs/drhs_multipass_mode=None) and use_original_residuals=False. + cfg = dot_general_qt.DotGeneralQtConfig( + lhs_qtype=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + multipass_mode='triangular', + use_original_residuals=False, + dlhs_grad_qtype=jnp.float8_e4m3fn, + drhs_grad_qtype=jnp.float8_e4m3fn, + dlhs_residual_qtype=jnp.float8_e4m3fn, + drhs_residual_qtype=jnp.float8_e4m3fn, + ) + + def loss_fn(a, b): + out = dot_general_qt.dot_general_qt(a, b, dnums, cfg) + return 0.5 * jnp.sum((out - target) ** 2) + + ga, gb = jax.grad(loss_fn, argnums=(0, 1))(lhs, rhs) + self.assertFalse(jnp.isnan(ga).any()) + self.assertFalse(jnp.isnan(gb).any()) + self.assertEqual(ga.shape, lhs.shape) + self.assertEqual(gb.shape, rhs.shape) + + 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 fp32 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=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + dlhs_grad_qtype=jnp.float8_e4m3fn, + drhs_grad_qtype=jnp.float8_e4m3fn, + dlhs_residual_qtype=jnp.float8_e4m3fn, + drhs_residual_qtype=jnp.float8_e4m3fn, + ) + + 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=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + multipass_mode='triangular', + dlhs_grad_qtype=jnp.float8_e4m3fn, + drhs_grad_qtype=jnp.float8_e4m3fn, + dlhs_residual_qtype=jnp.float8_e4m3fn, + drhs_residual_qtype=jnp.float8_e4m3fn, + 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_dot.multipass_dot_general( + zeros_l, zeros_r, dnums, mode=mode + ) + 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_dot.multipass_dot_general( + normal_l, zeros_r, dnums, mode=mode + ) + 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_dot.multipass_dot_general( + small_l, large_r, dnums, mode=mode + ) + 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_dot.multipass_dot_general( + lhs, rhs, dnums, mode='triangular' + ) + + # Transposition symmetry: (A @ B).T == B.T @ A.T + res_t = multipass_dot.multipass_dot_general( + rhs.T, lhs.T, dnums, mode='triangular' + ) + 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_dot.multipass_dot_general( + lhs * 2.0, rhs, dnums, mode='triangular' + ) + 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_dot.multipass_dot_general( + lhs, rhs, dnums, mode='triangular' + ) + + jitted_dot = jax.jit( + multipass_dot.multipass_dot_general, + static_argnames=( + 'dimension_numbers', + 'mode', + 'lhs_qtype', + 'rhs_qtype', + 'tile_size', + ), + ) + jit_res = jitted_dot(lhs, rhs, dnums, mode='triangular') + + np.testing.assert_allclose(eager_res, jit_res, rtol=1e-5, atol=1e-5) + + 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=jnp.float8_e4m3fn, + tile_size=None, + ) + how_k = dot_general.get_how_to_quantize( + dimension_numbers=qk_dnums, + ndims=(4, 4), + for_lhs=False, + qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + 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_dot.multipass_dot_general( + 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, 48.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_dot.multipass_dot_general( + weights, + value, + dimension_numbers=av_dnums, + mode='triangular', + lhs_qtype=jnp.float8_e4m3fn, + rhs_qtype=jnp.float8_e4m3fn, + tile_size=None, + ) + 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,)), ((), ())) + + # FP32 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + 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_dot.multipass_dot_general( + 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()