nx_f32_llm_v4.nx source
↩ module page · 288 lines · 13787 B
1// nx_f32_llm_v4.nx -- lazy-aware model binder + LLM forward.
2//
3// Closes the lazy Q4_K path end-to-end:
4// * Model binder uses nx_f32_llama_layer_lazy_load_from_gguf
5// for each layer (Q4_K weights stay as offsets; F32/F16 eagerly
6// widened).
7// * Embed + lm_head + gamma_out loaded eagerly to f32 (small + 1D
8// for gamma_out; embed/lm_head typically F32 or F16 in real
9// GGUFs).
10// * Forward dispatches through nx_f32_llama_stack_forward_v4.
11//
12// Reuses NxF32LlamaModel struct from nx_f32_llm.nx; the layers
13// array now holds *NxF32LlamaLayerLazy values (cast through i64).
14// Caller MUST use v4 forward consistently with v4-loaded model.
15//
16// genealogy_id: standard_llama_forward + lazy_dispatch
17// lineage_id: substrate_f32_llm_v4
18
19import "nx_syscalls.nx"
20import "nx_tier.nx"
21import "nx_gguf.nx"
22import "nx_gguf_load_f32.nx"
23import "nx_f32.nx"
24import "nx_f32_rmsnorm.nx"
25import "nx_f32_matmul.nx"
26import "nx_f32_matmul_t.nx"
27import "nx_f32_kv_cache.nx"
28import "nx_f32_lazy_weight.nx"
29import "nx_f32_llama_block_v4.nx"
30import "nx_f32_llama_stack_v4.nx"
31import "nx_f32_llama_layer_lazy_load.nx"
32import "nx_f32_llm.nx"
33
34const NX_FLV4_OK: nx_int = 0
35const NX_FLV4_ERR_NULL: nx_int = 1
36const NX_FLV4_ERR_BAD_DIM: nx_int = 2
37const NX_FLV4_ERR_NOT_FOUND: nx_int = 3
38const NX_FLV4_ERR_LAYER: nx_int = 4
39const NX_FLV4_ERR_STACK: nx_int = 5
40const NX_FLV4_ERR_TOKEN: nx_int = 6
41const NX_FLV4_N_VERDICTS: nx_int = 7
42
43func nx_flv4_verdict_is_valid(v: nx_int) -> nx_int {
44 if v < 0 { return 0 }
45 if v >= NX_FLV4_N_VERDICTS { return 0 }
46 return 1
47}
48
49// Load weights using the lazy-aware per-layer binder.
50// Top-level (embed / output_norm / output) eagerly dequanted to f32.
51
52func nx_f32_llm_load_weights_v4_from_gguf(buf: *u8, hdr: *NxGgufHeader,
53 model: *NxF32LlamaModel,
54 out_err: *i64) -> nx_int {
55 if model == (0 as *NxF32LlamaModel) { out_err[0] = NX_FLV4_ERR_NULL; return NX_FLV4_ERR_NULL }
56 if model.n_layers <= 0 { out_err[0] = NX_FLV4_ERR_BAD_DIM; return NX_FLV4_ERR_BAD_DIM }
57
58 let nv_out: *i64 = sys_mmap(8) as *i64
59 let inner_err: *i64 = sys_mmap(8) as *i64
60
61 // token_embd.weight (eager).
62 let n_te: *u8 = sys_mmap(17)
63 n_te[0]=0x74 as u8; n_te[1]=0x6f as u8; n_te[2]=0x6b as u8; n_te[3]=0x65 as u8
64 n_te[4]=0x6e as u8; n_te[5]=0x5f as u8; n_te[6]=0x65 as u8; n_te[7]=0x6d as u8
65 n_te[8]=0x62 as u8; n_te[9]=0x64 as u8; n_te[10]=0x2e as u8; n_te[11]=0x77 as u8
66 n_te[12]=0x65 as u8; n_te[13]=0x69 as u8; n_te[14]=0x67 as u8; n_te[15]=0x68 as u8
67 n_te[16]=0x74 as u8
68 let embed: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_te, 17, nv_out, inner_err)
69 if inner_err[0] != NX_GLF_OK { out_err[0] = NX_FLV4_ERR_NOT_FOUND; return NX_FLV4_ERR_NOT_FOUND }
70 model.embed_weights = embed
71
72 // output_norm.weight (eager).
73 let n_on: *u8 = sys_mmap(18)
74 n_on[0]=0x6f as u8; n_on[1]=0x75 as u8; n_on[2]=0x74 as u8; n_on[3]=0x70 as u8
75 n_on[4]=0x75 as u8; n_on[5]=0x74 as u8; n_on[6]=0x5f as u8; n_on[7]=0x6e as u8
76 n_on[8]=0x6f as u8; n_on[9]=0x72 as u8; n_on[10]=0x6d as u8; n_on[11]=0x2e as u8
77 n_on[12]=0x77 as u8; n_on[13]=0x65 as u8; n_on[14]=0x69 as u8; n_on[15]=0x67 as u8
78 n_on[16]=0x68 as u8; n_on[17]=0x74 as u8
79 let gout: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_on, 18, nv_out, inner_err)
80 if inner_err[0] != NX_GLF_OK { out_err[0] = NX_FLV4_ERR_NOT_FOUND; return NX_FLV4_ERR_NOT_FOUND }
81 model.gamma_out = gout
82
83 // output.weight (eager, tied-embed fallback).
84 let n_ow: *u8 = sys_mmap(13)
85 n_ow[0]=0x6f as u8; n_ow[1]=0x75 as u8; n_ow[2]=0x74 as u8; n_ow[3]=0x70 as u8
86 n_ow[4]=0x75 as u8; n_ow[5]=0x74 as u8; n_ow[6]=0x2e as u8; n_ow[7]=0x77 as u8
87 n_ow[8]=0x65 as u8; n_ow[9]=0x69 as u8; n_ow[10]=0x67 as u8; n_ow[11]=0x68 as u8
88 n_ow[12]=0x74 as u8
89 let lm: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, n_ow, 13, nv_out, inner_err)
90 if inner_err[0] == NX_GLF_OK {
91 model.lm_head = lm
92 } else {
93 model.lm_head = embed
94 }
95
96 // v4 (2026-07-08): if output.weight is Q8_0, ALSO build a Q8_0 LAZY weight
97 // so the forward runs the SIMD __f32_i8dot32 dequant-dot for lm_head
98 // (PROVEN 10.2x vs the eager-F32 matmul). Falls back to eager F32 (above)
99 // if not Q8_0 or absent. Row j (vocab) = hidden Q8_0 values contiguous;
100 // the dispatcher uses k=hidden,n=vocab from the forward call.
101 model.lm_head_q8 = 0
102 let ow_idx: nx_int = nx_gguf_find_tensor(hdr, n_ow, 13)
103 if ow_idx >= 0 {
104 let ow_ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, ow_idx)
105 let ow_rows: nx_int = ow_ti.dim_0 as nx_int
106 var ow_cols: nx_int = 1
107 if ow_ti.n_dims >= 2 { ow_cols = ow_ti.dim_1 as nx_int }
108 let ow_off: i64 = hdr.data_off + ow_ti.offset
109 if ow_ti.ggml_type == NX_GGML_TYPE_Q8_0 {
110 let lmq: *NxF32LazyWeight = nx_f32_lazy_weight_new_q8_0(buf, ow_off, ow_rows, ow_cols)
111 model.lm_head_q8 = lmq as i64
112 }
113 // ---- Q4_K / Q5_0 lm_head ALSO stay quantized (2026-08-01) --------------------------------------
114 // WHY THIS BRANCH HAD TO EXIST: the Q8_0 case above was wired 2026-07-08 and PROVEN 10.2x, but it was
115 // the ONLY quantized lm_head the binder recognised -- anything else fell through to the eager-F32
116 // path. That is the wrong half of the field. Q4_K_M is the 2026 CPU-inference STANDARD (llama.cpp /
117 // GGUF, ~4.5 bits/weight), and in a Q4_K_M file output.weight is Q4_K or Q6_K, never F32 -- so on
118 // exactly the models the ecosystem is meant to run, the biggest matmul in the forward was being
119 // widened to F32 and read at 8 bytes/value.
120 // SCALE: nx_matmul_t_pool_gate measures lm_head at ~1.09GB as F32. Q8_0 is ~273MB and Q4_K ~137MB of
121 // the SAME tensor, and single-token decode is MEMORY-BANDWIDTH bound (our own roofline: decode matvec
122 // Q4 AI=3.55 < ridge 8.53), so bytes-moved IS the runtime. This is the "keep weights QUANTIZED and
123 // dequant inside the dot" lever the perf lane named, applied to the tensor that dominates it.
124 // SAFE BY CONSTRUCTION: the dtype_tag dispatcher in nx_f32_lazy_matmul already implements Q4_K and
125 // Q5_0 fused dequant-dot (NX_LW_DTYPE_Q4_K / _Q5_0) and is exercised by nx_f32_lazy_weight_test; this
126 // adds NO new numerics, it only stops discarding a path that was already built and tested. Reusing
127 // the lm_head_q8 slot is deliberate -- the forward branches on it being nonzero and then dispatches
128 // on dtype_tag, so one slot serves every quantized format and the eager-F32 fallback is untouched
129 // for genuinely F32/F16 files (rule 19).
130 // u00e2u02dcu2026A DISPATCHER THAT HANDLES FOUR FORMATS BEHIND A BINDER THAT RECOGNISES ONE IS A ONE-FORMAT SYSTEM.
131 if ow_ti.ggml_type == NX_GGML_TYPE_Q4_K {
132 let lmq4: *NxF32LazyWeight = nx_f32_lazy_weight_new_q4k(buf, ow_off, ow_rows, ow_cols)
133 model.lm_head_q8 = lmq4 as i64
134 }
135 if ow_ti.ggml_type == NX_GGML_TYPE_Q5_0 {
136 let lmq5: *NxF32LazyWeight = nx_f32_lazy_weight_new_q5_0(buf, ow_off, ow_rows, ow_cols)
137 model.lm_head_q8 = lmq5 as i64
138 }
139 }
140
141 // Per-layer LAZY binding.
142 model.layers = sys_mmap(model.n_layers * 8) as *i64
143 var L: nx_int = 0
144 while L < model.n_layers {
145 let layer: *NxF32LlamaLayerLazy = nx_f32_llama_layer_lazy_alloc()
146 let v: nx_int = nx_f32_llama_layer_lazy_load_from_gguf(
147 buf, hdr, L, layer, inner_err)
148 if v != NX_FLLL_OK { out_err[0] = NX_FLV4_ERR_LAYER; return NX_FLV4_ERR_LAYER }
149 model.layers[L] = layer as i64
150 L = L + 1
151 }
152
153 out_err[0] = NX_FLV4_OK
154 return NX_FLV4_OK
155}
156
157// v4 LLM forward: embed lookup + lazy stack + final norm + LM head.
158
159// ★ONE IMPLEMENTATION, TWO ENTRY POINTS. This core carries the extra hidden_out parameter; the public
160// nx_f32_llm_forward_v4 below is a thin wrapper that passes 0, so ALL 58 EXISTING CALL SITES compile
161// untouched (rule 19 -- adding capability must not break a contract 58 places depend on). A second copy
162// of the forward pass would have been the alternative and it is the worse one: two things to keep
163// correct, guaranteed to diverge the first time either is edited.
164func _flv4_core(
165 model: *NxF32LlamaModel,
166 token_ids: *i64,
167 n_tokens: nx_int,
168 cache: *NxF32KVCache,
169 eps: i64,
170 attn_scale: i64,
171 rope_log_base: i64,
172 apply_rope: nx_int,
173 logits: *i64,
174 hidden_out: *i64) -> nx_int {
175
176 if model == (0 as *NxF32LlamaModel) { return NX_FLV4_ERR_NULL }
177 if token_ids == (0 as *i64) { return NX_FLV4_ERR_NULL }
178 if cache == (0 as *NxF32KVCache) { return NX_FLV4_ERR_NULL }
179 // logits may be NULL ONLY on the embedding path, where it is never written. Generation callers pass
180 // hidden_out = 0 and are validated exactly as before -- the contract they rely on is unchanged.
181 if hidden_out == (0 as *i64) { if logits == (0 as *i64) { return NX_FLV4_ERR_NULL } }
182 if n_tokens <= 0 { return NX_FLV4_ERR_BAD_DIM }
183
184 let hidden_dim: nx_int = model.hidden_dim
185 let vocab_size: nx_int = model.vocab_size
186
187 let x_embed: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
188 let x_stacked: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
189 let x_normed: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
190
191 var t: nx_int = 0
192 while t < n_tokens {
193 let tok: nx_int = token_ids[t] as nx_int
194 if tok < 0 { return NX_FLV4_ERR_TOKEN }
195 if tok >= vocab_size { return NX_FLV4_ERR_TOKEN }
196 var d: nx_int = 0
197 while d < hidden_dim {
198 x_embed[t * hidden_dim + d] = model.embed_weights[tok * hidden_dim + d]
199 d = d + 1
200 }
201 t = t + 1
202 }
203
204 let v_stk: nx_int = nx_f32_llama_stack_forward_v4(
205 x_embed, x_stacked, n_tokens, model.n_layers,
206 hidden_dim, model.n_heads, model.n_kv_heads, model.head_dim, model.ffn_dim,
207 model.layers, cache, eps, attn_scale, rope_log_base, apply_rope)
208 if v_stk != NX_STK4_OK { return NX_FLV4_ERR_STACK }
209
210 var t2: nx_int = 0
211 while t2 < n_tokens {
212 let src: *i64 = (((x_stacked as i64) + t2 * hidden_dim * 8)) as *i64
213 let dst: *i64 = (((x_normed as i64) + t2 * hidden_dim * 8)) as *i64
214 nx_f32_rmsnorm(src, model.gamma_out, hidden_dim, eps, dst)
215 t2 = t2 + 1
216 }
217
218 // EMBEDDING TAP (rung 1): x_normed IS the sentence representation. hidden_out != 0 => SUM-POOLED
219 // x_normed and lm_head SKIPPED. Rationale + measurements in nx_f32_embed_gate.nx (kept THERE, not
220 // here: this closure sits at the nx_cc source-size cap, see debt 1785449xxx).
221 if hidden_out != (0 as *i64) {
222 var hd: nx_int = 0
223 while hd < hidden_dim { hidden_out[hd] = 0; hd = hd + 1 }
224 var tp: nx_int = 0
225 while tp < n_tokens {
226 let row: *i64 = (((x_normed as i64) + tp * hidden_dim * 8)) as *i64
227 var d2: nx_int = 0
228 while d2 < hidden_dim {
229 hidden_out[d2] = nx_f32_add(hidden_out[d2], row[d2])
230 d2 = d2 + 1
231 }
232 tp = tp + 1
233 }
234 return NX_FLV4_OK
235 }
236
237 // lm_head = the forward's biggest single matmul (hidden x vocab =
238 // 896 x 151936 = 136M MACs/token). MEASURED 2026-07-08
239 // (nx_matmul_t_pool_gate): serial mmt_range = 0.755s/token = ~22%
240 // of the token; threaded on the shared pool = 0.143s (5.29x). The
241 // pooled path is bit-exact vs serial (flat-range bands, identical
242 // per-cell order). One shared process pool (nx_lw_shared_pool),
243 // single-submitter (the forward drives from one thread).
244 if model.lm_head_q8 != 0 {
245 // SIMD Q8_0 dequant-dot lm_head (10.2x, reads 1 B/val vs 8 B/val F32).
246 nx_f32_lazy_matmul(x_normed, model.lm_head_q8 as *NxF32LazyWeight, logits,
247 n_tokens, hidden_dim, vocab_size)
248 } else {
249 nx_f32_matmul_t_pool(nx_lw_shared_pool(), x_normed, model.lm_head, logits,
250 n_tokens, hidden_dim, vocab_size)
251 }
252
253 return NX_FLV4_OK
254}
255
256// PUBLIC GENERATION ENTRY -- byte-for-byte the contract the 58 existing call sites already use.
257func nx_f32_llm_forward_v4(
258 model: *NxF32LlamaModel,
259 token_ids: *i64,
260 n_tokens: nx_int,
261 cache: *NxF32KVCache,
262 eps: i64,
263 attn_scale: i64,
264 rope_log_base: i64,
265 apply_rope: nx_int,
266 logits: *i64) -> nx_int {
267 return _flv4_core(model, token_ids, n_tokens, cache, eps, attn_scale, rope_log_base,
268 apply_rope, logits, 0 as *i64)
269}
270
271// PUBLIC EMBEDDING ENTRY (rung 1). Fills hidden_out with the SUM-POOLED final-normed hidden state
272// (hidden_dim f32 values) and SKIPS the lm_head matmul entirely. Cosine is scale-invariant, so summing
273// rather than averaging changes no ranking; normalise once at the caller.
274// Returns NX_FLV4_OK, or the same error codes as the generation path.
275func nx_f32_llm_embed_v4(
276 model: *NxF32LlamaModel,
277 token_ids: *i64,
278 n_tokens: nx_int,
279 cache: *NxF32KVCache,
280 eps: i64,
281 attn_scale: i64,
282 rope_log_base: i64,
283 apply_rope: nx_int,
284 hidden_out: *i64) -> nx_int {
285 if hidden_out == (0 as *i64) { return NX_FLV4_ERR_NULL }
286 return _flv4_core(model, token_ids, n_tokens, cache, eps, attn_scale, rope_log_base,
287 apply_rope, 0 as *i64, hidden_out)
288}