NF4 Quantization

The problem, from first principles

Uniform 4-bit quantization spaces its 16 levels evenly, but trained weights are not spread evenly: they follow a bell curve, with most values near zero and few near the extremes. Evenly spaced levels waste resolution on the sparse tails and give too little to the crowded middle. NF4 (NormalFloat 4) places its 16 levels at the quantiles of a standard normal distribution, so each level is used about equally often. Weights are processed in blocks: each block is divided by its own absolute maximum so it lies in [-1, 1], snapped to the nearest NF4 level, and stored as a 4-bit code plus one scale per block.

From theory to code

Implement nf4_quantize and nf4_dequantize.

Constraints

• Use the fixed 16-entry table NF4_LEVELS provided in the starter, sorted ascending from -1.0 to 1.0 with an exact 0.0.
• nf4_quantize(w, block_size): w is a 1-D float array whose length is divisible by block_size. For each consecutive block compute scale = max(|block|), normalize block / scale (use scale 1.0 for an all-zero block), and assign each value the index of the nearest NF4 level (lowest index on ties). Return (codes, scales) with codes an integer array of the same length as w and scales of length len(w) / block_size.
• nf4_dequantize(codes, scales, block_size) returns NF4_LEVELS[codes] * scale_of_its_block.

Hints