int4 QArray stores at 1.0 bytes/element (no nibble-packing) — zero memory benefit
- Dominant language
- Python
- Stars
- 134
- Forks
- 27
- Avg merge
- 2d 2h
- Merged PRs (30d)
- 12
Description
## 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)
```python
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_
Contributor guide
Research direction
Start at qwix.quantize_api and trace how jnp.int4 values are stored in QArray.qvalue, then inspect the round-trip path for channelwise_axes=[0]. Verify the result with qvalue.nbytes and exact dequantization: two int4 values should occupy each uint8 byte, giving 0.5 B/element without changing round-trip behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100