cpu-attn

Single-token decode attention over an INT8-quantized KV cache, aarch64-first. 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. Versus an fp16 cache this halves both cache RAM and per-token cache bandwidth, which is the decode wall on LPDDR4-class boards. Companions on the Hub: phanerozoic/bitnet-cpu and phanerozoic/quant-matmul (linears), phanerozoic/decode-ops (fused glue and sampling).

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.

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

GQA follows from shapes: query head h reads kv head h // (H // Hkv).

API

Performance

Raspberry Pi 5 (4x Cortex-A76 2.4 GHz, SDOT path) and Raspberry Pi 4 Model B (4x Cortex-A72 1.8 GHz, NEON path), 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.

Requirements

• Head dim a multiple of 16, at most 1024; f32 queries; 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.

Scope

The KV-cache half of a quantized CPU inference stack: quantized linears (bitnet-cpu, quant-matmul) make weights cheap to stream, this kernel does the same for the cache, and long-context decode stops being the first thing that dies on an 8 GB board.

License

Apache-2.0.