| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # 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. |
|
|
|  |
|
|
| *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 |
|
|
| ```python |
| 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. |
|
|
| ## Requirements and limits |
|
|
| - NVIDIA GPU with compute capability 8.0+ (Ampere, Ada, Hopper). |
| - `K` divisible by 32; fused GEMV path requires `K <= 16384`. |
| - bf16 activations and output; int8 quantized activations; uint8 packed |
| weights. |
|
|
| ## References |
|
|
| Ma et al., "The Era of 1-bit LLMs" (BitNet b1.58, 2024); Microsoft |
| bitnet.cpp (the shared packing); NVIDIA IMMA tensor-core instructions. |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
