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. The inner product uses the identity dot(w, a) = dot(w + 1, a) - sum(a), which maps ternary GEMM directly onto unsigned x signed multiply-accumulate instructions:

The path is chosen once at runtime from cpuid; no flags needed.

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.

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 need x86-64 (AVX2 or newer); other architectures use 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.