~/blogternary-weights-tile-contract.md
cchu@nycu:~/blog$ cat ternary-weights-tile-contract.md
2026.10.146 min[language-models][quantization][cuda][testing]

Ternary Weights, Sixteen Rows at a Time

A CPU replay of BitNet's GPU packing stages shows how tile alignment, bit interleaving, and kernel dispatch become one inference contract.

I packed a synthetic 16×33 ternary matrix through BitNet's NumPy helpers. They retained 512 of its 528 values. The final shape calculation still accepted the packed buffer. The missing values were the last column, one per row.

That result sits outside the model shapes this GPU path targets. It is useful when adapting the implementation: the weight format carries a tile contract, and a plausible final tensor shape does not prove that conversion preserved every input weight.

I examined revision 0b341e582afbf9e1011f24744b554c96a3477eb5 of BitNet. The repository contains CPU C++ kernels, generated lookup-table kernels, a model-conversion pipeline, and a separate CUDA implementation. The engineering value is visible in the GPU decoder: its storage layout lets one packed word supply byte lanes for integer dot products. I tested the packing stages on Linux with Python 3.12.14 and NumPy 2.3.5. I did not run CUDA, load a checkpoint, or measure inference speed.

A 16 by 32 tile becomes 128 packed bytes through coordinate permutation, two-bit packing, and word interleaving. The CUDA consumer expects this layout.

The Stored Format Has Two Bits per Weight

The checkpoint converter rounds and clamps the selected linear-layer weights to -1, 0, or 1. It stores scales separately. The GPU packing wrapper adds two before conversion, so those weights become codes 1, 2, and 3.

The three possible values explain the 1.58-bit terminology: log2(3) is about 1.585. This implementation stores each code in a two-bit field. Four codes occupy one byte; sixteen occupy a 32-bit word. The remaining code, zero, decodes to -2 in the GPU helper. The converter's ternary clamp avoids producing it.

That is a representation detail, not a whole-model memory estimate. Scales, embeddings, caches, activations, and any tensors left in another format still consume memory. My byte-ratio checks compare synthetic int8 weight arrays with their packed weight buffers only.

Permutation Before Packing

The packing module divides the matrix into tiles with 16 rows and 32 columns. For a coordinate (i, j) inside one tile, it computes:

thread_id = i * 2 + j // 16
row = (thread_id // 16) * 8 + (thread_id % 8)
col = (j % 16) + 16 * ((thread_id % 16) // 8)

I enumerated all 512 coordinates. The mapping produced 512 distinct source coordinates, covering the complete tile. That establishes a bijection for one aligned tile, rather than assuming that a reshaped array is a permutation.

The next helper packs four codes into each byte. The interleaver views those bytes as 32-bit words and rearranges their two-bit fields. For original element position offset, the destination bit shift is:

shift = (offset % 4) * 8 + (offset // 4) * 2

Positions 0, 1, 2, and 3 therefore occupy the low two bits of four separate byte lanes. The CUDA decoder shifts a word by 2 * group, selects lanes with 0x03030303, and subtracts two from each byte. Its kernel then uses __dp4a to accumulate four int8 products at a time. I inspected that consumer code; the CPU probe uses an independent scalar decoder for its bit-selection arithmetic and does not execute the CUDA instructions.

The layout ties conversion to the consumer's access pattern. A buffer with the right number of bits still needs the right coordinate and field ordering.

Thirty-Two Round Trips

The downloadable probe checks the exact source blob before executing four unchanged function definitions from pack_weight.py. It selects those definitions through Python's AST and supplies NumPy. It does not import the module's PyTorch dependency or call its Torch conversion wrapper.

I ran it against the pinned source file retrieved for inspection:

python public/blog/ternary-weights-tile-contract/packing-probe.py /tmp/bitnet-note/pack_weight.py

The process exited with code zero:

{"python": "3.12.14", "numpy": "2.3.5", "byteorder": "little"}
{"tileCoordinates": 512, "uniqueCoordinates": 512}
{"roundTripCases": 32, "reconstructionErrors": 0, "int8ToPackedByteRatio": 4}

The fixtures combine row counts 16, 32, 48, and 64 with column counts 32, 64, 96, and 128. Each shape receives an all-zero matrix and a seeded random ternary matrix. The reference decoder extracts the fields, subtracts two, and restores the original coordinates. Every reconstructed array matched its input exactly. Every packed buffer used one quarter of the original int8 array's bytes.

This checks codec correctness for those fixtures on a little-endian host. It does not check GPU addressing, activation quantization, floating-point scales, model quality, or throughput.

The Tail Can Disappear Before Reshape

The permutation allocates its tile grid using N // 16 and K // 32. There is no remainder tile in these helpers. The probe exercised three misaligned shapes:

Input shapeInput valuesValues retained by tile selectionValues droppedFinal NumPy reshape
16×3352851216Accepted
17×3254451232Rejected
16×314960496Rejected

For the last column in the 16×33 case, integer division reduces both the tile count and the packed output width. The retained 512 values become 128 bytes. Reshaping those bytes to (16, 33 // 4) asks for exactly 128 elements, so the arithmetic agrees despite the dropped column. The probe also decoded that buffer and confirmed it matched the first 32 columns of the input.

The reshape column in this table is explicitly a NumPy replay of the wrapper's shape calculation. PyTorch was unavailable in this environment. The original wrapper calls torch.reshape; I did not execute that call or demonstrate an end-to-end checkpoint failure. The observed loss happens earlier, inside the unchanged NumPy tile selector.

The appropriate admission check for this representation is N % 16 == 0 and K % 32 == 0 before selecting tiles. A reshape check alone cannot enforce it. That condition would reject the synthetic boundary fixtures; I left the upstream code unchanged.

Dispatch Adds Another Shape Contract

Alignment is only the packing condition. The CUDA dispatcher launches specialized kernels for ten explicit (M, N, K) combinations, all with M == 1. Its unmatched branch prints the required dimensions and launches no kernel.

The Python caller allocates a zero-filled output, calls that function through ctypes, and returns the buffer without checking a returned status. Those are source-inspection findings. I did not call the dispatcher with an unsupported shape on a GPU. They show why adapting this path to a different batch or model width requires dispatch validation in addition to packing validation.

To repeat the CPU checks from a checkout:

git clone https://github.com/microsoft/BitNet BitNet
git -C BitNet checkout 0b341e582afbf9e1011f24744b554c96a3477eb5
python packing-probe.py ./BitNet

The probe needs NumPy and finishes with:

PASS: 512-coordinate bijection, 32 round trips, and three shape-boundary fixtures

For this revision, aligned synthetic weights survive the packing and reference-decoding path exactly. Misaligned inputs can lose values before the final shape check, and aligned inputs still need a supported kernel dispatch. Those are two concrete contracts to preserve when changing the model dimensions or reusing this weight format.