nx_kv_cache.nx
buildroot/runtime/nx_kv_cache.nx
about
nx_kv_cache.nx -- autoregressive attention KV cache.
THE foundation primitive of efficient autoregressive generation
(LLM serving, forever-generating world models, autoregressive
diffusion). Without it, each new token recomputes attention over
every past token: O(n^2) per step. With it: O(n) per step.
For n=1000 generation steps that's a 1000x speedup with bit-
identical output (algorithmic correctness, not approximation).
Per the min-hardware-floor + algo-led cardinal: this brick is
what makes a $50 SBC run autoregressive generation that a naive
O(n^2) recompute couldn't finish in a year.
Composes with:
* nx_sparse_tensor -- KV cache rows are mostly zero post-softmax;
store sparse, compute sparse
* nx_quant_block -- cached K and V can be 4-bit; 10x less
memory bandwidth at each new-token step
* nx_compute_graph -- KV cache is a typed CONTROL node in the
graph; runner persists state across runs of the same graph
Storage layout:
k_storage[head * max_seq_len * head_dim + pos * head_dim + d]
v_storage[head * max_seq_len * head_dim + pos * head_dim + d]
Multi-head attention works by treating each head as an independent
per-position dot product space. We don't materialise the [Q, K, V]
triple as a single tensor here -- the cache OWNS K and V, the
caller's Q changes every step.
genealogy_id: vaswani_2017_transformer_kv + flashattention_2_kv +
pope_2023_efficient_serving + mistral_sliding_window
lineage_id: substrate_kv_cache_v1
dependencies 2 imports · 1 importers
imports: nx_syscalls.nxnx_tier.nx
imported by: nx_kv_cache_test.nx
structs
| 63 | struct NxKvCache |
consts
| 43 | const NX_MAGIC_1024: i64 = 1024 |
| 47 | const NX_KV_OK: nx_int = 0 |
| 48 | const NX_KV_ERR_FULL: nx_int = 1 // append past max_seq_len |
| 49 | const NX_KV_ERR_BAD_POSITION: nx_int = 2 // read past current length |
| 50 | const NX_KV_ERR_BAD_HEAD: nx_int = 3 // head index out of range |
| 51 | const NX_KV_ERR_BAD_DIM: nx_int = 4 // dim index out of range |
| 52 | const NX_KV_ERR_SHAPE_MISMATCH: nx_int = 5 // input wrong size |
| 53 | const NX_KV_N_VERDICTS: nx_int = 6 |
| 72 | const NX_KV_BYTES: nx_int = 48 // 6 fields * 8 |
| 198 | const NX_KV_Q10: nx_int = 1024 |
functions
| 55 | func nx_kv_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 76 | func nx_kv_alloc(n_heads: nx_int, head_dim: nx_int, max_seq_len: nx_int) -> *NxKvCache |
| 97 | func nx_kv_clear(kv: *NxKvCache) -> nx_int called by 1: main |
| 110 | func nx_kv_append(kv: *NxKvCache, k_in: *i64, v_in: *i64) -> nx_int called by 1: main |
| 131 | func nx_kv_get_k(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int called by 1: main |
| 142 | func nx_kv_get_v(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int called by 1: main |
| 164 | func nx_kv_score(kv: *NxKvCache, q_in: *i64, scores_out: *i64) -> nx_int called by 1: main |
| 200 | func nx_kv_apply(kv: *NxKvCache, weights_in: *i64, out: *i64) -> nx_int called by 1: main |
| 235 | func nx_kv_fill_ratio_q10(kv: *NxKvCache) -> nx_int called by 1: main |
| 246 | func nx_kv_bytes_per_token(kv: *NxKvCache) -> nx_int called by 1: main |