nx_kvcache.nx source

1// nx_kvcache.nx -- sovereign PAGED KV-CACHE (census OPPORTUNITY axis 2// "KV-cache / paged attention", momentum 3; vLLM-class, Kwon 2023). 3// 4// The contiguous NxF32KVCache allocates max_seq_len rows per sequence up 5// front: fragmentation, no sharing, one sequence per buffer. Paging fixes 6// all three: the cache is a POOL of fixed-size BLOCKS (BS tokens); a 7// sequence is a PAGE TABLE of block ids; blocks allocate on demand and 8// refcount for sharing. 9// 10// KEY STRUCTURAL FACT exploited here: decode is APPEND-ONLY -- past rows 11// are never mutated -- so full blocks are IMMUTABLE and can be shared 12// forever; only the PARTIAL TAIL block ever needs copy-on-append. fork() 13// is therefore O(page-table) + refcounts: prefill a prompt ONCE, fork N 14// sequences that share the prompt blocks (the self-consistency / best-of-N 15// serving pattern), each fork copies at most ONE partial block on its 16// first append. 17// 18// Pool layout (f32 bits in i64 slots, engine convention): 19// K_pool/V_pool: [n_blocks][n_layers][BS][kv_dim] 20// row t of seq -> block pt[t / BS], offset t % BS. 21// Fail-fast: pool exhaustion returns an error id (-1) -- NO silent 22// fallback to contiguous. 23// 24// genealogy_id: kwon_2023_pagedattention 25// lineage_id: substrate_paged_kv_v1 26 27import "nx_syscalls.nx" 28import "nx_tier.nx" 29 30const NX_PKV_BS: nx_int = 16 // tokens per block 31 32const NX_PKV_OK: nx_int = 0 33const NX_PKV_ERR_NULL: nx_int = 1 34const NX_PKV_ERR_BAD_DIM: nx_int = 2 35const NX_PKV_ERR_EXHAUST: nx_int = 3 36const NX_PKV_ERR_BOUNDS: nx_int = 4 37 38// ===== Pool ========================================================= 39 40struct NxPagedPool { 41 n_blocks: nx_int, 42 n_layers: nx_int, 43 kv_dim: nx_int, 44 K_pool: *i64, // [n_blocks][n_layers][BS][kv_dim] 45 V_pool: *i64, 46 refcnt: *i64, // per block; 0 = free 47 free_st: *i64, // stack of free block ids 48 n_free: i64 49} 50 51const NX_PKV_POOL_BYTES: nx_int = 64 52 53func nx_pkv_pool_new(n_blocks: nx_int, n_layers: nx_int, 54 kv_dim: nx_int) -> *NxPagedPool { 55 if n_blocks <= 0 { return 0 as *NxPagedPool } 56 if n_layers <= 0 { return 0 as *NxPagedPool } 57 if kv_dim <= 0 { return 0 as *NxPagedPool } 58 let p: *NxPagedPool = sys_mmap(NX_PKV_POOL_BYTES) as *NxPagedPool 59 p.n_blocks = n_blocks 60 p.n_layers = n_layers 61 p.kv_dim = kv_dim 62 let blk_elems: i64 = n_layers * NX_PKV_BS * kv_dim 63 p.K_pool = sys_mmap(n_blocks * blk_elems * 8) as *i64 64 p.V_pool = sys_mmap(n_blocks * blk_elems * 8) as *i64 65 p.refcnt = sys_mmap(n_blocks * 8) as *i64 66 p.free_st = sys_mmap(n_blocks * 8) as *i64 67 // free stack: push descending so ids allocate 0,1,2,... (readable KATs) 68 var i: nx_int = 0 69 while i < n_blocks { 70 p.free_st[i] = (n_blocks - 1 - i) as i64 71 i = i + 1 72 } 73 p.n_free = n_blocks as i64 74 return p 75} 76 77func nx_pkv_alloc_block(p: *NxPagedPool) -> i64 { 78 if p.n_free <= 0 { return 0 - 1 } // fail-fast: exhausted 79 p.n_free = p.n_free - 1 80 let b: i64 = p.free_st[p.n_free] 81 p.refcnt[b] = 1 82 return b 83} 84 85func nx_pkv_ref_block(p: *NxPagedPool, b: i64) -> i64 { 86 p.refcnt[b] = p.refcnt[b] + 1 87 return 0 88} 89 90func nx_pkv_unref_block(p: *NxPagedPool, b: i64) -> i64 { 91 p.refcnt[b] = p.refcnt[b] - 1 92 if p.refcnt[b] <= 0 { 93 p.refcnt[b] = 0 94 p.free_st[p.n_free] = b 95 p.n_free = p.n_free + 1 96 } 97 return 0 98} 99 100// element index of (block b, layer L, row r, dim 0) 101func _pkv_base(p: *NxPagedPool, b: i64, L: nx_int, r: nx_int) -> i64 { 102 return ((b * p.n_layers + L) * NX_PKV_BS + r) * p.kv_dim 103} 104 105// ===== Sequence (page table) ======================================== 106 107struct NxPagedSeq { 108 pool: *NxPagedPool, 109 pt: *i64, // page table: block ids 110 max_pt: nx_int, 111 n_pt: nx_int, 112 seq_len: nx_int 113} 114 115const NX_PKV_SEQ_BYTES: nx_int = 48 116 117func nx_pkv_seq_new(p: *NxPagedPool, max_tokens: nx_int) -> *NxPagedSeq { 118 if p == (0 as *NxPagedPool) { return 0 as *NxPagedSeq } 119 let s: *NxPagedSeq = sys_mmap(NX_PKV_SEQ_BYTES) as *NxPagedSeq 120 let mp: nx_int = (max_tokens + NX_PKV_BS - 1) / NX_PKV_BS + 1 121 s.pool = p 122 s.pt = sys_mmap(mp * 8) as *i64 123 s.max_pt = mp 124 s.n_pt = 0 125 s.seq_len = 0 126 return s 127} 128 129// fork: share every block (refcount++), same seq_len. O(page table). 130func nx_pkv_seq_fork(src: *NxPagedSeq) -> *NxPagedSeq { 131 if src == (0 as *NxPagedSeq) { return 0 as *NxPagedSeq } 132 let p: *NxPagedPool = src.pool 133 let s: *NxPagedSeq = sys_mmap(NX_PKV_SEQ_BYTES) as *NxPagedSeq 134 s.pool = p 135 s.pt = sys_mmap(src.max_pt * 8) as *i64 136 s.max_pt = src.max_pt 137 s.n_pt = src.n_pt 138 s.seq_len = src.seq_len 139 var i: nx_int = 0 140 while i < src.n_pt { 141 s.pt[i] = src.pt[i] 142 nx_pkv_ref_block(p, src.pt[i]) 143 i = i + 1 144 } 145 return s 146} 147 148func nx_pkv_seq_free(s: *NxPagedSeq) -> i64 { 149 if s == (0 as *NxPagedSeq) { return 0 } 150 var i: nx_int = 0 151 while i < s.n_pt { nx_pkv_unref_block(s.pool, s.pt[i]); i = i + 1 } 152 s.n_pt = 0 153 s.seq_len = 0 154 return 0 155} 156 157// Rewind (speculative-decode truncation): drop rows > new_len; unref any 158// now-unused whole blocks. The (possibly shared) tail keeps its data -- 159// rows past seq_len are semantically dead and copy-on-append protects 160// sharers. 161func nx_pkv_seq_truncate(s: *NxPagedSeq, new_len: nx_int) -> nx_int { 162 if new_len > s.seq_len { return NX_PKV_ERR_BOUNDS } 163 let need_pt: nx_int = (new_len + NX_PKV_BS - 1) / NX_PKV_BS 164 while s.n_pt > need_pt { 165 s.n_pt = s.n_pt - 1 166 nx_pkv_unref_block(s.pool, s.pt[s.n_pt]) 167 } 168 s.seq_len = new_len 169 return NX_PKV_OK 170} 171 172// ===== Capacity + copy-on-append ==================================== 173// Called ONCE per multi-token step (stack level, before per-layer 174// appends). Ensures blocks exist for rows [seq_len, seq_len+n_new) and 175// PRIVATIZES the partial tail block if shared (copy rows < tail_off for 176// ALL layers into a fresh block -- the only copy sharing ever costs). 177 178func nx_pkv_ensure_append(s: *NxPagedSeq, n_new: nx_int) -> nx_int { 179 if s == (0 as *NxPagedSeq) { return NX_PKV_ERR_NULL } 180 if n_new <= 0 { return NX_PKV_ERR_BAD_DIM } 181 let p: *NxPagedPool = s.pool 182 183 // 1) privatize the partial tail if shared. 184 let tail_off: nx_int = s.seq_len - (s.seq_len / NX_PKV_BS) * NX_PKV_BS 185 if tail_off > 0 { 186 let tb: i64 = s.pt[s.n_pt - 1] 187 if p.refcnt[tb] > 1 { 188 let nb: i64 = nx_pkv_alloc_block(p) 189 if nb < 0 { return NX_PKV_ERR_EXHAUST } 190 var L: nx_int = 0 191 while L < p.n_layers { 192 var r: nx_int = 0 193 while r < tail_off { 194 let src_base: i64 = _pkv_base(p, tb, L, r) 195 let dst_base: i64 = _pkv_base(p, nb, L, r) 196 var d: nx_int = 0 197 while d < p.kv_dim { 198 p.K_pool[dst_base + d] = p.K_pool[src_base + d] 199 p.V_pool[dst_base + d] = p.V_pool[src_base + d] 200 d = d + 1 201 } 202 r = r + 1 203 } 204 L = L + 1 205 } 206 nx_pkv_unref_block(p, tb) 207 s.pt[s.n_pt - 1] = nb 208 } 209 } 210 211 // 2) allocate blocks to cover seq_len + n_new rows. 212 let need_pt: nx_int = (s.seq_len + n_new + NX_PKV_BS - 1) / NX_PKV_BS 213 if need_pt > s.max_pt { return NX_PKV_ERR_BOUNDS } 214 while s.n_pt < need_pt { 215 let b2: i64 = nx_pkv_alloc_block(p) 216 if b2 < 0 { return NX_PKV_ERR_EXHAUST } 217 s.pt[s.n_pt] = b2 218 s.n_pt = s.n_pt + 1 219 } 220 return NX_PKV_OK 221} 222 223// Per-layer append (mirrors nx_f32_kv_cache_append_layer): write rows 224// [seq_len, seq_len+n) of layer L. Blocks must be pre-ensured. Does NOT 225// advance seq_len (the stack advances once after all layers). 226 227func nx_pkv_append_layer(s: *NxPagedSeq, L: nx_int, 228 K_new: *i64, V_new: *i64, n_new: nx_int) -> nx_int { 229 if s == (0 as *NxPagedSeq) { return NX_PKV_ERR_NULL } 230 let p: *NxPagedPool = s.pool 231 var t: nx_int = 0 232 while t < n_new { 233 let row: nx_int = s.seq_len + t 234 let b: i64 = s.pt[row / NX_PKV_BS] 235 let off: nx_int = row - (row / NX_PKV_BS) * NX_PKV_BS 236 let base: i64 = _pkv_base(p, b, L, off) 237 var d: nx_int = 0 238 while d < p.kv_dim { 239 p.K_pool[base + d] = K_new[t * p.kv_dim + d] 240 p.V_pool[base + d] = V_new[t * p.kv_dim + d] 241 d = d + 1 242 } 243 t = t + 1 244 } 245 return NX_PKV_OK 246} 247 248func nx_pkv_advance(s: *NxPagedSeq, n_new: nx_int) -> nx_int { 249 s.seq_len = s.seq_len + n_new 250 return NX_PKV_OK 251} 252 253// Gather one kv-head stripe of rows [0, total_k) into dst[total_k][head_dim] 254// (the paged twin of the contiguous per-head gather in attn_with_cache). 255func nx_pkv_gather_head(s: *NxPagedSeq, L: nx_int, kv_head: nx_int, 256 head_dim: nx_int, total_k: nx_int, 257 dstK: *i64, dstV: *i64) -> nx_int { 258 let p: *NxPagedPool = s.pool 259 var t: nx_int = 0 260 while t < total_k { 261 let b: i64 = s.pt[t / NX_PKV_BS] 262 let off: nx_int = t - (t / NX_PKV_BS) * NX_PKV_BS 263 let base: i64 = _pkv_base(p, b, L, off) + kv_head * head_dim 264 var d: nx_int = 0 265 while d < head_dim { 266 dstK[t * head_dim + d] = p.K_pool[base + d] 267 dstV[t * head_dim + d] = p.V_pool[base + d] 268 d = d + 1 269 } 270 t = t + 1 271 } 272 return NX_PKV_OK 273}