code wiki / (root) / nx_kv_cache.nx

nx_kv_cache.nx

buildroot/runtime/nx_kv_cache.nx

8934 B248 linesdepth 2pulls 2 transitivereach 1 importersview sourcekind librarytopic kv
docsdependenciesstructsconstsfunctions

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

nx_syscalls.nx nx_tier.nx nx_kv_cache.nx nx_kv_cache_test.nx

imports: nx_syscalls.nxnx_tier.nx

imported by: nx_kv_cache_test.nx

structs

63struct NxKvCache

consts

43const NX_MAGIC_1024: i64 = 1024
47const NX_KV_OK: nx_int = 0
48const NX_KV_ERR_FULL: nx_int = 1 // append past max_seq_len
49const NX_KV_ERR_BAD_POSITION: nx_int = 2 // read past current length
50const NX_KV_ERR_BAD_HEAD: nx_int = 3 // head index out of range
51const NX_KV_ERR_BAD_DIM: nx_int = 4 // dim index out of range
52const NX_KV_ERR_SHAPE_MISMATCH: nx_int = 5 // input wrong size
53const NX_KV_N_VERDICTS: nx_int = 6
72const NX_KV_BYTES: nx_int = 48 // 6 fields * 8
198const NX_KV_Q10: nx_int = 1024

functions

55func nx_kv_verdict_is_valid(v: nx_int) -> nx_int
called by 1: main
76func nx_kv_alloc(n_heads: nx_int, head_dim: nx_int, max_seq_len: nx_int) -> *NxKvCache
called by 1: main calls 1: sys_mmap
97func nx_kv_clear(kv: *NxKvCache) -> nx_int
called by 1: main
110func nx_kv_append(kv: *NxKvCache, k_in: *i64, v_in: *i64) -> nx_int
called by 1: main
131func nx_kv_get_k(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int
called by 1: main
142func nx_kv_get_v(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int
called by 1: main
164func nx_kv_score(kv: *NxKvCache, q_in: *i64, scores_out: *i64) -> nx_int
called by 1: main
200func nx_kv_apply(kv: *NxKvCache, weights_in: *i64, out: *i64) -> nx_int
called by 1: main
235func nx_kv_fill_ratio_q10(kv: *NxKvCache) -> nx_int
called by 1: main
246func nx_kv_bytes_per_token(kv: *NxKvCache) -> nx_int
called by 1: main