bitnet-cpu

Ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) on CPUs. This is the CPU member of the BitNet stack on the Hub; the CUDA member is phanerozoic/bitnet-tc. Both share the same 2-bit weight packing (compatible with Microsoft's bitnet.cpp) and the same Python API, so code written against one runs against the other.

Ternary weights {-1, 0, +1} are stored at 2 bits per value (8x smaller than BF16). Activations are quantized per-token to INT8. On x86-64 the inner product uses the identity dot(w, a) = dot(w + 1, a) - sum(a), mapping onto unsigned x signed multiply-accumulate instructions; 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; it never promotes.

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.

Drop-in nn.Module:

layer = bitnet.BitLinear(K, N)
layer.w_packed.copy_(w_packed)
layer.scale_wt.copy_(scale_wt)
y = layer(x)

API

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

Performance

Measured on a 16 vCPU x86-64 host (AVX2, no VNNI; HF Jobs cpu-xl), 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 throughput 0.78 -> 3.86 tok/s (5.0x) versus the stock forward on the same host, with identical greedy output over the measured window. AVX-VNNI and AVX-512 VNNI hosts run faster than this AVX2 baseline.

On Raspberry Pi (aarch64 variants, 64-bit Raspberry Pi OS, torch 2.13 CPU; Pi 5: 4x Cortex-A76 2.4 GHz on the SDOT path, Pi 4 Model B: 4x Cortex-A72 1.8 GHz on the NEON path):

The M=1 rows stream packed weights at 10-11 GB/s on the Pi 5; M=128 runs 295 GOPS. End-to-end microsoft/bitnet-b1.58-2B-4T-bf16 through transformers on the Pi 5 with all 210 BitLinear layers routed to the kernel: decode 0.035 -> 0.94 tok/s (27x) versus the stock forward, with identical greedy output over the measured window. The remaining step time sits outside the ternary GEMVs (bf16 attention, norms, and lm_head on a CPU with no bf16 fast path).

The integer path is exact; output differs from an f32 reference only by bf16 output rounding (max relative error ~4e-3 across all dispatch paths).

Requirements

• 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 but slow.

Scope

The purpose of this kernel is hub-loadable BitNet inference through kernels/transformers on ordinary CPUs, replacing the stock unquantized matmul path. 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 when you want get_kernel and the transformers ecosystem.

License

Apache-2.0.