Implement INT4 group quantization, nibble packing and a dequantizing matmul

hardPremium edge casesimplement int4_matmul

Before you code

Two signed 4-bit values occupy one byte. Predict which intermediate representation you need before multiplying them by a floating-point scale.

Review the prerequisite concept

The concept page shows its own access level. The statement and visible tests below are free.

Write the three pieces of a weight-only INT4 path.

quantize_int4(w, group=128) -> (packed, scales, K)
  w        [N, K] float32
  scales   [N, G] float16, G = ceil(K / group), scale = absmax of the group / 7
  packed   [N, ceil(K / 2)] uint8: q = round(w / scale) clipped to [-8, 7], as a
           two's-complement nibble; element 2j in the low nibble of byte j
           (an odd K pads one zero nibble; a group of zeros gets scale 0)
unpack_int4(packed, K) -> int8 [N, K], sign-extended
int4_matmul(x, packed, scales, K, group=128) -> float32 [M, N]
  y = x @ dequantize(W).T, computed one group at a time: multiply x by the group's
  integer weights first, then apply the group's scale to each output column

The visible tests cover the round-trip error and the byte layout. The full suite adds shapes with a partial last group and an odd K, a group of zeros, sign extension of -8, the matmul against a full dequantization, and the storage size.

YOUR IMPLEMENTATION · PYTHON 3.14 + NUMPY

Visible tests run for everyone; the edge cases and the solution are Premium.

The visible tests

import numpy as np


def test_round_trip_error_stays_within_half_a_step(ns):
    rng = np.random.default_rng(0)
    w = rng.normal(size=(8, 256)).astype(np.float32)
    packed, scale, K = ns["quantize_int4"](w, 128)
    q = np.asarray(ns["unpack_int4"](packed, K), dtype=np.float32)
    s = np.repeat(np.asarray(scale, np.float32), 128, axis=1)[:, :K]
    err = np.abs(q * s - w)
    assert (err <= s * 0.5 + 1e-3 * s + 1e-6).all(), f"worst error {err.max():.4f} exceeds half a quantization step"


def test_nibble_layout_is_even_low_odd_high(ns):
    w = np.array([[1.0, -1.0, 7.0, -7.0]], dtype=np.float32)    # absmax 7 -> scale 1
    packed, scale, K = ns["quantize_int4"](w, 4)
    assert float(np.asarray(scale)[0, 0]) == 1.0, f"absmax 7 over 7 is scale 1, got {np.asarray(scale)[0, 0]}"
    got = [int(b) for b in np.asarray(packed)[0]]
    assert got == [0xF1, 0x97], f"expected bytes [0xF1, 0x97] (1 low, -1 high; 7 low, -7 high), got {[hex(b) for b in got]}"

The reasoning behind this component: the companion page.