From f8437ef15e8977d81e8f91734983903e5639cbe3 Mon Sep 17 00:00:00 2001 From: James Chapman Date: Mon, 29 Jun 2026 16:03:08 -0700 Subject: [PATCH] Add fused Hadamard quantization Pallas kernel. PiperOrigin-RevId: 940083700 --- .../kernels/fused_hadamard_quantize.py | 182 ++++++++++++++++++ .../kernels/fused_hadamard_quantize_test.py | 54 ++++++ 2 files changed, 236 insertions(+) create mode 100644 qwix/contrib/kernels/fused_hadamard_quantize.py create mode 100644 tests/contrib/kernels/fused_hadamard_quantize_test.py diff --git a/qwix/contrib/kernels/fused_hadamard_quantize.py b/qwix/contrib/kernels/fused_hadamard_quantize.py new file mode 100644 index 0000000..e2a20d6 --- /dev/null +++ b/qwix/contrib/kernels/fused_hadamard_quantize.py @@ -0,0 +1,182 @@ +# 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. +"""Fused Hadamard Quantization. + +This kernel fuses together the Hadamard transform and the quantization step. + +Currently supports doing a Hadamard transform on the final axis and quantizes +with 2d scales. Does not use random keys. For more features, refer to the jax +implementation in qwix. Supports int8 quantization. +""" + +import functools + +import jax +import jax.experimental.pallas as pl +import jax.experimental.pallas.tpu as pltpu +import jax.numpy as jnp +from qwix.contrib import hadamard_rot + +_CORE_AXIS_NAME = "core" + + +def quantize_a_tile(x, dtype=jnp.int8): + """Quantizes a tile of LHS values.""" + s = jnp.max(jnp.abs(x), keepdims=True, axis=(-2, -1)) / ( + jnp.iinfo(dtype).max + 0.5 + ) + x = jnp.round(x / s).astype(dtype) + return x, s + + +def fused_hadamard_quantize_kernel( + x_hbm: jax.Ref, + had_mat_hbm: jax.Ref, + xq_hbm: jax.Ref, + s_hbm: jax.Ref, + *, + bm: int, + bn: int, + sm_global: int, + sn_global: int, +): + """Fused Hadamard Quantize Kernel. + + Args: + x_hbm: Reference to the input array in hbm (m, k) + had_mat_hbm: Reference to the Hadamard matrix in hbm (had_size, had_size) + xq_hbm: Reference to the output array in hbm (m, k) + s_hbm: Reference to the output scale array in hbm (sm, sn, 1, 1) + bm: blockspec for the m dimension + bn: blockspec for the n dimension + sm_global: Scale shape for the m dimension + sn_global: Scale shape for the n dimension + """ + # Grid + m, n = x_hbm.shape + grid = (pl.cdiv(m, bm), pl.cdiv(n, bn)) + + # Scale shapes within the kernel + sm = pl.cdiv(sm_global, grid[0]) + sn = pl.cdiv(sn_global, grid[1]) + + # Blockspecs for the kernel + x_spec = pl.BlockSpec((bm, bn), lambda a, b: (a, b)) + had_mat_spec = pl.BlockSpec(had_mat_hbm.shape, lambda a, b: (0, 0)) + xq_spec = pl.BlockSpec((bm, bn), lambda a, b: (a, b)) + s_spec = pl.BlockSpec((sm, sn, 1, 1), lambda a, b: (a, b, 0, 0)) + + # Tile sizes corresponding to scale entries + m_tile_size = pl.cdiv(bm, sm) + n_tile_size = pl.cdiv(bn, sn) + + # Kernel body + def kernel_body( + x_vmem: jax.Ref, + had_mat_vmem: jax.Ref, + xq_vmem: jax.Ref, + s_vmem: jax.Ref, + ): + # Initialize accumulation buffer + had_mat = had_mat_vmem[...] + + for nloop in range(sn): + data_n_slc = pl.Slice(nloop * n_tile_size, n_tile_size) + x = x_vmem[:, data_n_slc] + xrot = jnp.matmul(x, had_mat) + + # Quantize + xq, sx = quantize_a_tile(xrot.reshape(-1, m_tile_size, n_tile_size)) + + # Write results + xq_vmem[:, data_n_slc] = xq.reshape(-1, n_tile_size) + s_vmem[:, nloop] = sx.reshape(-1, 1, 1) + + # Call the kernel. + pltpu.emit_pipeline( + kernel_body, + grid=grid, + in_specs=[x_spec, had_mat_spec], + out_specs=[xq_spec, s_spec], + core_axis_name=_CORE_AXIS_NAME, + dimension_semantics=(pltpu.PARALLEL, pltpu.PARALLEL), + )(x_hbm, had_mat_hbm, xq_hbm, s_hbm) + + +def fused_hadamard_quantize( + x: jax.Array, + *, + bm: int, + bn: int, + sm: int, + sn: int, + dtype: jax.typing.DTypeLike = jnp.int8, + hadamard_power: int, +): + """Fused Hadamard Quantization. + + Args: + x: Input array (m, k) + bm: blockspec for the m dimension + bn: blockspec for the n dimension + sm: Scale shape for the m dimension + sn: Scale shape for the n dimension + dtype: Output quantized dtype + hadamard_power: Power of 2 for the Hadamard matrix + + Returns: + A tuple of the quantized array and the scales. + """ + # Create the tensor core mesh + tc_mesh = pltpu.create_tensorcore_mesh(axis_name=_CORE_AXIS_NAME) + + # Create the output type + xq_out_type = jax.core.ShapedArray(x.shape, dtype) + s_out_type = jax.core.ShapedArray((sm, sn, 1, 1), x.dtype) + + had_size = int(2**hadamard_power) + had_mat = hadamard_rot._create_hadamard_matrix( # pylint: disable=protected-access + hadamard_power, + None, + row_sign_flip=False, + col_sign_flip=False, + dtype=jnp.int8, + )[ + 0 + ] + n_tile_size = x.shape[1] // sn + if dtype != jnp.int8: + raise ValueError("Only int8 quantization is supported.") + if had_size != n_tile_size: + raise ValueError( + "Hadamard matrix size must match the tile size of the unquantized" + " matrix." + ) + + # Create the kernel with kwargs + kernel = functools.partial( + fused_hadamard_quantize_kernel, + bm=bm, + bn=bn, + sm_global=sm, + sn_global=sn, + ) + + # Call the kernel + xq, s = pl.kernel( + kernel, + out_type=[xq_out_type, s_out_type], + mesh=tc_mesh, + )(x, had_mat) + return xq, s.reshape(sm, sn) diff --git a/tests/contrib/kernels/fused_hadamard_quantize_test.py b/tests/contrib/kernels/fused_hadamard_quantize_test.py new file mode 100644 index 0000000..59b69b4 --- /dev/null +++ b/tests/contrib/kernels/fused_hadamard_quantize_test.py @@ -0,0 +1,54 @@ +from absl.testing import absltest +from absl.testing import parameterized +import jax +import jax.numpy as jnp +import numpy as np +from qwix.contrib import hadamard_rot +from qwix.contrib.kernels import fused_hadamard_quantize as fhq + + +class LHSFusedHadamardQuantizeTest(parameterized.TestCase): + + @parameterized.parameters( + (1024, 1024, 8, 8, 1024, 1024), + ) + def test_kernel_dot_general(self, m, n, sm, sn, bm, bn): + hadamard_power = 7 + key = jax.random.PRNGKey(0) + lhs = jax.random.normal(key, (m, n), dtype=jnp.float32) + + kernel_answer_xq, kernel_answer_sx = fhq.fused_hadamard_quantize( + lhs, + bm=bm, + bn=bn, + sm=sm, + sn=sn, + hadamard_power=hadamard_power, + ) + + had_mat = hadamard_rot._create_hadamard_matrix( # pylint: disable=protected-access + hadamard_power, + None, + row_sign_flip=False, + col_sign_flip=False, + dtype=jnp.int8, + )[ + 0 + ] + # Apply Hadamard transform to the LHS + lhs_jax = jnp.matmul(lhs.reshape(sm, m // sm, sn, n // sn), had_mat) + + # Quantize the LHS + vqfn = jax.vmap(fhq.quantize_a_tile, in_axes=0, out_axes=(0, 0)) + vvqfn = jax.vmap(vqfn, in_axes=2, out_axes=(2, 2)) + xq, sx = vvqfn(lhs_jax) + xq = xq.reshape(m, n) + sx = sx.reshape(sm, sn) + + # Check numerics + np.testing.assert_allclose(kernel_answer_xq, xq, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(kernel_answer_sx, sx, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + absltest.main()