kv-quant-attn

Single-token decode attention over an int8 or int4 KV cache, loadable through kernels. The reference baselines are fp32 attention on the dequantized cache (kernel residual at the bf16 rounding floor) and unquantized fp32 attention (quantization error 1.0e-2 at int8, cosine 0.99995). CPU sibling: cpu-attn.

On a device small enough to want a compressed model, the cache and not the weights is what bounds long context: a 27B-class model with 16 full-attention layers, 4 KV heads, and head dimension 256 spends 64 KiB per token in bf16, so its own 262,144-token context is 16 GiB of cache on top of a few GB of weights. Storing the cache at int8 halves that and int4 quarters it, and past a few thousand tokens reading half or a quarter of the bytes also makes decode faster than bf16 attention.

The context odometer sweeps to 262,144 tokens: the bf16 cache crosses an 8 GB board's budget while int4 idles at a quarter of it, and the measured decode-latency curves cross near 4,096 tokens, reaching 2.78 ms (bf16 SDPA) against 0.85 ms (int8 kernel) at 65,536.

Usage

import torch
from kernels import get_kernel

kv = get_kernel("phanerozoic/kv-quant-attn", version=1, trust_remote_code=True)

cache = kv.QuantizedKVCache(batch=1, heads_kv=4, head_dim=256,
                            max_tokens=262144, bits=8)

for step in range(n):
    k, v = project_kv(hidden)          # [B, Hkv, D] bf16
    cache.append(k, v)
    out = cache.attend(q)              # q is [B, Hq, D] bf16 -> [B, Hq, D]

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. One cache object per layer.

API

Method

Every key and value row carries its own absmax scale, computed when the row is written, so dequantization needs one scalar per row, appending never rewrites earlier rows, and no calibration pass is required. Decode is flash-decoding over the quantized cache: a block owns a (batch, KV head) pair, so under grouped-query attention all query heads sharing a KV head are resolved from one read of it; warps carry private (max, sumexp, accumulator) state with no barrier or atomic in the loop; and the sequence is split across blocks with partials merged in a fixed order, so the result is bitwise reproducible.

Measured

Bonsai 27B full-attention geometry (Hq = 24, Hkv = 4, D = 256, batch 1), against F.scaled_dot_product_attention on an unquantized bf16 cache of the same contents:

The crossover is near 4,000 tokens; below it the split-and-combine overhead exceeds the read it saves.

Cache footprint for the same model, scales included:

Correctness

• Against fp32 attention on the dequantized cache, which separates kernel correctness from quantization error, the residual is 1.6e-3 to 1.7e-3 relative across both bit widths and every shape tested (bf16 output rounding).
• Quantization error against unquantized fp32 attention is 1.0e-2 relative at int8 (cosine 0.99995) and 1.8e-1 at int4 (cosine 0.984); int4 is a real approximation.
• Grouped-query ratios 1:1 through 8:1 match a replicated-head reference; ragged lengths verified down to one token; appends round-trip within half a quantization step.
• Deterministic: no atomics and a fixed combine order, so repeated decodes are bitwise identical.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• head_dim at most 256 and a multiple of 32; grouped-query ratio at most 8; bf16 queries, keys and values.
• Decode only; prefill wants a tiled attention kernel, then hand the results to append.
• Slower below roughly 4,000 tokens; route short contexts to a dense path.
• Per-row absmax only: no group-wise scaling or outlier channels, which would narrow the int4 gap at some metadata cost.

References

Dao et al., "FlashAttention-2" (2023) and the flash-decoding split; Hooper et al., "KVQuant" (2024); Liu et al., "KIVI" (2024); Ainslie et al., "GQA" (2023).

License

Apache-2.0.