google / google/qwix

int4 QArray stores at 1.0 bytes/element (no nibble-packing) — zero memory benefit

Open
#328 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.