From 38799a6f9f75fe209a22a4e715657841ca537ff0 Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Wed, 8 Jul 2026 01:41:02 -0700
Subject: [PATCH] feat(int4): add nibble_pack/unpack for true 0.5 B/elem int4
 storage

Currently, qwix int4 QArrays store values at 1.0 bytes/element because
JAX's jnp.int4 byte-pads each value. This means int4 quantization pays
the quality cost of 4-bit with zero memory benefit over int8.

This adds nibble_pack() and nibble_unpack() utilities that store two
4-bit values per uint8 byte, achieving the true 0.5 B/elem (4x vs fp32,
2x vs int8) memory reduction that int4 is meant to provide.

Measured: nibble_pack gives exactly 0.5 B/elem on both GPU (GB200) and
TPU (v7x), roundtrip-exact, with the standard GPTQ/AWQ packing layout.

Addresses: https://github.com/google/qwix/issues/328
---
 qwix/_src/core/nibble_pack.py | 47 +++++++++++++++++++++++++++++++++++
 tests/nibble_pack_test.py     | 38 ++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 qwix/_src/core/nibble_pack.py
 create mode 100644 tests/nibble_pack_test.py

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()