nx_kv_cache.nx source

1// nx_kv_cache.nx -- autoregressive attention KV cache. 2// 3// THE foundation primitive of efficient autoregressive generation 4// (LLM serving, forever-generating world models, autoregressive 5// diffusion). Without it, each new token recomputes attention over 6// every past token: O(n^2) per step. With it: O(n) per step. 7// For n=1000 generation steps that's a 1000x speedup with bit- 8// identical output (algorithmic correctness, not approximation). 9// 10// Per the min-hardware-floor + algo-led cardinal: this brick is 11// what makes a $50 SBC run autoregressive generation that a naive 12// O(n^2) recompute couldn't finish in a year. 13// 14// Composes with: 15// * nx_sparse_tensor -- KV cache rows are mostly zero post-softmax; 16// store sparse, compute sparse 17// * nx_quant_block -- cached K and V can be 4-bit; 10x less 18// memory bandwidth at each new-token step 19// * nx_compute_graph -- KV cache is a typed CONTROL node in the 20// graph; runner persists state across runs of the same graph 21// 22// Storage layout: 23// k_storage[head * max_seq_len * head_dim + pos * head_dim + d] 24// v_storage[head * max_seq_len * head_dim + pos * head_dim + d] 25// 26// Multi-head attention works by treating each head as an independent 27// per-position dot product space. We don't materialise the [Q, K, V] 28// triple as a single tensor here -- the cache OWNS K and V, the 29// caller's Q changes every step. 30// 31// genealogy_id: vaswani_2017_transformer_kv + flashattention_2_kv + 32// pope_2023_efficient_serving + mistral_sliding_window 33// lineage_id: substrate_kv_cache_v1 34 35// nx_safety_envelope: 36// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 37// sil_target: SIL1 38// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 39// verdict: NOT_YET_EVALUATED 40 41import "nx_syscalls.nx" 42import "nx_tier.nx" 43const NX_MAGIC_1024: i64 = 1024 44 45// ===== Sealed-enum: KvVerdict ===================================== 46 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 54 55func nx_kv_verdict_is_valid(v: nx_int) -> nx_int { 56 if v < 0 { return 0 } 57 if v >= NX_KV_N_VERDICTS { return 0 } 58 return 1 59} 60 61// ===== Cache struct =============================================== 62 63struct NxKvCache { 64 n_heads: nx_int, 65 head_dim: nx_int, 66 max_seq_len: nx_int, 67 n_filled: nx_int, 68 k_storage: *i64, // [n_heads * max_seq_len * head_dim] 69 v_storage: *i64 // same shape 70} 71 72const NX_KV_BYTES: nx_int = 48 // 6 fields * 8 73 74// ===== Builder ==================================================== 75 76func nx_kv_alloc(n_heads: nx_int, head_dim: nx_int, max_seq_len: nx_int) -> *NxKvCache { 77 let kv: *NxKvCache = (sys_mmap(NX_KV_BYTES)) as *NxKvCache 78 kv.n_heads = n_heads 79 kv.head_dim = head_dim 80 kv.max_seq_len = max_seq_len 81 kv.n_filled = 0 82 let storage_count: nx_int = n_heads * max_seq_len * head_dim 83 let storage_bytes: nx_int = storage_count * NX_SIZEOF_NX_INT 84 kv.k_storage = (sys_mmap(storage_bytes)) as *i64 85 kv.v_storage = (sys_mmap(storage_bytes)) as *i64 86 // Zero-init so reads past n_filled return zeros instead of garbage 87 var i: nx_int = 0 88 while i < storage_count { 89 kv.k_storage[i] = 0 90 kv.v_storage[i] = 0 91 i = i + 1 92 } 93 return kv 94} 95 96// Reset (new sequence) 97func nx_kv_clear(kv: *NxKvCache) -> nx_int { 98 kv.n_filled = 0 99 return 0 100} 101 102// ===== Append a new token's K + V projections ==================== 103// 104// k_in / v_in: flat [n_heads * head_dim] -- each head's new vector 105// concatenated. The cache moves them into per-head per-position 106// storage so subsequent reads can re-walk past tokens. 107// 108// Returns the position the new token was placed at, or -1 if full. 109 110func nx_kv_append(kv: *NxKvCache, k_in: *i64, v_in: *i64) -> nx_int { 111 if kv.n_filled >= kv.max_seq_len { return 0 - 1 } 112 let pos: nx_int = kv.n_filled 113 var h: nx_int = 0 114 while h < kv.n_heads { 115 var d: nx_int = 0 116 while d < kv.head_dim { 117 let dst_idx: nx_int = h * kv.max_seq_len * kv.head_dim + pos * kv.head_dim + d 118 let src_idx: nx_int = h * kv.head_dim + d 119 kv.k_storage[dst_idx] = k_in[src_idx] 120 kv.v_storage[dst_idx] = v_in[src_idx] 121 d = d + 1 122 } 123 h = h + 1 124 } 125 kv.n_filled = kv.n_filled + 1 126 return pos 127} 128 129// ===== Read accessors ============================================= 130 131func nx_kv_get_k(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int { 132 if head < 0 { return 0 } 133 if head >= kv.n_heads { return 0 } 134 if pos < 0 { return 0 } 135 if pos >= kv.n_filled { return 0 } 136 if dim < 0 { return 0 } 137 if dim >= kv.head_dim { return 0 } 138 let idx: nx_int = head * kv.max_seq_len * kv.head_dim + pos * kv.head_dim + dim 139 return kv.k_storage[idx] 140} 141 142func nx_kv_get_v(kv: *NxKvCache, head: nx_int, pos: nx_int, dim: nx_int) -> nx_int { 143 if head < 0 { return 0 } 144 if head >= kv.n_heads { return 0 } 145 if pos < 0 { return 0 } 146 if pos >= kv.n_filled { return 0 } 147 if dim < 0 { return 0 } 148 if dim >= kv.head_dim { return 0 } 149 let idx: nx_int = head * kv.max_seq_len * kv.head_dim + pos * kv.head_dim + dim 150 return kv.v_storage[idx] 151} 152 153// ===== Compute Q @ K^T across the cache ========================== 154// 155// q_in: flat [n_heads * head_dim] -- the new token's query vectors. 156// scores_out: flat [n_heads * n_filled] -- pre-softmax attention 157// scores. Caller composes softmax (when we have it) and feeds back 158// into nx_kv_attention_apply. 159// 160// Standard attention shape: score[h, p] = dot(Q[h, :], K[h, p, :]). 161// This is O(n_heads * n_filled * head_dim) -- linear in current 162// sequence length, not quadratic. 163 164func nx_kv_score(kv: *NxKvCache, q_in: *i64, scores_out: *i64) -> nx_int { 165 if kv.n_filled <= 0 { return NX_KV_OK } // nothing to score 166 var h: nx_int = 0 167 while h < kv.n_heads { 168 var p: nx_int = 0 169 while p < kv.n_filled { 170 var acc: nx_int = 0 171 var d: nx_int = 0 172 while d < kv.head_dim { 173 let q_idx: nx_int = h * kv.head_dim + d 174 let k_idx: nx_int = h * kv.max_seq_len * kv.head_dim + p * kv.head_dim + d 175 acc = acc + q_in[q_idx] * kv.k_storage[k_idx] 176 d = d + 1 177 } 178 scores_out[h * kv.n_filled + p] = acc 179 p = p + 1 180 } 181 h = h + 1 182 } 183 return NX_KV_OK 184} 185 186// ===== Apply attention weights to V ============================== 187// 188// weights_in: flat [n_heads * n_filled] -- typically softmax(scores). 189// For testing we accept raw weights; the caller decides whether 190// to normalise. 191// 192// out: flat [n_heads * head_dim] -- weighted sum of V vectors: 193// out[h, d] = sum_p(weights[h, p] * V[h, p, d]) 194// 195// For Q10 substrate semantics, weights are typically Q10 already. 196// We do (acc) / Q10 at the end to normalise. 197 198const NX_KV_Q10: nx_int = 1024 199 200func nx_kv_apply(kv: *NxKvCache, weights_in: *i64, out: *i64) -> nx_int { 201 if kv.n_filled <= 0 { 202 // Nothing in cache -> output is zero 203 var z: nx_int = 0 204 while z < kv.n_heads * kv.head_dim { 205 out[z] = 0 206 z = z + 1 207 } 208 return NX_KV_OK 209 } 210 var h: nx_int = 0 211 while h < kv.n_heads { 212 var d: nx_int = 0 213 while d < kv.head_dim { 214 var acc: nx_int = 0 215 var p: nx_int = 0 216 while p < kv.n_filled { 217 let w_idx: nx_int = h * kv.n_filled + p 218 let v_idx: nx_int = h * kv.max_seq_len * kv.head_dim + p * kv.head_dim + d 219 acc = acc + weights_in[w_idx] * kv.v_storage[v_idx] 220 p = p + 1 221 } 222 out[h * kv.head_dim + d] = acc / NX_KV_Q10 223 d = d + 1 224 } 225 h = h + 1 226 } 227 return NX_KV_OK 228} 229 230// ===== Memory usage (Q10 fraction of capacity) =================== 231// 232// fill_ratio_q10 = n_filled / max_seq_len * 1024 233// Useful for "should we evict?" decisions in long-context serving. 234 235func nx_kv_fill_ratio_q10(kv: *NxKvCache) -> nx_int { 236 if kv.max_seq_len <= 0 { return 0 } 237 return (kv.n_filled * NX_MAGIC_1024) / kv.max_seq_len 238} 239 240// ===== Bytes-per-token (memory accounting) ======================= 241// 242// Per token: n_heads * head_dim * 8 bytes (K) + same (V) = 2 * 243// n_heads * head_dim * 8. Useful for capacity planning on small 244// hardware ("can I run 4096 tokens on 4 GB?"). 245 246func nx_kv_bytes_per_token(kv: *NxKvCache) -> nx_int { 247 return 2 * kv.n_heads * kv.head_dim * 8 248}