bitnet-cpu

Ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs, loadable through kernels. The CUDA member of this stack is bitnet-tc; both share the same 2-bit packing (compatible with Microsoft's bitnet.cpp) and the same Python API, so code written against one runs against the other. The reference baseline is an fp32 reference product, matched within bf16 output rounding.

A ternary-weight model is small enough to fit a single-board computer, and useless there if the matmul runs at framework speed: a 2B model's LM head alone takes 0.7 seconds per token in fp32 on a Raspberry Pi. This kernel multiplies the packed 2-bit weights directly against per-token INT8 activations, picking the widest instruction the CPU actually has at load time, so the same binary runs a 27B-geometry layer in 0.4 ms on a Pi 5's Cortex-A76 and still works correctly on a Pi 4, an old x86, or anything else.

Measured on the boards: decode layers at 72x to 98x the fp32 path with weights down from 1,252 MB to 79 MB on the LM head, and the runtime dispatch ladder from the same binary, 2.06 ms scalar, 0.58 NEON, 0.42 SDOT on a Pi 5's Cortex-A76, against 14.3 / 3.62 / 3.95 on a Pi 4's A72, which has no dotprod and correctly stays on base NEON.

Usage

import torch
from kernels import get_kernel

bitnet = get_kernel("phanerozoic/bitnet-cpu", version=1, trust_remote_code=True)

N, K = 4096, 4096
W = torch.randint(-1, 2, (N, K), dtype=torch.int8)
w_packed = bitnet.pack_weights(W)                       # [N, K//4] uint8
scale_wt = torch.ones(N, dtype=torch.bfloat16)

x = torch.randn(1, K, dtype=torch.bfloat16)
y = bitnet.bitnet_linear(x, w_packed, scale_wt)         # [1, N] bf16

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. BitLinear is the drop-in nn.Module; BitLinearKernel is the kernelize layer for the transformers BitNet modules.

API

Weight encoding: ternary {-1, 0, +1} -> 2-bit codes {1, 2, 3} packed four per byte (decode is byte - 2), the packing used by bitnet-tc and Microsoft's bitnet.cpp.

Method

On x86-64 the inner product uses dot(w, a) = dot(w + 1, a) - sum(a), mapping onto unsigned x signed multiply-accumulate; on aarch64 the packed codes multiply directly as signed int8 via dot(w, a) = dot(w + 2, a) - 2 * sum(a), and the decode path (M <= 16) reads weights straight from the packed bytes with no unpack buffer.

The path is chosen once at runtime from cpuid / HWCAP; no flags needed. BITNET_CPU_ISA (scalar, neon, avx2, ...) demotes the selection for A/B runs and never promotes.

Measured

Raspberry Pi 5 (4x Cortex-A76 2.4 GHz, SDOT path), 27B-class layer geometry, against torch fp32 on the same board:

Dispatch ladder, same binary, 1 x 6912 x 2560:

The A76 takes SDOT; the A72 has no dotprod, so auto correctly selects base NEON and the SDOT row is unavailable to it.

On a 16 vCPU x86-64 host (AVX2, no VNNI), torch 2.12 CPU, median of 15:

End-to-end microsoft/bitnet-b1.58-2B-4T-bf16 through transformers with all 210 BitLinear layers routed to the kernel: decode 0.78 -> 3.86 tok/s (5.0x) on the x86 host, 0.035 -> 0.94 tok/s (27x) on the Pi 5, with identical greedy output over the measured window.

Correctness

The integer path is exact; output differs from an fp32 reference only by bf16 output rounding (max relative error ~1e-2 across dispatch paths on the Pi measurements above, ~4e-3 on the x86 suite). Every tier produces the same result within that bound, verified by demoting through BITNET_CPU_ISA on both a Cortex-A76 and a Cortex-A72 board.

Requirements and limits

• K divisible by 32; bf16 or f32 activations; bf16 output; uint8 packed weights.
• Fast paths cover x86-64 (AVX2 or newer) and aarch64 (NEON; SDOT where the CPU has dotprod). Anything else uses the scalar fallback, which is correct and roughly 5x slower than NEON.
• Prefill (large M) narrows to about 5x; the win is concentrated in decode.
• Microsoft's bitnet.cpp remains faster on CPU via lookup-table kernels and is the right choice when a dedicated llama.cpp-style runtime is acceptable; this kernel is for get_kernel and the transformers ecosystem.

References

Ma et al., "The Era of 1-bit LLMs" (BitNet b1.58, 2024); Microsoft bitnet.cpp (the shared packing); Arm sdot and NEON integer dot products.

License

Apache-2.0.