nx_nofloat_prefill_gate.nx source
↩ module page · 166 lines · 9688 B
1// nx_nofloat_prefill_gate.nx -- R0r (2026-09-17): the BATCHED PREFILL is bit-identical to the sequential per-token prefill it
2// replaces, proven on a real model through the serve core's own generate path, and its speed is printed as values.
3// Two runs of nsv_generate on one prompt long enough for two full NF_PREFILL_M blocks plus a partial one:
4// A = g_nsv_prefill_batched 0 (the sequential control) B = 1 (the served default)
5// Teeth: the fixture reached the condition (blocks and a partial block), every KV cache row of every layer for every
6// prompt position identical, a planted cell read as a difference (the comparator is live), the generated bytes and the
7// token count identical, the served default is batched. Prefill ms for both paths and the speedup are values.
8// The model (2026-09-18): argv[1] > sp_engine_model -- the engine model THIS HOST has staged (NX_ENGINE_MODEL env >
9// knowledge/stage.conf engine_model=), a valid subject because the identity holds for ANY model the engine runs -- >
10// the stage's nx_coder_model.gguf > SKIP when none opens. A roster forks this gate BARE, so the host-declared key is how
11// the heavy roster reaches a real model without staging another model under the coder model's name (a label that lies).
12// The rung that answered prints on the MODEL line (model_source=argv|engine_model|stage|none) and as model_source_code.
13// expect_exit: 0 license_tier: ORIGINAL
14import "nx_syscalls.nx"
15import "nx_nofloat_serve_core.nx"
16import "nx_gate_verdict.nx"
17import "nx_stage_path.nx"
18
19const PG_OUT_CAP: i64 = 8192
20const PG_MAX_NEW: i64 = 6
21const PG_SEED: i64 = 12345
22const PG_PERMIL: i64 = 1000
23const PG_REPEATS: i64 = 12 // sentence repeats: ~230 tokens, so two full 32-blocks and a partial one
24const PG_PROMPT_CAP: i64 = 4096
25const PG_GP_SLOTS: i64 = 16
26const PG_META_SLOTS: i64 = 8
27const PG_MODE_I8: i64 = 1
28
29func pg_cat(d: *u8, o: i64, s: *u8) -> i64 { var i: i64 = 0; var p: i64 = o; while s[i] != (0 as u8) { d[p] = s[i]; p = p + 1; i = i + 1 } return p }
30// rows 0..nrows-1 of every layer's K and V, laid out layer by layer, into dst; returns the i64 count copied
31func pg_kv_snapshot(dst: *i64, nrows: i64) -> i64 {
32 let nl: i64 = g_nsv_dims[1]
33 let kvd: i64 = g_nsv_dims[6]
34 var off: i64 = 0
35 var L: i64 = 0
36 while L < nl {
37 cpy(((dst as i64)+off*8) as *i64, g_nsv_kvc[2*L] as *i64, nrows*kvd); off = off + nrows*kvd
38 cpy(((dst as i64)+off*8) as *i64, g_nsv_kvc[2*L+1] as *i64, nrows*kvd); off = off + nrows*kvd
39 L = L + 1
40 }
41 return off
42}
43func pg_same(a: *i64, b: *i64, n: i64) -> i64 { var i: i64 = 0; while i < n { if a[i] != b[i] { return 0 } i = i + 1 } return 1 }
44
45// model_source codes, emitted as the value model_source_code (the label is printed on the MODEL line)
46const PG_MSRC_NONE: i64 = 0
47const PG_MSRC_ARGV: i64 = 1
48const PG_MSRC_ENGINE: i64 = 2
49const PG_MSRC_STAGE: i64 = 3
50func pg_src_name(s: i64) -> *u8 {
51 if s == PG_MSRC_ARGV { return "argv" as *u8 }
52 if s == PG_MSRC_ENGINE { return "engine_model" as *u8 }
53 if s == PG_MSRC_STAGE { return "stage" as *u8 }
54 return "none" as *u8
55}
56// the MODEL line, printed before the model loads: which model, which rung answered, and what the engine_model key said
57func pg_announce(msrc: i64, mpath: *u8, esrc: i64, epath: *u8) -> i64 {
58 gv_puts("MODEL model_source=" as *u8); gv_puts(pg_src_name(msrc))
59 gv_puts(" model=" as *u8); gv_puts(mpath)
60 if msrc != PG_MSRC_ARGV {
61 gv_puts(" engine_model_rung=" as *u8); gv_puts(sp_engine_src_name(esrc))
62 if esrc != SP_SRC_NONE { gv_puts(" engine_model=" as *u8); gv_puts(epath); gv_puts(" engine_model_opens=" as *u8); gv_num(sp_present(epath)) }
63 }
64 gv_puts("\n (rungs: argv > engine_model [NX_ENGINE_MODEL env > knowledge/stage.conf engine_model=] > stage nx_coder_model.gguf > none = SKIP; model_source_code argv=1 engine_model=2 stage=3 none=0)\n\n" as *u8)
65 return 0
66}
67
68func main(argc: i64, argv: *i64) -> i64 {
69 let ctr: *i64 = sys_mmap(64) as *i64
70 gv_puts("nx_nofloat_prefill_gate -- batched prefill (one weight pass per block) is bit-identical to the sequential prefill\n\n" as *u8)
71 // THE MODEL: the first rung that names a file that OPENS answers -- argv[1] > the engine model this host has staged
72 // (sp_engine_model) > the stage's nx_coder_model.gguf > none = SKIP, before any check (the ordering law).
73 var mpath: *u8 = 0 as *u8
74 var msrc: i64 = PG_MSRC_NONE
75 let epath: *u8 = sys_mmap(SP_PATH_MAX)
76 var esrc: i64 = SP_SRC_NONE
77 if argc > 1 { mpath = argv[1] as *u8; msrc = PG_MSRC_ARGV } else {
78 esrc = sp_engine_model(epath)
79 if esrc != SP_SRC_NONE { if sp_present(epath) == 1 { mpath = epath; msrc = PG_MSRC_ENGINE } }
80 if msrc == PG_MSRC_NONE {
81 mpath = sp_path("nx_coder_model.gguf" as *u8, sys_mmap(SP_PATH_MAX))
82 if sp_present(mpath) == 1 { msrc = PG_MSRC_STAGE }
83 }
84 }
85 pg_announce(msrc, mpath, esrc, epath)
86 if msrc == PG_MSRC_NONE {
87 gv_kv("model_source_code" as *u8, msrc)
88 if esrc != SP_SRC_NONE { sp_engine_model_skip("NOFLOAT-PREFILL-GATE" as *u8, epath, esrc) }
89 sp_skip_unless("NOFLOAT-PREFILL-GATE" as *u8, mpath)
90 }
91 let ir: i64 = nsv_init_i8(mpath) // the served route: i8 caches only, memory-proportionate (no i32 cache beside them)
92 gv_need("model-initialised" as *u8, (ir == 0) as i64, ctr)
93 let prompt: *u8 = sys_mmap(PG_PROMPT_CAP)
94 var plen: i64 = 0
95 var r: i64 = 0
96 while r < PG_REPEATS { plen = pg_cat(prompt, plen, "The quick brown fox jumps over the lazy dog near the river bank while the sun sets behind the hills. " as *u8); r = r + 1 }
97 let outA: *u8 = sys_mmap(PG_OUT_CAP)
98 let outB: *u8 = sys_mmap(PG_OUT_CAP)
99 let metaA: *i64 = sys_mmap(PG_META_SLOTS*8) as *i64
100 let metaB: *i64 = sys_mmap(PG_META_SLOTS*8) as *i64
101 let gp: *i64 = sys_mmap(PG_GP_SLOTS*8) as *i64
102 gp[0] = prompt as i64
103 gp[1] = plen
104 gp[2] = PG_MAX_NEW
105 gp[3] = PG_MODE_I8
106 gp[7] = 0
107 gp[8] = 0
108 gp[9] = 0
109 gp[10] = 0
110 gp[11] = PG_SEED
111 gp[12] = 0
112 // A: the sequential control -- per-token prefill with SERIAL attention (R0v's control arm)
113 g_nsv_prefill_batched = 0
114 g_nf_attn_serial = 1
115 gp[4] = outA as i64; gp[5] = PG_OUT_CAP; gp[6] = metaA as i64
116 let nA: i64 = nsv_generate(gp)
117 let nprompt: i64 = metaA[0]
118 let nl: i64 = g_nsv_dims[1]
119 let kvd: i64 = g_nsv_dims[6]
120 let kvn: i64 = nl*2*nprompt*kvd
121 let kvA: *i64 = sys_mmap(kvn*8) as *i64
122 let kvB: *i64 = sys_mmap(kvn*8) as *i64
123 let cpA: i64 = pg_kv_snapshot(kvA, nprompt)
124 // B: the served default -- batched prefill with POOL attention and pool silu (R0v); T1/T2 compare it to A
125 g_nsv_prefill_batched = 1
126 g_nf_attn_serial = NF_ATTN_SERIAL_DEFAULT
127 gp[4] = outB as i64; gp[5] = PG_OUT_CAP; gp[6] = metaB as i64
128 let nB: i64 = nsv_generate(gp)
129 let cpB: i64 = pg_kv_snapshot(kvB, nprompt)
130 gv_kv("model_source_code" as *u8, msrc)
131 gv_kv("prompt_tokens" as *u8, nprompt)
132 gv_kv("kv_cells_compared" as *u8, cpA)
133 gv_kv("prefill_ms_sequential" as *u8, metaA[6])
134 gv_kv("prefill_ms_batched" as *u8, metaB[6])
135 var spd: i64 = 0
136 if metaB[6] > 0 { spd = (metaA[6] * PG_PERMIL) / metaB[6] }
137 gv_kv("prefill_speedup_permil" as *u8, spd)
138 gv_kv("total_ms_sequential" as *u8, metaA[2])
139 gv_kv("total_ms_batched" as *u8, metaB[2])
140 // R0s: where the batched prefill's time went (microseconds accumulated over the LAST batched call, all layers)
141 gv_kv("pf_us_norm" as *u8, g_nf_pf_us_norm)
142 gv_kv("pf_us_qkv" as *u8, g_nf_pf_us_qkv)
143 gv_kv("pf_us_attn" as *u8, g_nf_pf_us_attn)
144 gv_kv("pf_us_o" as *u8, g_nf_pf_us_o)
145 gv_kv("pf_us_gateup" as *u8, g_nf_pf_us_gu)
146 gv_kv("pf_us_silu" as *u8, g_nf_pf_us_silu)
147 gv_kv("pf_us_down" as *u8, g_nf_pf_us_down)
148 gv_kv("pf_us_quant_within" as *u8, g_nf_pf_us_quant)
149 gv_check("T5 R0v the control arm ran serial attention and the served arm the pool (the KV and byte identities above compared them)" as *u8, (NF_ATTN_SERIAL_DEFAULT == 0) as i64, ctr)
150 gv_check("T0 the fixture reached the condition: the prompt spans more than one block" as *u8, (nprompt > NF_PREFILL_M) as i64, ctr)
151 gv_check("T0a the fixture reached the condition: the last block is partial" as *u8, ((nprompt % NF_PREFILL_M) != 0) as i64, ctr)
152 gv_check("T0b both runs generated (no error code)" as *u8, ((nA > 0) as i64) * ((nB > 0) as i64), ctr)
153 gv_check_eq("T0c both snapshots copied the same cell count" as *u8, cpB, cpA, ctr)
154 gv_check("T1 every KV cache row of every layer for every prompt position is bit-identical" as *u8, pg_same(kvA, kvB, kvn), ctr)
155 kvB[kvd+1] = kvB[kvd+1] + 1
156 gv_check("neg-control-T1b the KV comparator reads a planted +1 as a difference" as *u8, (pg_same(kvA, kvB, kvn) == 0) as i64, ctr)
157 kvB[kvd+1] = kvB[kvd+1] - 1
158 var sameout: i64 = 1
159 if nA != nB { sameout = 0 }
160 if sameout == 1 { var k: i64 = 0; while k < nA { if outA[k] != outB[k] { sameout = 0 } k = k + 1 } }
161 gv_check("T2 the generated bytes are identical (greedy, six tokens)" as *u8, sameout, ctr)
162 gv_check_eq("T3 the generated token count is identical" as *u8, metaB[1], metaA[1], ctr)
163 gv_check_eq("T4 the served default is the batched prefill" as *u8, NSV_PREFILL_BATCHED_DEFAULT, 1, ctr)
164 let rc: i64 = gv_verdict("NOFLOAT-PREFILL-GATE" as *u8, ctr, "one weight pass per block of prompt positions, bit-identical to the per-token prefill on a real model" as *u8)
165 return gv_exit(rc) // exit_group through the base class, not the raw thread exit: the pool workers parked in futex wait must die with the verdict
166}