library_name: kernels
license: apache-2.0

binary-gemm

W1A8 matrix products for binary-weight language models, loadable through kernels. Each weight is one bit meaning +1 or -1, with one fp16 scale per group of G weights: at G = 128 that is 1.125 effective bits per weight, against 2 bits for a ternary layer and 16 for bf16. The reference baseline is the dequantized product; the integer path is torch.equal-exact where bf16 can represent the accumulator. Ternary siblings: bitnet-tc (CUDA) and bitnet-cpu (CPU).

At one bit per weight a model's memory is its scales plus raw sign bits, and decode, which reads every weight per token, becomes almost pure bit traffic. This kernel multiplies the packed bits directly: a 16-entry constant table expands a nibble of the pattern into four packed ±1 int8 values feeding __dp4a, so the weight is never materialized as a number and the decode wall drops with the bytes.

Actual packed sign bits of a layer (cyan +1, magenta -1), and decode measured with weights rotated past L2 in both paths: 2.4x on attention QKV, 2.8x on the MLP gate, 3.8x on the 248,320-row LM head.

Usage

import torch
from kernels import get_kernel

bg = get_kernel("phanerozoic/binary-gemm", version=1, trust_remote_code=True)

W = torch.where(torch.randn(N, K, device="cuda") >= 0, 1, -1).to(torch.int8)
wq = bg.pack_weights(W)                       # [N, K//32] int32
ws = group_scales.to(torch.float16)           # [N, K//128] fp16

y = bg.binary_linear(x, wq, ws, group_size=128)   # bf16 in, bf16 out

layer = bg.BinaryLinear.from_dense(nn.Linear(K, N))   # sign + group mean-abs scale

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Binary weights collapse the inner product: writing the stored bit as b = (w + 1) / 2,

dot(w, a) = 2 * sum_{i : b_i = 1} a_i  -  sum_i a_i

so the weight never has to exist as a number. The kernel keeps the int8 activation path and resolves four weights at a time through a 16-entry constant table that expands a nibble directly into four packed int8 ±1 values for __dp4a; the table is the whole decode. One warp owns each output column: lanes stride the weight groups, accumulate an int32 dot per group, apply that group's fp16 scale in fp32, and the warp reduces. The batch dimension is tiled at compile time so a weight word is fetched once per tile.

Measured

Bonsai 27B shapes, hidden 5120, G = 128, against cuBLAS bf16, weights rotated over enough copies to exceed twice the card's L2 on both paths (an 84 MB bf16 matrix inside a 96 MB L2 otherwise measures cache bandwidth):

Weight footprint at hidden 5120:

1.89x rather than the nominal 1.78x because the fp16 group scales are a fixed overhead both formats pay.

Correctness

• Packing is lossless: unpack(pack(W)) is torch.equal to W for every shape tested. Zeros are treated as -1, so a ternary tensor deliberately does not round-trip.
• The integer path is exact: with unit group scales the output is torch.equal to an exact int32 reference wherever bf16 represents the accumulator (|acc| < 256).
• The full path matches a dequantize-and-fp32 reference to below 2^-8 normwise, the bf16 output rounding step.
• Activation quantization is per-token absmax with reconstruction error under 1% of the row maximum; group size verified at 32, 64, 128, 256.
• Deterministic: repeated products are bitwise identical; fixed reduction tree, no atomics.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+ (__dp4a).
• K a multiple of 32 and of group_size; bf16 activations and output.
• Decode only: a warp-per-column GEMV built for M = 1. Above that it loses to cuBLAS (0.40x at M = 2, 0.02x at M = 256); route prefill and batch elsewhere.
• Weights must already be binary; from_dense is a sign quantizer for testing, not a compression method.

References

Rastegari et al., "XNOR-Net" (2016); Courbariaux et al., "BinaryConnect" (2015); Ma et al., "The Era of 1-bit LLMs" (2024); group-wise scaling with 1-bit weights as shipped in the Bonsai family (PrismML, 2026).

License

Apache-2.0.