Summary
qwix.quantize_api(array, jnp.int4, channelwise_axes=[0]) produces a QArray with qvalue.dtype=int4 but qvalue.nbytes = N*K (1.0 B/elem) — identical to int8. JAX's jnp.int4 pads each 4-bit value to a full byte in HBM. So int4 quantization pays the quality cost of 4-bit with zero memory benefit.
Measured (jax 0.10.2, GB200 + TPU v7x)
qa = qwix.quantize_api(W, jnp.int4, channelwise_axes=[0])
print(qa.qvalue.nbytes / (N*K)) # 1.0 — same as int8!
Expected: 0.5 B/elem (two int4 values packed per uint8 byte = the real 4× memory win).
The fix (we implemented externally)
Manual nibble-packing into uint8: lo = q[...,0::2] & 0xF; hi = q[...,1::2] << 4; packed = lo | hi → true 0.5 B/elem. Roundtrip-exact. Used in production on our GB200+TPU benchmarks.
Impact
Without packing, qwix int4 is strictly dominated by int8 (same memory, worse quality). This makes the int4 qtype effectively unusable for its intended purpose (weight compression). The fix is the standard GPTQ/AWQ nibble-packing layout.
— via Navi on behalf of @lokic233
Summary
qwix.quantize_api(array, jnp.int4, channelwise_axes=[0])produces a QArray withqvalue.dtype=int4butqvalue.nbytes = N*K(1.0 B/elem) — identical to int8. JAX'sjnp.int4pads each 4-bit value to a full byte in HBM. So int4 quantization pays the quality cost of 4-bit with zero memory benefit.Measured (jax 0.10.2, GB200 + TPU v7x)
Expected: 0.5 B/elem (two int4 values packed per uint8 byte = the real 4× memory win).
The fix (we implemented externally)
Manual nibble-packing into uint8:
lo = q[...,0::2] & 0xF; hi = q[...,1::2] << 4; packed = lo | hi→ true 0.5 B/elem. Roundtrip-exact. Used in production on our GB200+TPU benchmarks.Impact
Without packing, qwix int4 is strictly dominated by int8 (same memory, worse quality). This makes the int4 qtype effectively unusable for its intended purpose (weight compression). The fix is the standard GPTQ/AWQ nibble-packing layout.
— via Navi on behalf of @lokic233