nx_kv_arena.nx source
↩ module page · 342 lines · 13928 B
1// nx_kv_arena.nx -- PagedAttention page pool (bits-up).
2//
3// vLLM's PagedAttention (UC Berkeley, SOSP'23) is the killer
4// LLM-serving algorithm: instead of one contiguous K/V buffer
5// per sequence, split into fixed-size PAGES. Each sequence owns
6// a list of page indices; pages can be SHARED across sequences
7// for prefix caching (e.g., shared system prompt = one cached
8// page); free pages return to the pool. vLLM published 24x
9// throughput over HuggingFace TGI on the same hardware.
10//
11// THIS primitive is the page-pool substrate. K/V tensor data
12// itself is stored elsewhere (composes with shipped nx_kv_cache
13// for the tensor mechanics). This is the ALLOCATOR + INDEX
14// layer that makes paged-attention possible.
15//
16// V1 scope:
17// - Fixed pool of NX_KV_ARENA_MAX_PAGES pages
18// - Each page is a fixed-size byte block (caller picks page_bytes
19// at arena_new time; typical 4 KiB for 16 tokens of f16 KV per
20// head)
21// - Sequence holds a list of page indices it owns (no sharing
22// in V1; prefix caching V2 adds refcount-on-pages)
23// - Pool exhaustion -> NULL append; caller handles (caller's
24// responsibility to evict via nx_lysosome or yield via
25// nx_yield_protocol)
26// - Canary-bracketed for tamper detection
27//
28// Deferred to V2:
29// - Cross-sequence page sharing (refcount per page; prefix cache)
30// - Page eviction policy (LRU? attention-class? composes with
31// nx_attention_class)
32// - Dynamic page allocation (grow arena on demand)
33// - Multi-tier pages (GPU VRAM vs RAM vs NVMe; composes with
34// conductor cardinal)
35//
36// genealogy_id: vllm_paged_attention_2023_sosp + s_lora_punica_2023 +
37// cardinal_2026-05-20_elder_ai_off_docker +
38// cardinal_2026-05-20_bits_up_nishi_not_linux
39// lineage_id: substrate_kv_arena_v1
40//
41// nx_capability_manifest:
42// variant_class: kv_paged_arena
43// variant_id: kv_arena_v1_fixed_pool
44// requires_isa: [rv32i, rv64imac, x86_64, cortex_m, armv7a, aarch64, avr, wasm32]
45// requires_syscalls: [mmap]
46// requires_ram_min_b: 16384 // arena + per-page overhead
47// tier_floor: NX_TIER_INF_MOBILE
48// tier_ceiling: NX_TIER_INF_HPC
49// cost_model:
50// flops_per_n: 1.0 // O(1) alloc per page
51// bytes_per_n: 4096.0 // typical page size
52// syscalls_per_n: 0.0 // pool pre-allocated
53// adversary_class: THREAT_AI_ADVERSARY
54//
55// nx_safety_envelope:
56// intended_use: "Page pool for PagedAttention-style KV cache;
57// bits-up; host-agnostic; composes with
58// nx_kv_cache tensor layer"
59// sil_target: SIL2
60// evidence: [canary_bracketed, fixed_pool_no_fragmentation,
61// refusal_on_exhaustion, free_list_reuse]
62// verdict: NOT_YET_EVALUATED
63
64import "nx_syscalls.nx"
65
66// ===== Constants =================================================
67const NX_KV_ARENA_MAX_PAGES: i64 = 256
68const NX_KV_ARENA_MAX_SEQ_PAGES: i64 = 64 // max pages per sequence (V1)
69const NX_KV_ARENA_DEFAULT_PAGE_B: i64 = 4096 // 4 KiB
70
71// Verdicts.
72const NX_KV_OK: i64 = 0
73const NX_KV_BAD_INPUT: i64 = 1
74const NX_KV_POOL_EXHAUSTED: i64 = 2
75const NX_KV_SEQ_FULL: i64 = 3 // per-seq page count exceeded
76const NX_KV_TAMPER: i64 = 4
77const NX_KV_NOT_FOUND: i64 = 5
78const NX_KV_N_VERDICTS: i64 = 6
79
80func nx_kv_verdict_is_valid(v: i64) -> i64 {
81 if v < 0 { return 0 }
82 if v >= NX_KV_N_VERDICTS { return 0 }
83 return 1
84}
85
86// Canary distinct from other primitive canaries.
87const NX_KV_ARENA_CANARY_PRE: i64 = 0x4E584B56415250 // "NXKVARP\0"
88const NX_KV_ARENA_CANARY_POST: i64 = 0x4E584B564152454E // "NXKVAREN"
89const NX_KV_SEQ_CANARY_PRE: i64 = 0x4E584B5653455150 // "NXKVSEQP"
90const NX_KV_SEQ_CANARY_POST: i64 = 0x4E584B5653454E45 // "NXKVSENE"
91
92// ===== Page state (sealed enum) =================================
93const NX_KV_PAGE_FREE: i64 = 0
94const NX_KV_PAGE_IN_USE: i64 = 1
95
96// ===== NxKvArena ================================================
97// page_refcount[i]: number of sequences currently referencing page i.
98// 0 -> page FREE; can be allocated by any sequence
99// 1 -> page IN_USE by exactly one sequence (the
100// legacy non-shared case)
101// 2..N -> page SHARED by N sequences (H4 prefix-caching)
102//
103// page_states[i] is kept as a sealed-enum view derived from refcount:
104// refcount == 0 <-> state == NX_KV_PAGE_FREE
105// refcount > 0 <-> state == NX_KV_PAGE_IN_USE
106// Both are updated together; consistency is part of arena validity.
107struct NxKvArena {
108 canary_pre: i64,
109 n_pages: i64,
110 max_pages: i64,
111 page_bytes: i64,
112 n_in_use: i64,
113 n_free: i64,
114 page_states: *i64, // *i64 array of NX_KV_PAGE_* per page
115 page_refcount: *i64, // *i64 array of refcounts per page (H4)
116 page_storage: *u8, // contiguous byte buffer; page i at
117 // offset i * page_bytes
118 canary_post: i64,
119}
120
121// ===== NxKvSequence =============================================
122// Holds a list of page indices that belong to this sequence.
123// Indices point into the parent arena's page_states / page_storage.
124struct NxKvSequence {
125 canary_pre: i64,
126 arena: *NxKvArena,
127 seq_id: i64,
128 n_pages_used: i64,
129 max_pages: i64,
130 page_indices: *i64, // *i64 array of page indices owned by this seq
131 canary_post: i64,
132}
133
134// ===== Construction =============================================
135// Allocate an arena with `n_pages` pages of `page_bytes` each.
136// Caller picks both; choose page_bytes to match the K/V layout
137// (e.g., 16 tokens * n_heads * head_dim * 2 (K+V) * sizeof(f16)).
138
139func nx_kv_arena_new(n_pages: i64, page_bytes: i64) -> *NxKvArena {
140 if n_pages <= 0 { return (0 as i64) as *NxKvArena }
141 if n_pages > NX_KV_ARENA_MAX_PAGES { return (0 as i64) as *NxKvArena }
142 if page_bytes <= 0 { return (0 as i64) as *NxKvArena }
143
144 let a: *NxKvArena = (sys_mmap(128)) as *NxKvArena
145 a.canary_pre = NX_KV_ARENA_CANARY_PRE
146 a.canary_post = NX_KV_ARENA_CANARY_POST
147 a.n_pages = n_pages
148 a.max_pages = n_pages
149 a.page_bytes = page_bytes
150 a.n_in_use = 0
151 a.n_free = n_pages
152 a.page_states = (sys_mmap(n_pages * 8)) as *i64
153 a.page_refcount = (sys_mmap(n_pages * 8)) as *i64
154 a.page_storage = sys_mmap(n_pages * page_bytes)
155 var i: i64 = 0
156 while i < n_pages {
157 a.page_states[i] = NX_KV_PAGE_FREE
158 a.page_refcount[i] = 0
159 i = i + 1
160 }
161 return a
162}
163
164// ===== Validity gate ============================================
165func nx_kv_arena_is_valid(a: *NxKvArena) -> i64 {
166 if (a as i64) == 0 { return 0 }
167 if a.canary_pre != NX_KV_ARENA_CANARY_PRE { return 0 }
168 if a.canary_post != NX_KV_ARENA_CANARY_POST { return 0 }
169 if a.n_pages < 0 { return 0 }
170 if a.n_pages > a.max_pages { return 0 }
171 if a.n_in_use < 0 { return 0 }
172 if a.n_free < 0 { return 0 }
173 if (a.n_in_use + a.n_free) != a.n_pages { return 0 }
174 return 1
175}
176
177// ===== Find first free page =====================================
178// O(n_pages) linear scan; acceptable for V1 fixed-pool sizes.
179// V2 will maintain an explicit free-list head for O(1) alloc.
180func _kv_arena_find_free(a: *NxKvArena) -> i64 {
181 var i: i64 = 0
182 while i < a.n_pages {
183 if a.page_states[i] == NX_KV_PAGE_FREE { return i }
184 i = i + 1
185 }
186 return -1
187}
188
189// ===== Sequence construction ===================================
190func nx_kv_sequence_new(arena: *NxKvArena, seq_id: i64) -> *NxKvSequence {
191 if nx_kv_arena_is_valid(arena) != 1 { return (0 as i64) as *NxKvSequence }
192 let s: *NxKvSequence = (sys_mmap(64)) as *NxKvSequence
193 s.canary_pre = NX_KV_SEQ_CANARY_PRE
194 s.canary_post = NX_KV_SEQ_CANARY_POST
195 s.arena = arena
196 s.seq_id = seq_id
197 s.n_pages_used = 0
198 s.max_pages = NX_KV_ARENA_MAX_SEQ_PAGES
199 s.page_indices = (sys_mmap(NX_KV_ARENA_MAX_SEQ_PAGES * 8)) as *i64
200 var i: i64 = 0
201 while i < NX_KV_ARENA_MAX_SEQ_PAGES {
202 s.page_indices[i] = -1
203 i = i + 1
204 }
205 return s
206}
207
208// ===== Sequence validity ========================================
209func nx_kv_sequence_is_valid(s: *NxKvSequence) -> i64 {
210 if (s as i64) == 0 { return 0 }
211 if s.canary_pre != NX_KV_SEQ_CANARY_PRE { return 0 }
212 if s.canary_post != NX_KV_SEQ_CANARY_POST { return 0 }
213 if s.n_pages_used < 0 { return 0 }
214 if s.n_pages_used > s.max_pages { return 0 }
215 if nx_kv_arena_is_valid(s.arena) != 1 { return 0 }
216 return 1
217}
218
219// ===== Append a page to a sequence =============================
220// Returns the page INDEX on success, or a negative verdict. Caller
221// uses the index to compute the byte offset into arena.page_storage:
222// byte_off = idx * arena.page_bytes
223// page_ptr = (arena.page_storage as i64) + byte_off
224func nx_kv_sequence_append_page(s: *NxKvSequence) -> i64 {
225 if nx_kv_sequence_is_valid(s) != 1 { return 0 - NX_KV_TAMPER }
226 if s.n_pages_used >= s.max_pages { return 0 - NX_KV_SEQ_FULL }
227 let a: *NxKvArena = s.arena
228 if a.n_free <= 0 { return 0 - NX_KV_POOL_EXHAUSTED }
229 let free_idx: i64 = _kv_arena_find_free(a)
230 if free_idx < 0 { return 0 - NX_KV_POOL_EXHAUSTED }
231 a.page_states[free_idx] = NX_KV_PAGE_IN_USE
232 a.page_refcount[free_idx] = 1
233 a.n_free = a.n_free - 1
234 a.n_in_use = a.n_in_use + 1
235 s.page_indices[s.n_pages_used] = free_idx
236 s.n_pages_used = s.n_pages_used + 1
237 return free_idx
238}
239
240// ===== H4 prefix caching: share an existing page ================
241// Adds src_page_idx to this sequence's page list AND increments
242// the page's refcount. The page must already be IN_USE
243// (refcount > 0); sharing a free page is rejected (caller should
244// use append_page for fresh allocation).
245//
246// Use case: shared system prompt. Sequence A appends pages for
247// the system prompt; sequences B / C / D share those same pages
248// instead of duplicating the K/V data. Each subsequent free
249// decrements the refcount; the page returns to the pool only
250// when the LAST owner frees it.
251func nx_kv_sequence_share_page(s: *NxKvSequence, src_page_idx: i64) -> i64 {
252 if nx_kv_sequence_is_valid(s) != 1 { return 0 - NX_KV_TAMPER }
253 if s.n_pages_used >= s.max_pages { return 0 - NX_KV_SEQ_FULL }
254 let a: *NxKvArena = s.arena
255 if src_page_idx < 0 { return 0 - NX_KV_BAD_INPUT }
256 if src_page_idx >= a.n_pages { return 0 - NX_KV_BAD_INPUT }
257 // Page must already be in use; sharing a FREE page is wrong.
258 if a.page_refcount[src_page_idx] <= 0 { return 0 - NX_KV_NOT_FOUND }
259 if a.page_states[src_page_idx] != NX_KV_PAGE_IN_USE { return 0 - NX_KV_NOT_FOUND }
260 // Increment refcount; do NOT touch n_in_use / n_free
261 // (page was already counted as in_use when first allocated).
262 a.page_refcount[src_page_idx] = a.page_refcount[src_page_idx] + 1
263 s.page_indices[s.n_pages_used] = src_page_idx
264 s.n_pages_used = s.n_pages_used + 1
265 return src_page_idx
266}
267
268// ===== Refcount query for a page ===============================
269// Returns the page's current refcount, or -1 on invalid input.
270// Useful for asserts + tests; in production the substrate doesn't
271// usually need this (the abstraction is "I own a page until I
272// free it").
273func nx_kv_arena_page_refcount(a: *NxKvArena, page_idx: i64) -> i64 {
274 if nx_kv_arena_is_valid(a) != 1 { return -1 }
275 if page_idx < 0 { return -1 }
276 if page_idx >= a.n_pages { return -1 }
277 return a.page_refcount[page_idx]
278}
279
280// ===== Free a sequence (return all its pages to the pool) ======
281// Resets the sequence's page list AND returns each page to the
282// arena's free pool. Idempotent: calling twice on the same
283// sequence after first free is a no-op (n_pages_used is 0).
284func nx_kv_sequence_free(s: *NxKvSequence) -> i64 {
285 if nx_kv_sequence_is_valid(s) != 1 { return NX_KV_TAMPER }
286 let a: *NxKvArena = s.arena
287 var i: i64 = 0
288 while i < s.n_pages_used {
289 let idx: i64 = s.page_indices[i]
290 if idx >= 0 {
291 if idx < a.n_pages {
292 if a.page_refcount[idx] > 0 {
293 // Decrement refcount. Only return to pool when
294 // the LAST owner frees (refcount hits 0).
295 a.page_refcount[idx] = a.page_refcount[idx] - 1
296 if a.page_refcount[idx] == 0 {
297 a.page_states[idx] = NX_KV_PAGE_FREE
298 a.n_in_use = a.n_in_use - 1
299 a.n_free = a.n_free + 1
300 }
301 }
302 }
303 }
304 s.page_indices[i] = -1
305 i = i + 1
306 }
307 s.n_pages_used = 0
308 return NX_KV_OK
309}
310
311// ===== Get a page byte pointer by index =========================
312// Returns the start of the page's byte buffer; NULL on bad input.
313// Caller writes/reads K/V data into this buffer per their layout
314// (composes with nx_kv_cache for tensor mechanics).
315func nx_kv_arena_page_bytes(a: *NxKvArena, page_idx: i64) -> *u8 {
316 if nx_kv_arena_is_valid(a) != 1 { return (0 as i64) as *u8 }
317 if page_idx < 0 { return (0 as i64) as *u8 }
318 if page_idx >= a.n_pages { return (0 as i64) as *u8 }
319 let off: i64 = page_idx * a.page_bytes
320 return ((a.page_storage as i64) + off) as *u8
321}
322
323// ===== Stats accessors =========================================
324func nx_kv_arena_n_in_use(a: *NxKvArena) -> i64 {
325 if nx_kv_arena_is_valid(a) != 1 { return -1 }
326 return a.n_in_use
327}
328
329func nx_kv_arena_n_free(a: *NxKvArena) -> i64 {
330 if nx_kv_arena_is_valid(a) != 1 { return -1 }
331 return a.n_free
332}
333
334func nx_kv_arena_n_pages(a: *NxKvArena) -> i64 {
335 if nx_kv_arena_is_valid(a) != 1 { return -1 }
336 return a.n_pages
337}
338
339func nx_kv_sequence_n_pages(s: *NxKvSequence) -> i64 {
340 if nx_kv_sequence_is_valid(s) != 1 { return -1 }
341 return s.n_pages_used
342}