nx_hackers_algo_compose_test.nx source
↩ module page · 176 lines · 8697 B
1// nx_hackers_algo_compose_test.nx -- end-to-end composition of the
2// ten hackers-algo primitives. Proves that the substrate primitives
3// shipped 2026-05-20 compose without integration bugs.
4//
5// Scenario: serving two concurrent multi-modal requests against a
6// shared backbone with GQA attention, multi-LoRA adapters, paged
7// KV with prefix-sharing, continuous batching, speculative
8// decoding, and a token-merge plan. All bits-up; zero Linux
9// delegation. Substrate composition only -- no tensor kernels,
10// no model weights -- the goal here is to verify primitive
11// interaction, not to claim end-to-end inference correctness.
12
13import "nx_syscalls.nx"
14import "nx_kv_arena.nx"
15import "nx_batch_scheduler.nx"
16import "nx_speculative_verify.nx"
17import "nx_lora_pool.nx"
18import "nx_modality_router.nx"
19import "nx_gqa_head_map.nx"
20import "nx_token_merge.nx"
21
22func main() -> i64 {
23 // ===== Set up shared backbone + modality router =====
24 let backbone_weights: *i64 = (sys_mmap(512)) as *i64
25 let head_text: *i64 = (sys_mmap(128)) as *i64
26 let head_image: *i64 = (sys_mmap(128)) as *i64
27 let adapter_a_data: *i64 = (sys_mmap(64)) as *i64
28 let adapter_a_b_data: *i64 = (sys_mmap(64)) as *i64
29 let adapter_i_data: *i64 = (sys_mmap(64)) as *i64
30 let adapter_i_b_data: *i64 = (sys_mmap(64)) as *i64
31
32 let router: *NxModalityRouter = nx_modality_router_new(4, 8)
33 if nx_modality_register_backbone(router, 1, backbone_weights, 4096, 4096, 32) < 0 { return 1 }
34 if nx_modality_register_head(router, 100, 1, head_text, 4096, NX_MODALITY_TEXT) < 0 { return 2 }
35 if nx_modality_register_head(router, 101, 1, head_image, 4096, NX_MODALITY_IMAGE) < 0 { return 3 }
36 // Backbone refcount must equal modality count.
37 if nx_modality_backbone_refcount(router, 1) != 2 { return 4 }
38
39 // ===== H6 GQA: 32 Q heads, 8 KV heads (Llama-3-8B style) =====
40 let gqa: *NxGqaHeadMap = nx_gqa_head_map_new(32, 8)
41 if nx_gqa_head_map_kind(gqa) != NX_ATTN_GQA { return 5 }
42 if nx_gqa_head_map_group_size(gqa) != 4 { return 6 }
43
44 // ===== H9 LoRA pool: one adapter per modality =====
45 let lora: *NxLoraPool = nx_lora_pool_new(8)
46 if nx_lora_register(lora, 5001, NX_LORA_RANK_8, 32, adapter_a_data, adapter_a_b_data, 65536) < 0 { return 7 }
47 if nx_lora_register(lora, 5002, NX_LORA_RANK_16, 32, adapter_i_data, adapter_i_b_data, 131072) < 0 { return 8 }
48
49 // ===== H1+H4 KV arena: 16 pages, 1024 bytes each =====
50 let arena: *NxKvArena = nx_kv_arena_new(16, 1024)
51 if nx_kv_arena_n_pages(arena) != 16 { return 9 }
52 if nx_kv_arena_n_free(arena) != 16 { return 10 }
53
54 // ===== H2 batch scheduler: 4 slots, 8 pending =====
55 let sch: *NxBatchScheduler = nx_batch_scheduler_new(4, 8)
56
57 // ===== Admit request A: TEXT modality, 16 tokens, adapter 5001 =====
58 let slot_a: i64 = nx_batch_scheduler_admit(sch, 9001, 16)
59 if slot_a < 0 { return 11 }
60 if nx_lora_acquire(lora, 5001) != 1 { return 12 }
61 let seq_a: *NxKvSequence = nx_kv_sequence_new(arena, 9001)
62 if nx_kv_sequence_append_page(seq_a) < 0 { return 13 }
63 let spec_a: *NxSpeculativeSession = nx_speculative_session_new(9001, 8)
64
65 // ===== Admit request B: IMAGE modality, 8 tokens, adapter 5002 =====
66 let slot_b: i64 = nx_batch_scheduler_admit(sch, 9002, 8)
67 if slot_b < 0 { return 14 }
68 if nx_lora_acquire(lora, 5002) != 1 { return 15 }
69 let seq_b: *NxKvSequence = nx_kv_sequence_new(arena, 9002)
70 let page_b: i64 = nx_kv_sequence_append_page(seq_b)
71 if page_b < 0 { return 16 }
72 let spec_b: *NxSpeculativeSession = nx_speculative_session_new(9002, 8)
73
74 // ===== H4 prefix sharing: request A shares request B's first page =====
75 if nx_kv_sequence_share_page(seq_a, page_b) < 0 { return 17 }
76 if nx_kv_arena_page_refcount(arena, page_b) != 2 { return 18 }
77
78 // ===== Both slots resolve to the SAME backbone via the router =====
79 let bb_text: *NxBackbone = nx_modality_resolve_backbone(router, 100)
80 let bb_image: *NxBackbone = nx_modality_resolve_backbone(router, 101)
81 if (bb_text as i64) == 0 { return 19 }
82 if (bb_image as i64) == 0 { return 20 }
83 if (bb_text as i64) != (bb_image as i64) { return 21 } // shared backbone
84
85 // ===== H3 speculative round for request A (5 drafts, 4 accepted) =====
86 let draft_a: *i64 = (sys_mmap(128)) as *i64
87 let flags_a: *i64 = (sys_mmap(128)) as *i64
88 let out_a: *i64 = (sys_mmap(128)) as *i64
89 draft_a[0] = 1001
90 draft_a[1] = 1002
91 draft_a[2] = 1003
92 draft_a[3] = 1004
93 draft_a[4] = 1005
94 flags_a[0] = 1
95 flags_a[1] = 1
96 flags_a[2] = 1
97 flags_a[3] = 1
98 flags_a[4] = 0
99 let committed_a: i64 = nx_speculative_verify(spec_a, draft_a, flags_a, 5, 2000, out_a)
100 if committed_a != 5 { return 22 } // 4 accepted + 1 bonus
101 if nx_spec_total_accepted(spec_a) != 4 { return 23 }
102
103 // ===== Tick request A: consumed 5 tokens, 11 remaining =====
104 let rem_a: i64 = nx_batch_scheduler_tick(sch, slot_a, committed_a)
105 if rem_a != 11 { return 24 }
106 if nx_batch_scheduler_slot_state(sch, slot_a) != NX_SLOT_RUNNING { return 25 }
107
108 // ===== H8 token merge plan: shorten request A's 11-token tail to 8 =====
109 let merge_plan: *NxTokenMergePlan = nx_token_merge_plan_new(11)
110 let pairs: *i64 = (sys_mmap(256)) as *i64
111 pairs[0] = 0
112 pairs[1] = 1
113 pairs[2] = 99
114 pairs[3] = 2
115 pairs[4] = 3
116 pairs[5] = 98
117 pairs[6] = 4
118 pairs[7] = 5
119 pairs[8] = 97
120 let merges: i64 = nx_token_merge_apply(merge_plan, pairs, 3, 3)
121 if merges != 3 { return 26 }
122 if nx_token_merge_n_kept(merge_plan) != 8 { return 27 } // 11 - 3 == 8
123
124 // ===== Slot B completes early =====
125 let rem_b: i64 = nx_batch_scheduler_tick(sch, slot_b, 8)
126 if rem_b != 0 { return 28 }
127 if nx_batch_scheduler_slot_state(sch, slot_b) != NX_SLOT_DONE { return 29 }
128
129 // Release B's adapter + free its KV sequence.
130 if nx_lora_release(lora, 5002) != 0 { return 30 }
131 let freed_b: i64 = nx_kv_sequence_free(seq_b)
132 if freed_b < 0 { return 31 }
133 // page_b had refcount 2 (B's own + A's share); freeing B drops to 1, not back to free pool.
134 if nx_kv_arena_page_refcount(arena, page_b) != 1 { return 32 }
135
136 // ===== Backbone refcount drops because B's head is still attached =====
137 // (modality_router doesn't auto-detach on batch slot completion;
138 // the consumer would call unregister_head if the modality is being
139 // torn down. For the smoke we just verify refcount math is consistent.)
140 if nx_modality_backbone_refcount(router, 1) != 2 { return 33 }
141
142 // ===== Recycle the scheduler -- B's slot reclaimed, no queue to drain =====
143 let admitted: i64 = nx_batch_scheduler_recycle(sch)
144 if admitted != 0 { return 34 } // queue empty
145 if nx_batch_scheduler_n_running(sch) != 1 { return 35 } // A still running
146
147 // ===== Canary invariant audit at the end =====
148 if nx_modality_router_is_valid(router) != 1 { return 36 }
149 if nx_lora_pool_is_valid(lora) != 1 { return 37 }
150 if nx_kv_arena_is_valid(arena) != 1 { return 38 }
151 if nx_batch_scheduler_is_valid(sch) != 1 { return 39 }
152 if nx_gqa_head_map_is_valid(gqa) != 1 { return 40 }
153 if nx_spec_session_is_valid(spec_a) != 1 { return 41 }
154 if nx_spec_session_is_valid(spec_b) != 1 { return 42 }
155 if nx_token_merge_plan_is_valid(merge_plan) != 1 { return 43 }
156 if nx_kv_sequence_is_valid(seq_a) != 1 { return 44 }
157
158 // ===== Cross-primitive invariants =====
159 // 1) modality router backbone refcount equals number of attached heads.
160 if nx_modality_backbone_refcount(router, 1) != nx_modality_count_heads_for_backbone(router, 1) { return 45 }
161 // 2) lora pool refcount equals number of acquires minus releases.
162 // Request A acquired adapter 5001 once and never released -> refcount 1.
163 if nx_lora_refcount(lora, 5001) != 1 { return 46 }
164 // Request B acquired 5002 once and released once -> refcount 0.
165 if nx_lora_refcount(lora, 5002) != 0 { return 47 }
166 // 3) GQA invariant: group_size * n_kv_heads == n_q_heads.
167 if (nx_gqa_head_map_group_size(gqa) * nx_gqa_head_map_n_kv_heads(gqa)) != nx_gqa_head_map_n_q_heads(gqa) { return 48 }
168 // 4) Token merge invariant: kept + merged == n_tokens.
169 if (nx_token_merge_n_kept(merge_plan) + nx_token_merge_n_merged(merge_plan)) != nx_token_merge_n_tokens(merge_plan) { return 49 }
170 // 5) Speculative monotonic: accepted <= drafted.
171 if nx_spec_total_accepted(spec_a) > nx_spec_total_drafted(spec_a) { return 50 }
172 // 6) KV arena: page_refcount on shared page is consistent.
173 if nx_kv_arena_page_refcount(arena, page_b) != 1 { return 51 }
174
175 return 0
176}