diff --git a/qwix/_src/core/nibble_pack.py b/qwix/_src/core/nibble_pack.py new file mode 100644 index 0000000..bc633fb --- /dev/null +++ b/qwix/_src/core/nibble_pack.py @@ -0,0 +1,47 @@ +"""Nibble-packing for int4 QArrays — stores two 4-bit values per uint8 byte. + +Without packing, jnp.int4 arrays consume 1.0 bytes/element (JAX/XLA byte-pads sub-byte dtypes). +This module provides pack/unpack that achieve the true 0.5 bytes/element memory reduction — +the intended benefit of int4 quantization. + +Usage: + packed = nibble_pack(q_int4_or_int8) # [N, K] -> [N, K/2] uint8 + unpacked = nibble_unpack(packed) # [N, K/2] uint8 -> [N, K] int8 in [-8, 7] +""" + +import jax.numpy as jnp + + +def nibble_pack(q): + """Pack int4/int8 values (in [-8, 7]) into uint8 (two nibbles per byte). + + Args: + q: Array with values in [-8, 7], any shape with even last dimension. + dtype can be int4, int8, or int32. + + Returns: + uint8 array with last dimension halved. Each byte stores two 4-bit values: + low nibble = q[..., 0::2], high nibble = q[..., 1::2]. + """ + n = (q.astype(jnp.int32) & 0xF).astype(jnp.uint8) + lo = n[..., 0::2] + hi = n[..., 1::2] + return (lo | (hi << 4)).astype(jnp.uint8) + + +def nibble_unpack(packed): + """Unpack uint8 nibble-packed array back to int8 values in [-8, 7]. + + Args: + packed: uint8 array from nibble_pack. + + Returns: + int8 array with last dimension doubled. Values are sign-extended 4-bit + (range [-8, 7]). + """ + lo = packed & 0xF + hi = (packed >> 4) & 0xF + # Sign-extend from 4-bit two's complement + lo = jnp.where(lo >= 8, lo.astype(jnp.int32) - 16, lo.astype(jnp.int32)).astype(jnp.int8) + hi = jnp.where(hi >= 8, hi.astype(jnp.int32) - 16, hi.astype(jnp.int32)).astype(jnp.int8) + return jnp.stack([lo, hi], axis=-1).reshape(*packed.shape[:-1], packed.shape[-1] * 2) diff --git a/tests/nibble_pack_test.py b/tests/nibble_pack_test.py new file mode 100644 index 0000000..41d7378 --- /dev/null +++ b/tests/nibble_pack_test.py @@ -0,0 +1,38 @@ +"""Tests for nibble_pack/unpack correctness and memory savings.""" + +import numpy as np +from absl.testing import absltest +import jax.numpy as jnp + +from qwix._src.core.nibble_pack import nibble_pack, nibble_unpack + + +class NibblePackTest(absltest.TestCase): + + def test_roundtrip_exact(self): + """Pack then unpack recovers original values exactly.""" + rng = np.random.default_rng(42) + q = jnp.asarray(rng.integers(-8, 8, (64, 128)), jnp.int8) + packed = nibble_pack(q) + unpacked = nibble_unpack(packed) + np.testing.assert_array_equal(np.asarray(unpacked), np.asarray(q)) + + def test_memory_reduction(self): + """Packed array is exactly half the elements (0.5 B/elem vs 1.0).""" + q = jnp.zeros((1024, 4096), jnp.int8) + packed = nibble_pack(q) + self.assertEqual(packed.shape, (1024, 2048)) + self.assertEqual(packed.dtype, jnp.uint8) + # True memory: packed is half the bytes of the original + self.assertEqual(packed.nbytes, q.nbytes // 2) + + def test_boundary_values(self): + """Correctly handles the full int4 range [-8, 7].""" + vals = jnp.asarray(list(range(-8, 8)) * 2, jnp.int8).reshape(1, 32) + packed = nibble_pack(vals) + unpacked = nibble_unpack(packed) + np.testing.assert_array_equal(np.asarray(unpacked), np.asarray(vals)) + + +if __name__ == '__main__': + absltest.main()