nx_batch_scheduler.nx source
↩ module page · 369 lines · 14317 B
1// nx_batch_scheduler.nx -- H2 continuous batching (bits-up).
2//
3// Per NISHI_ELDER_AI_OFF_DOCKER_2026_05_20.md ยง2.1 H2: pack pending
4// requests at TOKEN level not request level so the GPU stays
5// saturated. vLLM + TGI ship this pattern at scale (~10x throughput
6// over naive request-by-request scheduling).
7//
8// V1 mechanics: N-slot fixed scheduler. Each slot holds at most
9// one active request (sequence_id + tokens_remaining + state).
10// admit() places a request in a free slot OR queues it. step()
11// reports the slots currently in RUNNING state -- those are the
12// active batch for the next forward pass. complete() marks a
13// slot DONE. recycle() reclaims DONE slots and drains the
14// pending queue into the newly-free slots.
15//
16// The "continuous" part: at each scheduler step, the active batch
17// composition CAN CHANGE. A slot whose request finished at token
18// N gets reclaimed and refilled with a queued request at token
19// 0 -- so the GPU forward pass never has to wait for the whole
20// batch to finish.
21//
22// Pure substrate logic. No Linux features. Composes with shipped
23// nx_kv_arena (request's sequence holds its KV pages) and the
24// future actor cardinal (sequences serve text + image-prompt +
25// any other endpoint).
26//
27// V1 honest scope:
28// - Fixed slot count (N <= NX_BATCH_MAX_SLOTS = 32)
29// - Fixed pending queue (M <= NX_BATCH_MAX_PENDING = 64)
30// - First-come-first-served queue ordering (no priority yet)
31// - No per-sequence preemption (a slot runs to DONE; preemption
32// queued for V2 when attention-class arbitration lands)
33//
34// genealogy_id: vllm_continuous_batching_2023 + tgi_2023 +
35// cardinal_2026-05-20_elder_ai_off_docker +
36// cardinal_2026-05-20_bits_up_nishi_not_linux
37// lineage_id: substrate_batch_scheduler_v1
38//
39// nx_capability_manifest:
40// variant_class: batch_scheduler
41// variant_id: batch_scheduler_v1_slot_fcfs
42// requires_isa: [rv32i, rv64imac, x86_64, cortex_m, armv7a, aarch64, wasm32]
43// requires_syscalls: [mmap]
44// requires_ram_min_b: 4096
45// tier_floor: NX_TIER_INF_MOBILE
46// tier_ceiling: NX_TIER_INF_HPC
47// cost_model:
48// flops_per_n: 1.0 // O(n_slots) scan per step
49// bytes_per_n: 32.0 // per-slot bookkeeping
50// syscalls_per_n: 0.0
51// adversary_class: THREAT_AI_ADVERSARY
52//
53// nx_safety_envelope:
54// intended_use: "Continuous-batching request scheduler for
55// LLM serving; bits-up; composes with nx_kv_arena
56// and future actor cardinal"
57// sil_target: SIL2
58// evidence: [canary_bracketed, slot_state_sealed_enum,
59// refusal_on_queue_full, fcfs_drain_order]
60// verdict: NOT_YET_EVALUATED
61
62import "nx_syscalls.nx"
63
64// ===== Constants =================================================
65const NX_BATCH_MAX_SLOTS: i64 = 32
66const NX_BATCH_MAX_PENDING: i64 = 64
67
68// Slot state sealed enum.
69const NX_SLOT_FREE: i64 = 0
70const NX_SLOT_RUNNING: i64 = 1
71const NX_SLOT_DONE: i64 = 2
72const NX_SLOT_N_STATES: i64 = 3
73
74func nx_slot_state_is_valid(s: i64) -> i64 {
75 if s < 0 { return 0 }
76 if s >= NX_SLOT_N_STATES { return 0 }
77 return 1
78}
79
80// Verdicts (admit / complete / recycle).
81// admit() returns:
82// slot_idx >= 0 -> directly placed in slot
83// 0 - NX_BATCH_QUEUED -> placed in queue (will admit on recycle)
84// 0 - NX_BATCH_FULL -> slots + queue both full; REJECTED
85// 0 - NX_BATCH_BAD_INPUT -> n_tokens <= 0
86// 0 - NX_BATCH_TAMPER -> canary corrupted
87const NX_BATCH_OK: i64 = 0
88const NX_BATCH_BAD_INPUT: i64 = 1
89const NX_BATCH_QUEUED: i64 = 2 // admitted to queue, not slot
90const NX_BATCH_FULL: i64 = 3 // slots + queue both full
91const NX_BATCH_SLOT_NOT_RUNNING: i64 = 4
92const NX_BATCH_BAD_SLOT: i64 = 5
93const NX_BATCH_TAMPER: i64 = 6
94const NX_BATCH_N_VERDICTS: i64 = 7
95
96func nx_batch_verdict_is_valid(v: i64) -> i64 {
97 if v < 0 { return 0 }
98 if v >= NX_BATCH_N_VERDICTS { return 0 }
99 return 1
100}
101
102// Canaries.
103const NX_BATCH_CANARY_PRE: i64 = 0x4E5842415453504C // "NXBATSPL"
104const NX_BATCH_CANARY_POST: i64 = 0x4E58424154454E44 // "NXBATEND"
105
106// ===== Slot ======================================================
107struct NxBatchSlot {
108 state: i64, // NX_SLOT_*
109 request_id: i64, // operator-supplied request id
110 n_tokens_remaining: i64,
111}
112
113// ===== Scheduler ================================================
114struct NxBatchScheduler {
115 canary_pre: i64,
116 max_slots: i64,
117 n_running: i64,
118 n_done: i64,
119 slots: *i64, // *i64 array of NxBatchSlot pointers
120 // Pending queue: ring buffer of request_ids + token counts.
121 queue_request_ids: *i64,
122 queue_tokens: *i64,
123 queue_head: i64,
124 queue_tail: i64,
125 queue_size: i64,
126 queue_max: i64,
127 canary_post: i64,
128}
129
130// ===== Construction =============================================
131func nx_batch_scheduler_new(max_slots: i64, max_pending: i64) -> *NxBatchScheduler {
132 if max_slots <= 0 { return (0 as i64) as *NxBatchScheduler }
133 if max_slots > NX_BATCH_MAX_SLOTS { return (0 as i64) as *NxBatchScheduler }
134 if max_pending < 0 { return (0 as i64) as *NxBatchScheduler }
135 if max_pending > NX_BATCH_MAX_PENDING { return (0 as i64) as *NxBatchScheduler }
136
137 let sch: *NxBatchScheduler = (sys_mmap(96)) as *NxBatchScheduler
138 sch.canary_pre = NX_BATCH_CANARY_PRE
139 sch.canary_post = NX_BATCH_CANARY_POST
140 sch.max_slots = max_slots
141 sch.n_running = 0
142 sch.n_done = 0
143 sch.slots = (sys_mmap(max_slots * 8)) as *i64
144 var i: i64 = 0
145 while i < max_slots {
146 let slot: *NxBatchSlot = (sys_mmap(32)) as *NxBatchSlot
147 slot.state = NX_SLOT_FREE
148 slot.request_id = -1
149 slot.n_tokens_remaining = 0
150 sch.slots[i] = slot as i64
151 i = i + 1
152 }
153 sch.queue_max = max_pending
154 sch.queue_request_ids = (sys_mmap(max_pending * 8)) as *i64
155 sch.queue_tokens = (sys_mmap(max_pending * 8)) as *i64
156 sch.queue_head = 0
157 sch.queue_tail = 0
158 sch.queue_size = 0
159 return sch
160}
161
162// ===== Validity gate ============================================
163func nx_batch_scheduler_is_valid(sch: *NxBatchScheduler) -> i64 {
164 if (sch as i64) == 0 { return 0 }
165 if sch.canary_pre != NX_BATCH_CANARY_PRE { return 0 }
166 if sch.canary_post != NX_BATCH_CANARY_POST { return 0 }
167 if sch.max_slots < 0 { return 0 }
168 if sch.max_slots > NX_BATCH_MAX_SLOTS { return 0 }
169 if sch.n_running < 0 { return 0 }
170 if sch.n_running > sch.max_slots { return 0 }
171 if sch.queue_size < 0 { return 0 }
172 if sch.queue_size > sch.queue_max { return 0 }
173 return 1
174}
175
176// ===== Find first free slot =====================================
177// O(n_slots) linear scan; n_slots <= 32 so this is fine.
178func _batch_find_free_slot(sch: *NxBatchScheduler) -> i64 {
179 var i: i64 = 0
180 while i < sch.max_slots {
181 let slot: *NxBatchSlot = (sch.slots[i]) as *NxBatchSlot
182 if slot.state == NX_SLOT_FREE { return i }
183 i = i + 1
184 }
185 return -1
186}
187
188// ===== Admit a request ==========================================
189// If a free slot exists -> place directly in slot (RUNNING).
190// Else if queue has room -> push to queue.
191// Else -> NO_SLOTS_FREE (caller must wait + retry).
192// Returns slot_idx on direct admit, or 0 - NX_BATCH_* verdict.
193func nx_batch_scheduler_admit(
194 sch: *NxBatchScheduler,
195 request_id: i64,
196 n_tokens: i64
197) -> i64 {
198 if nx_batch_scheduler_is_valid(sch) != 1 { return 0 - NX_BATCH_TAMPER }
199 if n_tokens <= 0 { return 0 - NX_BATCH_BAD_INPUT }
200 let free_idx: i64 = _batch_find_free_slot(sch)
201 if free_idx >= 0 {
202 let slot: *NxBatchSlot = (sch.slots[free_idx]) as *NxBatchSlot
203 slot.state = NX_SLOT_RUNNING
204 slot.request_id = request_id
205 slot.n_tokens_remaining = n_tokens
206 sch.n_running = sch.n_running + 1
207 return free_idx
208 }
209 // No free slot -> try queue.
210 if sch.queue_size >= sch.queue_max { return 0 - NX_BATCH_FULL }
211 sch.queue_request_ids[sch.queue_tail] = request_id
212 sch.queue_tokens[sch.queue_tail] = n_tokens
213 sch.queue_tail = (sch.queue_tail + 1) % sch.queue_max
214 sch.queue_size = sch.queue_size + 1
215 return 0 - NX_BATCH_QUEUED
216}
217
218// ===== Step: enumerate currently-RUNNING slots ==================
219// Writes up to max_out slot indices that are RUNNING into out_buf.
220// Returns the number written. Caller uses these to assemble the
221// batch for the next forward pass.
222func nx_batch_scheduler_step(
223 sch: *NxBatchScheduler,
224 out_slot_ids: *i64,
225 max_out: i64
226) -> i64 {
227 if nx_batch_scheduler_is_valid(sch) != 1 { return 0 - NX_BATCH_TAMPER }
228 if (out_slot_ids as i64) == 0 { return 0 - NX_BATCH_BAD_INPUT }
229 if max_out <= 0 { return 0 - NX_BATCH_BAD_INPUT }
230 var n_out: i64 = 0
231 var i: i64 = 0
232 while i < sch.max_slots {
233 if n_out >= max_out { i = sch.max_slots; continue }
234 let slot: *NxBatchSlot = (sch.slots[i]) as *NxBatchSlot
235 if slot.state == NX_SLOT_RUNNING {
236 out_slot_ids[n_out] = i
237 n_out = n_out + 1
238 }
239 i = i + 1
240 }
241 return n_out
242}
243
244// ===== Decrement a slot's tokens (call after each forward pass) =
245// If tokens drop to 0, slot transitions to DONE. Returns the new
246// tokens-remaining count or 0 - NX_BATCH_* verdict.
247func nx_batch_scheduler_tick(
248 sch: *NxBatchScheduler,
249 slot_idx: i64,
250 n_tokens_produced: i64
251) -> i64 {
252 if nx_batch_scheduler_is_valid(sch) != 1 { return 0 - NX_BATCH_TAMPER }
253 if slot_idx < 0 { return 0 - NX_BATCH_BAD_SLOT }
254 if slot_idx >= sch.max_slots { return 0 - NX_BATCH_BAD_SLOT }
255 if n_tokens_produced < 0 { return 0 - NX_BATCH_BAD_INPUT }
256 let slot: *NxBatchSlot = (sch.slots[slot_idx]) as *NxBatchSlot
257 if slot.state != NX_SLOT_RUNNING { return 0 - NX_BATCH_SLOT_NOT_RUNNING }
258 slot.n_tokens_remaining = slot.n_tokens_remaining - n_tokens_produced
259 if slot.n_tokens_remaining <= 0 {
260 slot.n_tokens_remaining = 0
261 slot.state = NX_SLOT_DONE
262 sch.n_running = sch.n_running - 1
263 sch.n_done = sch.n_done + 1
264 }
265 return slot.n_tokens_remaining
266}
267
268// ===== Complete a slot explicitly ===============================
269// Forces a slot to DONE regardless of remaining tokens (caller
270// might cancel a request).
271func nx_batch_scheduler_complete(
272 sch: *NxBatchScheduler,
273 slot_idx: i64
274) -> i64 {
275 if nx_batch_scheduler_is_valid(sch) != 1 { return NX_BATCH_TAMPER }
276 if slot_idx < 0 { return NX_BATCH_BAD_SLOT }
277 if slot_idx >= sch.max_slots { return NX_BATCH_BAD_SLOT }
278 let slot: *NxBatchSlot = (sch.slots[slot_idx]) as *NxBatchSlot
279 if slot.state != NX_SLOT_RUNNING { return NX_BATCH_SLOT_NOT_RUNNING }
280 slot.state = NX_SLOT_DONE
281 slot.n_tokens_remaining = 0
282 sch.n_running = sch.n_running - 1
283 sch.n_done = sch.n_done + 1
284 return NX_BATCH_OK
285}
286
287// ===== Recycle DONE slots + drain queue =========================
288// Sweeps DONE slots to FREE, then pulls queued requests into the
289// newly-free slots (FCFS order). Returns the number of queued
290// requests that got admitted this recycle. Caller invokes this
291// AT THE END OF EACH FORWARD PASS to keep the slot table tight.
292func nx_batch_scheduler_recycle(sch: *NxBatchScheduler) -> i64 {
293 if nx_batch_scheduler_is_valid(sch) != 1 { return 0 - NX_BATCH_TAMPER }
294 // Sweep DONE -> FREE.
295 var i: i64 = 0
296 while i < sch.max_slots {
297 let slot: *NxBatchSlot = (sch.slots[i]) as *NxBatchSlot
298 if slot.state == NX_SLOT_DONE {
299 slot.state = NX_SLOT_FREE
300 slot.request_id = -1
301 slot.n_tokens_remaining = 0
302 sch.n_done = sch.n_done - 1
303 }
304 i = i + 1
305 }
306 // Drain queue into free slots.
307 var n_admitted: i64 = 0
308 var go: i64 = 1
309 while go == 1 {
310 if sch.queue_size <= 0 { go = 0 }
311 if go == 1 {
312 let free_idx: i64 = _batch_find_free_slot(sch)
313 if free_idx < 0 { go = 0 }
314 if go == 1 {
315 let req_id: i64 = sch.queue_request_ids[sch.queue_head]
316 let toks: i64 = sch.queue_tokens[sch.queue_head]
317 sch.queue_head = (sch.queue_head + 1) % sch.queue_max
318 sch.queue_size = sch.queue_size - 1
319 let slot: *NxBatchSlot = (sch.slots[free_idx]) as *NxBatchSlot
320 slot.state = NX_SLOT_RUNNING
321 slot.request_id = req_id
322 slot.n_tokens_remaining = toks
323 sch.n_running = sch.n_running + 1
324 n_admitted = n_admitted + 1
325 }
326 }
327 }
328 return n_admitted
329}
330
331// ===== Accessors ===============================================
332func nx_batch_scheduler_n_running(sch: *NxBatchScheduler) -> i64 {
333 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
334 return sch.n_running
335}
336
337func nx_batch_scheduler_n_pending(sch: *NxBatchScheduler) -> i64 {
338 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
339 return sch.queue_size
340}
341
342func nx_batch_scheduler_n_done(sch: *NxBatchScheduler) -> i64 {
343 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
344 return sch.n_done
345}
346
347func nx_batch_scheduler_slot_state(sch: *NxBatchScheduler, slot_idx: i64) -> i64 {
348 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
349 if slot_idx < 0 { return -1 }
350 if slot_idx >= sch.max_slots { return -1 }
351 let slot: *NxBatchSlot = (sch.slots[slot_idx]) as *NxBatchSlot
352 return slot.state
353}
354
355func nx_batch_scheduler_slot_request_id(sch: *NxBatchScheduler, slot_idx: i64) -> i64 {
356 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
357 if slot_idx < 0 { return -1 }
358 if slot_idx >= sch.max_slots { return -1 }
359 let slot: *NxBatchSlot = (sch.slots[slot_idx]) as *NxBatchSlot
360 return slot.request_id
361}
362
363func nx_batch_scheduler_slot_tokens_remaining(sch: *NxBatchScheduler, slot_idx: i64) -> i64 {
364 if nx_batch_scheduler_is_valid(sch) != 1 { return -1 }
365 if slot_idx < 0 { return -1 }
366 if slot_idx >= sch.max_slots { return -1 }
367 let slot: *NxBatchSlot = (sch.slots[slot_idx]) as *NxBatchSlot
368 return slot.n_tokens_remaining
369}