| --- |
| library_name: kernels |
| license: apache-2.0 |
| --- |
| |
| # cpu-attn |
|
|
| Single-token decode attention over an INT8-quantized KV cache, aarch64 |
| first, loadable through `kernels`. The reference baseline is fp32 attention |
| on the dequantized cache, matched within 8e-3. Companions: |
| [bitnet-cpu](https://huggingface.co/kernels/phanerozoic/bitnet-cpu) and |
| [quant-matmul](https://huggingface.co/kernels/phanerozoic/quant-matmul) |
| (linears), |
| [decode-ops](https://huggingface.co/kernels/phanerozoic/decode-ops) (fused |
| glue and sampling). |
|
|
| Once the linear layers of a CPU inference stack are quantized, the KV cache |
| is what remains: on an 8 GB single-board computer it is both the memory that |
| runs out and the bandwidth that sets decode latency. Storing keys and values |
| at int8 halves both, and this kernel does the decode step directly against |
| that quantized cache, with exact integer logit dots, so long-context decode |
| on a Raspberry Pi stops being the first thing that dies. |
|
|
|  |
|
|
| *Measured on the boards, 8 heads at head dimension 128: 0.10, 0.42, and |
| 1.80 ms at 512, 2,048, and 8,192 tokens on a Pi 5's Cortex-A76, against |
| 1.18, 4.30, and 10.51 ms on a Pi 4's A72, with the cache down from 32 MB to |
| 16 MB at the longest context and output within 8e-3 of fp32 attention on the |
| dequantized cache.* |
|
|
| ## Usage |
|
|
| ```python |
| import torch |
| from kernels import get_kernel |
| |
| ca = get_kernel("phanerozoic/cpu-attn", version=1, trust_remote_code=True) |
| |
| cache = ca.Int8KVCache(n_kv_heads=8, max_seq=4096, head_dim=128) |
| for k, v in kv_stream: # [Hkv, D] f32 per token |
| cache.append(k, v) |
| out = cache.decode(q) # q [H, D] -> [H, D] f32 |
| ``` |
|
|
| `version` selects the release branch; `trust_remote_code` is required by |
| `kernels` for publishers without the trusted-publisher mark. GQA follows |
| from shapes: query head `h` reads kv head `h // (H // Hkv)`. |
|
|
| ## API |
|
|
| | Symbol | Purpose | |
| |---|---| |
| | `kv_append(k_cache, k_scale, v_cache, v_scale, k_new, v_new, pos)` | quantize one token into the caches | |
| | `attn_decode(q, k_cache, k_scale, v_cache, v_scale, seq_len, scale)` | decode attention for one token | |
| | `Int8KVCache(n_kv_heads, max_seq, head_dim)` | allocating wrapper with `append` / `decode` / `dequant_k` / `dequant_v` | |
|
|
| ## Method |
|
|
| K and V quantize per token per head (absmax int8) at append time; decode |
| runs integer q.k logits, an f32 softmax, and f32 accumulation over the int8 |
| V stream: |
|
|
| ``` |
| logits[s] = sdot(q_i8, k_i8[s]) * q_scale * k_scale[s] * softmax_scale |
| out[d] = sum_s softmax(logits)[s] * v_scale[s] * v_i8[s, d] |
| ``` |
|
|
| Logit dots are exact int32, accumulation order is fixed, and the NEON |
| polynomial exp satisfies `exp(0) == 1` exactly, so uniform-logit |
| constructions are bit-exact end to end. |
|
|
| | path | instruction | selected when | |
| |---|---|---| |
| | AVX2 | 16-bit widening `vpmaddwd` + Cephes exp | x86-64 with AVX2 and FMA | |
| | SDOT | `sdot` (dotprod) | aarch64 with `asimddp`: Cortex-A76+, Neoverse, Apple silicon | |
| | NEON | `smull` + `sadalp` | any other aarch64 (Pi 4, Pi Zero 2) | |
| | scalar | portable C++ | everything else | |
|
|
| Chosen once at runtime from CPUID / HWCAP; `CA_CPU_ISA` demotes for A/B |
| runs. The int8 q.k dot is signed x signed, which AVX-VNNI does not |
| accelerate, so both x86 tiers share the AVX2 path. |
|
|
| ## Measured |
|
|
| Raspberry Pi 5 (4x Cortex-A76 2.4 GHz, SDOT) and Raspberry Pi 4 Model B (4x |
| Cortex-A72 1.8 GHz, NEON), 64-bit Raspberry Pi OS, torch 2.13 CPU, D = 128: |
|
|
| | H, Hkv, S | Pi 5 | Pi 4 | cache bf16 -> int8 | |
| |---|---|---|---| |
| | 8, 8, 512 | 0.10 ms | 1.18 ms | 2 -> 1 MB | |
| | 8, 8, 2048 | 0.42 ms | 4.30 ms | 8 -> 4 MB | |
| | 8, 8, 8192 | 1.80 ms | 10.51 ms | 32 -> 16 MB | |
| | 32, 8, 2048 | 1.12 ms | 3.88 ms | 8 -> 4 MB | |
|
|
| The Pi 5 streams KV at 9-12 GB/s, the board's memory ceiling; an fp16 cache |
| would move twice the bytes for the same context. |
|
|
| ## Correctness |
|
|
| Output matches fp32 attention on the dequantized cache within 8e-3 relative |
| across the measured contexts on both boards, which separates kernel |
| correctness from quantization error. Head dim a multiple of 16 up to 1024; |
| `H % Hkv == 0`. |
|
|
| ## Requirements and limits |
|
|
| - f32 queries; head dim a multiple of 16, at most 1024; `H % Hkv == 0`. |
| - This kernel is the decode path; prefill attention stays in torch SDPA. |
| - Fast paths cover aarch64 (NEON; SDOT with dotprod) and x86-64 (AVX2 + |
| FMA); anything else uses the correct scalar fallback. |
| - Per-token per-head absmax only, no group-wise scaling within a row. |
|
|
| ## References |
|
|
| Dao et al., "FlashAttention-2" (2023); Hooper et al., "KVQuant" (2024); |
| Ainslie et al., "GQA" (2023); Arm `sdot`. |
|
|
| ## License |
|
|
| Apache-2.0. |
|
|
