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 and quant-matmul (linears), 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

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

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.

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:

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.