bitnet-tc

Tensor-core ternary x INT8 GEMM for BitNet b1.58 (W1.58 A8) models on NVIDIA Ampere and newer, loadable through kernels. The reference baseline is cuBLAS bf16 on the same shapes, and end-to-end the stock transformers forward, which it doubles.

Ternary-weight language models store each weight as {-1, 0, +1} in 2 bits, 8x smaller than bf16, but without a kernel that multiplies packed ternary weights directly, inference dequantizes them back to floating point and the size advantage never becomes speed. This kernel runs the product on the INT8 tensor cores against the packed weights themselves: autoregressive decode, where every token reads every weight, gets several times faster, and the LM head, the largest single matrix in the model, gets the most.

$One decode step's linear layers as stacked bars: the bf16 tower compresses to a fraction of its height under the ternary kernel$

The linear layers of one LLaMA-7B decode step at M = 1, each measured live: 0.76 ms of cuBLAS bf16 collapses to 0.16 ms (4.8x), the LM head alone 9.4x, with the weight footprint down from 636 MB to 80 MB.

Usage

import torch
from kernels import get_kernel

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

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

x = torch.randn(1, K, dtype=torch.bfloat16, device="cuda")
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

Symbol	Purpose
`pack_weights(W)`	ternary `{-1,0,+1}` int8 `[N,K]` -> packed uint8 `[N,K//4]`
`quantize_activation(x)`	bf16 `[..,K]` -> (`int8`, per-row bf16 scale)
`bitnet_gemm(x_int8, w_packed, scale_act, scale_wt)`	INT8 activations x packed ternary weights -> bf16
`bitnet_gemv_fused(x_bf16, w_packed, scale_wt)`	fused quantize + GEMV, M < 16, K <= 16384
`bitnet_linear(x_bf16, w_packed, scale_wt)`	one-shot forward, auto-dispatch
`BitLinear(in, out)`	`nn.Module` wrapper
`BitLinearKernel`	`kernelize` layer for transformers BitNet `BitLinear` / `AutoBitLinear`

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

Method

Compute paths dispatch by batch size:

M	path	tensor cores
1	fused quantize + GEMV (BN=32)	dp4a
2-15	fused multi-M GEMV (shared B reads)	dp4a
16-383	BM=32 split-K IMMA, INT32 atomic reduction (non-split when N fills the grid)	m16n8k32
>= 384	IMMA, 128x128 tile, 3-stage cp.async	m16n8k32

Measured

Against cuBLAS bf16 on real LLaMA-2-7B and BitNet b1.58 2B-4T layer shapes:

Shape (M, N, K)	Layer	speedup vs cuBLAS bf16
1, 32000, 4096	LLaMA-7B LM head	~10x
1, 12288, 4096	LLaMA-7B QKV proj	~6-7x
1, 11008, 4096	LLaMA-7B FFN gate	~6-7x
1, 4096, 4096	LLaMA-7B attn O proj	~2-3x
128, 11008, 4096	mid-batch, large N	~1.1x
512, 11008, 4096	LLaMA-7B FFN gate (prefill)	~1.5x
512, 4096, 4096	square prefill	parity (~1.0x)

Top measured throughput is ~150 TFLOPS at M=512 on the large-N shapes, on top of the 8x weight-memory reduction. End-to-end through transformers (microsoft/bitnet-b1.58-2B-4T, all 210 BitLinear layers routed to the kernel): decode 7.5 -> 15.3 tok/s (2.05x) versus the stock bf16 online-quant path, ~2.3x on an L4; wikitext-2 perplexity +0.086% vs stock.

Small-N mid-batch (M=64-256, N=4096) is the remaining weak corner at ~0.6-0.75x of cuBLAS, where the problem is too small to fill the GPU at INT8 throughput.