nx_gguf_load_block.nx source
↩ module page · 325 lines · 13593 B
1// nx_gguf_load_block.nx -- per-layer Llama-class weight bundle loader.
2//
3// L5 brick: composes nx_gguf_load_tensor + nx_rope_compute_inv_freq +
4// NxTransformerBlockWeights to read the 9 named tensors for one
5// transformer layer from a parsed GGUF and populate the bundle struct
6// the transformer block forward consumes.
7//
8// This is the brick that takes nx_llm_run from v1 "API composition
9// proof" to v2 "real-weight runner" -- caller can now load a Llama
10// Q8_0 GGUF, ask for layer N's weights, and feed them to
11// nx_transformer_block_forward unchanged.
12//
13// ===== Llama GGUF tensor naming convention (public format spec) ===
14//
15// Per ggml docs (Gerganov 2024) -- clean-room implementation from
16// public spec; no copied code.
17//
18// blk.{N}.attn_norm.weight -> gamma_attn (RMS-norm scale, 1-D)
19// blk.{N}.attn_q.weight -> W_q
20// blk.{N}.attn_k.weight -> W_k
21// blk.{N}.attn_v.weight -> W_v
22// blk.{N}.attn_output.weight -> W_o
23// blk.{N}.ffn_norm.weight -> gamma_ffn (RMS-norm scale, 1-D)
24// blk.{N}.ffn_gate.weight -> W_gate
25// blk.{N}.ffn_up.weight -> W_up
26// blk.{N}.ffn_down.weight -> W_down
27//
28// inv_freq is COMPUTED from rope_base via nx_rope_compute_inv_freq,
29// not loaded -- GGUF stores rope_freq_base as metadata, not as a
30// tensor. Caller supplies the base.
31//
32// Bits-up composition:
33// nx_gguf.nx -- header + tensor_info table
34// nx_gguf_load.nx -- per-tensor load + dequant (F32/F16/Q8_0)
35// nx_strconv.nx -- decimal layer-index formatting
36// nx_tensor.nx -- NxTensor return shape
37// nx_rope.nx -- inv_freq computation
38// nx_transformer_block.nx -- NxTransformerBlockWeights struct
39//
40// genealogy_id: llama_gguf_naming_gerganov_2024 +
41// touvron_2023_llama + standard_rope_inv_freq
42// lineage_id: substrate_gguf_block_loader_v1_llama_shape
43
44// nx_safety_envelope:
45// intended_use: "Populate one NxTransformerBlockWeights
46// bundle from a parsed Llama-shape GGUF;
47// single-layer scope"
48// sil_target: SIL2
49// asil_target: QM
50// dal_target: DAL C
51// evidence: [gguf_naming_spec_public, clean_room,
52// composes_only_shipped_canonical_bricks]
53// hazard_register: [bug-tape-name-collision-on-malformed-gguf,
54// bug-tape-dim-mismatch-vs-model-spec,
55// bug-tape-layer-index-overflow]
56// residual_risk: "Caller must supply correct hidden_dim /
57// head_dim / ffn_dim matching the GGUF; this
58// brick does not parse GGUF metadata to derive
59// the spec (queued for metadata-walker brick)"
60// verdict: NOT_YET_EVALUATED
61
62import "nx_syscalls.nx"
63import "nx_tier.nx"
64import "nx_loop.nx"
65import "nx_strconv.nx"
66import "nx_tensor.nx"
67import "nx_rope.nx"
68import "nx_gguf.nx"
69import "nx_gguf_load.nx"
70import "nx_transformer_block.nx"
71
72// ===== Sealed-enum: BlockLoadVerdict ==============================
73
74const NX_GBL_OK: nx_int = 0
75const NX_GBL_ERR_NOT_FOUND: nx_int = 1
76const NX_GBL_ERR_BAD_LAYER: nx_int = 2
77const NX_GBL_ERR_BAD_DIM: nx_int = 3
78const NX_GBL_ERR_OOM: nx_int = 4
79const NX_GBL_ERR_LOAD_FAIL: nx_int = 5
80const NX_GBL_ERR_NAME_TOO_BIG: nx_int = 6
81const NX_GBL_ERR_ROPE_FAIL: nx_int = 7
82const NX_GBL_N_VERDICTS: nx_int = 8
83
84func nx_gbl_verdict_is_valid(v: nx_int) -> nx_int {
85 if v < 0 { return 0 }
86 if v >= NX_GBL_N_VERDICTS { return 0 }
87 return 1
88}
89
90// Max tensor-name length we ever need to format. "blk." (4) +
91// max i64 decimal (20) + "." (1) + longest suffix "attn_output.weight"
92// (18) + safety margin = 64. GGUF v3 already caps NX_GGUF_MAX_NAME
93// at 64; we honor the same bound.
94const NX_GBL_MAX_NAME: nx_int = 64
95
96// ===== Helper: format "blk.{N}.{suffix}" ==========================
97//
98// Writes "blk.{layer_idx}.{suffix}" into out_name and returns the
99// total byte count written (NOT including any trailing 0; caller
100// supplies len explicitly so 0-terminator is not required). Returns
101// negative verdict on overflow.
102
103func _gbl_fmt_blk_name(layer_idx: nx_int,
104 suffix: *u8, suffix_len: nx_int,
105 out_name: *u8) -> nx_int {
106 // "blk." = 4 bytes
107 out_name[0] = 0x62 // 'b'
108 out_name[1] = 0x6c // 'l'
109 out_name[2] = 0x6b // 'k'
110 out_name[3] = 0x2e // '.'
111 var off: nx_int = 4
112
113 // Layer index in decimal.
114 let n_dec: i64 = nx_strconv_format_i64(layer_idx, out_name + off)
115 off = off + n_dec
116
117 // Bounds check.
118 if off + 1 + suffix_len > NX_GBL_MAX_NAME {
119 return 0 - NX_GBL_ERR_NAME_TOO_BIG
120 }
121
122 // '.'
123 out_name[off] = 0x2e
124 off = off + 1
125
126 // Suffix bytes (literal copy).
127 var i: nx_int = 0
128 var iter: nx_int = 0
129 var verdict: nx_int = NX_LOOP_RUNNING
130 let BUDGET: nx_int = suffix_len
131 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
132 out_name[off + i] = suffix[i]
133 i = i + 1
134 iter = iter + 1
135 }
136 off = off + suffix_len
137
138 return off
139}
140
141// ===== Helper: load one named per-layer tensor =====================
142//
143// Composes _gbl_fmt_blk_name + nx_gguf_load_tensor for the common
144// pattern of "blk.{N}.{suffix}".
145
146func _gbl_load_blk_tensor(buf: *u8, hdr: *NxGgufHeader,
147 layer_idx: nx_int,
148 suffix: *u8, suffix_len: nx_int,
149 out_err: *i64) -> *NxTensor {
150 let name: *u8 = sys_mmap(NX_GBL_MAX_NAME)
151 let name_len: nx_int = _gbl_fmt_blk_name(layer_idx, suffix, suffix_len, name)
152 if name_len < 0 {
153 out_err[0] = NX_GBL_ERR_NAME_TOO_BIG
154 return 0 as *NxTensor
155 }
156 let t: *NxTensor = nx_gguf_load_tensor(buf, hdr, name, name_len, out_err)
157 return t
158}
159
160// ===== Public: load one transformer block's weights ===============
161//
162// Composes 9 nx_gguf_load_tensor calls + 1 nx_rope_compute_inv_freq
163// into a populated NxTransformerBlockWeights. Caller pre-allocates
164// the bundle struct + provides the model spec dims.
165//
166// rope_base is the RoPE theta base (Llama-2: 10000; Llama-3: 500000).
167// Caller supplies; this brick does not parse GGUF metadata to derive
168// it (that walker is a separate brick).
169//
170// Args (8 -- under the 16-arg NishiLang limit):
171// buf, hdr, layer_idx, head_dim, rope_base, w_out, out_err
172// (hidden_dim / ffn_dim are not strictly needed by THIS brick --
173// they are validated against loaded tensor shapes downstream by
174// the transformer-block forward; we keep the signature compact)
175
176func nx_gguf_load_block_weights(buf: *u8, hdr: *NxGgufHeader,
177 layer_idx: nx_int,
178 head_dim: nx_int,
179 rope_base: nx_int,
180 w_out: *NxTransformerBlockWeights,
181 out_err: *i64) -> nx_int {
182 if layer_idx < 0 {
183 out_err[0] = NX_GBL_ERR_BAD_LAYER
184 return NX_GBL_ERR_BAD_LAYER
185 }
186 if head_dim <= 0 {
187 out_err[0] = NX_GBL_ERR_BAD_DIM
188 return NX_GBL_ERR_BAD_DIM
189 }
190 if rope_base <= 1 {
191 out_err[0] = NX_GBL_ERR_BAD_DIM
192 return NX_GBL_ERR_BAD_DIM
193 }
194
195 // ----- Suffix constants ------------------------------------------
196 //
197 // We pre-allocate each suffix as a mmap'd byte string. Each is
198 // an ASCII literal stored once; reused across layers if caller
199 // batches calls (each call here allocs its own copy -- v1 cost).
200
201 let s_attn_norm: *u8 = sys_mmap(15)
202 s_attn_norm[0]=0x61; s_attn_norm[1]=0x74; s_attn_norm[2]=0x74; s_attn_norm[3]=0x6e
203 s_attn_norm[4]=0x5f; s_attn_norm[5]=0x6e; s_attn_norm[6]=0x6f; s_attn_norm[7]=0x72
204 s_attn_norm[8]=0x6d; s_attn_norm[9]=0x2e; s_attn_norm[10]=0x77; s_attn_norm[11]=0x65
205 s_attn_norm[12]=0x69; s_attn_norm[13]=0x67; s_attn_norm[14]=0x68
206 // "attn_norm.weight" -- but length 16; rewrite as 16 bytes.
207
208 // Cleaner: allocate the full string each time with a byte-by-byte
209 // pattern. We write each suffix inline below to avoid the off-by-one
210 // mistake just made. Functions hold one suffix; alloc + write +
211 // load + record.
212
213 // ----- attn_norm.weight (16 bytes) -----
214 let suf_an: *u8 = sys_mmap(16)
215 suf_an[0]=0x61; suf_an[1]=0x74; suf_an[2]=0x74; suf_an[3]=0x6e
216 suf_an[4]=0x5f; suf_an[5]=0x6e; suf_an[6]=0x6f; suf_an[7]=0x72
217 suf_an[8]=0x6d; suf_an[9]=0x2e; suf_an[10]=0x77; suf_an[11]=0x65
218 suf_an[12]=0x69; suf_an[13]=0x67; suf_an[14]=0x68; suf_an[15]=0x74
219 let t_an: *NxTensor = _gbl_load_blk_tensor(buf, hdr, layer_idx,
220 suf_an, 16, out_err)
221 if out_err[0] != NX_GL_OK {
222 out_err[0] = NX_GBL_ERR_NOT_FOUND
223 return NX_GBL_ERR_NOT_FOUND
224 }
225 w_out.gamma_attn = t_an.storage as *i64
226
227 // ----- attn_q.weight (13 bytes) -----
228 let suf_q: *u8 = sys_mmap(13)
229 suf_q[0]=0x61; suf_q[1]=0x74; suf_q[2]=0x74; suf_q[3]=0x6e
230 suf_q[4]=0x5f; suf_q[5]=0x71; suf_q[6]=0x2e; suf_q[7]=0x77
231 suf_q[8]=0x65; suf_q[9]=0x69; suf_q[10]=0x67; suf_q[11]=0x68
232 suf_q[12]=0x74
233 w_out.W_q = _gbl_load_blk_tensor(buf, hdr, layer_idx,
234 suf_q, 13, out_err)
235 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
236
237 // ----- attn_k.weight (13 bytes) -----
238 let suf_k: *u8 = sys_mmap(13)
239 suf_k[0]=0x61; suf_k[1]=0x74; suf_k[2]=0x74; suf_k[3]=0x6e
240 suf_k[4]=0x5f; suf_k[5]=0x6b; suf_k[6]=0x2e; suf_k[7]=0x77
241 suf_k[8]=0x65; suf_k[9]=0x69; suf_k[10]=0x67; suf_k[11]=0x68
242 suf_k[12]=0x74
243 w_out.W_k = _gbl_load_blk_tensor(buf, hdr, layer_idx,
244 suf_k, 13, out_err)
245 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
246
247 // ----- attn_v.weight (13 bytes) -----
248 let suf_v: *u8 = sys_mmap(13)
249 suf_v[0]=0x61; suf_v[1]=0x74; suf_v[2]=0x74; suf_v[3]=0x6e
250 suf_v[4]=0x5f; suf_v[5]=0x76; suf_v[6]=0x2e; suf_v[7]=0x77
251 suf_v[8]=0x65; suf_v[9]=0x69; suf_v[10]=0x67; suf_v[11]=0x68
252 suf_v[12]=0x74
253 w_out.W_v = _gbl_load_blk_tensor(buf, hdr, layer_idx,
254 suf_v, 13, out_err)
255 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
256
257 // ----- attn_output.weight (18 bytes) -----
258 let suf_o: *u8 = sys_mmap(18)
259 suf_o[0]=0x61; suf_o[1]=0x74; suf_o[2]=0x74; suf_o[3]=0x6e
260 suf_o[4]=0x5f; suf_o[5]=0x6f; suf_o[6]=0x75; suf_o[7]=0x74
261 suf_o[8]=0x70; suf_o[9]=0x75; suf_o[10]=0x74; suf_o[11]=0x2e
262 suf_o[12]=0x77; suf_o[13]=0x65; suf_o[14]=0x69; suf_o[15]=0x67
263 suf_o[16]=0x68; suf_o[17]=0x74
264 w_out.W_o = _gbl_load_blk_tensor(buf, hdr, layer_idx,
265 suf_o, 18, out_err)
266 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
267
268 // ----- ffn_norm.weight (15 bytes) -----
269 let suf_fn: *u8 = sys_mmap(15)
270 suf_fn[0]=0x66; suf_fn[1]=0x66; suf_fn[2]=0x6e; suf_fn[3]=0x5f
271 suf_fn[4]=0x6e; suf_fn[5]=0x6f; suf_fn[6]=0x72; suf_fn[7]=0x6d
272 suf_fn[8]=0x2e; suf_fn[9]=0x77; suf_fn[10]=0x65; suf_fn[11]=0x69
273 suf_fn[12]=0x67; suf_fn[13]=0x68; suf_fn[14]=0x74
274 let t_fn: *NxTensor = _gbl_load_blk_tensor(buf, hdr, layer_idx,
275 suf_fn, 15, out_err)
276 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
277 w_out.gamma_ffn = t_fn.storage as *i64
278
279 // ----- ffn_gate.weight (15 bytes) -----
280 let suf_fg: *u8 = sys_mmap(15)
281 suf_fg[0]=0x66; suf_fg[1]=0x66; suf_fg[2]=0x6e; suf_fg[3]=0x5f
282 suf_fg[4]=0x67; suf_fg[5]=0x61; suf_fg[6]=0x74; suf_fg[7]=0x65
283 suf_fg[8]=0x2e; suf_fg[9]=0x77; suf_fg[10]=0x65; suf_fg[11]=0x69
284 suf_fg[12]=0x67; suf_fg[13]=0x68; suf_fg[14]=0x74
285 w_out.W_gate = _gbl_load_blk_tensor(buf, hdr, layer_idx,
286 suf_fg, 15, out_err)
287 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
288
289 // ----- ffn_up.weight (13 bytes) -----
290 let suf_fu: *u8 = sys_mmap(13)
291 suf_fu[0]=0x66; suf_fu[1]=0x66; suf_fu[2]=0x6e; suf_fu[3]=0x5f
292 suf_fu[4]=0x75; suf_fu[5]=0x70; suf_fu[6]=0x2e; suf_fu[7]=0x77
293 suf_fu[8]=0x65; suf_fu[9]=0x69; suf_fu[10]=0x67; suf_fu[11]=0x68
294 suf_fu[12]=0x74
295 w_out.W_up = _gbl_load_blk_tensor(buf, hdr, layer_idx,
296 suf_fu, 13, out_err)
297 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
298
299 // ----- ffn_down.weight (15 bytes) -----
300 let suf_fd: *u8 = sys_mmap(15)
301 suf_fd[0]=0x66; suf_fd[1]=0x66; suf_fd[2]=0x6e; suf_fd[3]=0x5f
302 suf_fd[4]=0x64; suf_fd[5]=0x6f; suf_fd[6]=0x77; suf_fd[7]=0x6e
303 suf_fd[8]=0x2e; suf_fd[9]=0x77; suf_fd[10]=0x65; suf_fd[11]=0x69
304 suf_fd[12]=0x67; suf_fd[13]=0x68; suf_fd[14]=0x74
305 w_out.W_down = _gbl_load_blk_tensor(buf, hdr, layer_idx,
306 suf_fd, 15, out_err)
307 if out_err[0] != NX_GL_OK { return NX_GBL_ERR_NOT_FOUND }
308
309 // ----- inv_freq (computed, NOT loaded) ------------------------
310 //
311 // RoPE inv_freq depends on head_dim and rope_base. GGUF stores
312 // rope_freq_base as metadata, not as a tensor; we compute from
313 // head_dim + caller-supplied base.
314
315 let inv_freq: *i64 = sys_mmap((head_dim / 2) * 8) as *i64
316 let v_rope: nx_int = nx_rope_compute_inv_freq(head_dim, rope_base, inv_freq)
317 if v_rope != NX_ROPE_OK {
318 out_err[0] = NX_GBL_ERR_ROPE_FAIL
319 return NX_GBL_ERR_ROPE_FAIL
320 }
321 w_out.inv_freq = inv_freq
322
323 out_err[0] = NX_GBL_OK
324 return NX_GBL_OK
325}