nx_genweights.nx source
↩ module page · 290 lines · 11217 B
1// nx_genweights.nx -- read model weights straight from a GGUF. The engine's weight source.
2//
3// Until now this lane read weights from fixtures the oracle tap had dumped. That was right for
4// building a RULER -- the tap's f32 conversion is ggml's own dequantization, which is exactly what
5// a differential test needs. It is wrong for an ENGINE: it means the sovereign path cannot run
6// without first booting a third-party binary to export its weights, and it does not scale (one
7// Z-Image layer is ~190MB of raw blocks; 36 layers is ~7GB of fixtures nobody should be dumping).
8//
9// ★ A RULER MAY DEPEND ON THE INCUMBENT. A RUNTIME MAY NOT.
10//
11// The GGUF's tensor names turn out to be EXACTLY the names sd.cpp uses internally
12// ("model.diffusion_model.layers.0.attention.qkv.weight"), so this is a drop-in swap for the
13// fixture loader -- verified by running the same block both ways and comparing.
14//
15// The file is mmapped whole (sys_map_file), not read: a 6.2GB model is paged in on demand and
16// only the tensors actually touched ever reach RAM.
17//
18// HANDLE LAYOUT (i64 slots): [0]=mapped base, [1]=file length, [2]=*NxGgufHeader
19// license_tier: ORIGINAL
20
21import "nx_syscalls.nx"
22import "nx_le.nx"
23import "nx_f32.nx"
24import "nx_tensor.nx"
25import "nx_gguf.nx"
26import "nx_gguf_load.nx"
27import "nx_q8_0_to_f32.nx"
28import "nx_safetensors_load.nx"
29
30const NX_GW_TYPE_F32: i64 = 0
31const NX_GW_TYPE_F16: i64 = 1
32const NX_GW_TYPE_Q8_0: i64 = 8
33
34// Backend tag in slot 3. One API, two container formats: GGUF (the quantized release format) and
35// safetensors (what community checkpoints ship as). Without both, "hot-swappable" means
36// "swappable for another file we happened to convert".
37const NX_GW_FMT_GGUF: i64 = 0
38const NX_GW_FMT_ST: i64 = 1
39
40// safetensors dtype tags, kept distinct from ggml's numbering
41const NX_GW_ST_F32: i64 = 100
42const NX_GW_ST_F16: i64 = 101
43const NX_GW_ST_BF16: i64 = 102
44const NX_GW_ST_F8E4M3: i64 = 103
45
46// FP8 E4M3 -> f32. 1 sign, 4 exponent (bias 7), 3 mantissa; no infinities, and exp=15,mant=7 is
47// the only NaN. Subnormals are real here (exp=0) and carry weight values, so they are normalized
48// properly rather than flushed -- flushing them would quietly zero part of the model.
49func nx_gw_f8e4m3_to_f32(b: i64) -> i64 {
50 let s: i64 = (b >> 7) & 1
51 let e: i64 = (b >> 3) & 0xF
52 let m: i64 = b & 7
53 if e == 0 {
54 if m == 0 { return s << 31 }
55 // value = m * 2^-9, normalized: m = 1.xxx * 2^k
56 var k: i64 = 0
57 if m >= 4 { k = 2 } else { if m >= 2 { k = 1 } }
58 let mant: i64 = (m - (1 << k)) << (23 - k)
59 return (s << 31) | ((k + 118) << 23) | mant
60 }
61 if e == 15 {
62 if m == 7 { return (s << 31) | 0x7FC00000 }
63 }
64 return (s << 31) | ((e + 120) << 23) | (m << 20)
65}
66
67func nx_gw_open(path: *u8) -> *i64 {
68 let lenp: *i64 = sys_mmap(32) as *i64
69 lenp[0] = 0
70 let base: *u8 = sys_map_file(path, lenp)
71 if (base as i64) == 0 { return 0 as *i64 }
72
73 // Format by MAGIC, not by filename extension -- an extension is a claim, a magic is evidence.
74 if base[0] != (0x47 as u8) || base[1] != (0x47 as u8) || base[2] != (0x55 as u8) || base[3] != (0x46 as u8) {
75 let hl: i64 = stl_header_len(base)
76 if hl <= 0 { return 0 as *i64 }
77 if 8 + hl > lenp[0] { return 0 as *i64 }
78 if base[8] != (0x7B as u8) { return 0 as *i64 } // header must open with '{'
79 let g2: *i64 = sys_mmap(64) as *i64
80 g2[0] = base as i64
81 g2[1] = lenp[0]
82 g2[2] = 0
83 g2[3] = NX_GW_FMT_ST
84 g2[4] = hl
85 g2[5] = stl_data_start(base)
86 return g2
87 }
88 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
89 // The metadata block sits at the front, so parsing over the whole mapping touches only the
90 // first pages -- no 6.2GB read.
91 let v: nx_int = nx_gguf_parse(base, lenp[0], hdr)
92 if v != NX_GGUF_OK { return 0 as *i64 }
93 if hdr.version != 3 { return 0 as *i64 }
94 let gw: *i64 = sys_mmap(64) as *i64
95 gw[0] = base as i64
96 gw[1] = lenp[0]
97 gw[2] = hdr as i64
98 gw[3] = NX_GW_FMT_GGUF
99 return gw
100}
101
102// ===== safetensors helpers ==========================================
103//
104// Locate a tensor and read dtype + shape + byte offsets. `stl_tensor` gives dtype and offsets but
105// not shape, and the engine needs shape to derive the architecture, so the shape array is parsed
106// here.
107//
108// ⚠SHAPE ORDER IS REVERSED BETWEEN THE TWO FORMATS. safetensors stores torch row-major
109// [out_features, in_features]; GGUF ne is [in, out] with ne0 the contiguous dim. This API returns
110// GGUF semantics for BOTH -- dim0 = INPUT width -- so callers do not have to know which container
111// they were handed. Getting this backwards transposes every weight and still produces a picture.
112func _gw_st_find(gw: *i64, name: *u8, nl: i64, out: *i64) -> i64 {
113 let base: *u8 = gw[0] as *u8
114 let hend: i64 = 8 + gw[4]
115 let np: i64 = stl_find(base, hend, 8, name)
116 if np < 0 { return 0 - 1 }
117 let dp: i64 = stl_find(base, hend, np, "\"dtype\":\"" as *u8)
118 if dp < 0 { return 0 - 2 }
119 var i: i64 = dp + 9
120 let dt: *u8 = sys_mmap(32)
121 var k: i64 = 0
122 while base[i] != (34 as u8) { dt[k] = base[i]; i = i + 1; k = k + 1 }
123 dt[k] = 0 as u8
124 var ty: i64 = 0 - 1
125 if stl_streq(dt, "F32" as *u8) == 1 { ty = NX_GW_ST_F32 }
126 if stl_streq(dt, "F16" as *u8) == 1 { ty = NX_GW_ST_F16 }
127 if stl_streq(dt, "BF16" as *u8) == 1 { ty = NX_GW_ST_BF16 }
128 if stl_streq(dt, "F8_E4M3" as *u8) == 1 { ty = NX_GW_ST_F8E4M3 }
129 if ty < 0 { return 0 - 3 }
130
131 let sp: i64 = stl_find(base, hend, np, "\"shape\":[" as *u8)
132 if sp < 0 { return 0 - 4 }
133 let ep: *i64 = sys_mmap(16) as *i64
134 var p: i64 = sp + 9
135 let d0: i64 = stl_parse_int(base, p, ep)
136 var d1: i64 = 1
137 p = ep[0]
138 while base[p] == (32 as u8) { p = p + 1 }
139 if base[p] == (44 as u8) {
140 p = p + 1
141 while base[p] == (32 as u8) { p = p + 1 }
142 d1 = stl_parse_int(base, p, ep)
143 }
144
145 let op: i64 = stl_find(base, hend, np, "\"data_offsets\":[" as *u8)
146 if op < 0 { return 0 - 5 }
147 let o0: i64 = stl_parse_int(base, op + 16, ep)
148 p = ep[0]
149 while base[p] < (48 as u8) { p = p + 1 }
150 let o1: i64 = stl_parse_int(base, p, ep)
151
152 out[0] = ty
153 // REVERSE to GGUF semantics: safetensors [out, in] -> dim0 = in, dim1 = out
154 if d1 == 1 { out[1] = d0; out[2] = 1 } else { out[1] = d1; out[2] = d0 }
155 out[3] = o0
156 out[4] = o1
157 return 0
158}
159
160func nx_gw_ntensors(gw: *i64) -> i64 {
161 if gw[3] == NX_GW_FMT_ST { return 0 - 1 } // safetensors has no index; unknown, not zero
162 let hdr: *NxGgufHeader = gw[2] as *NxGgufHeader
163 return hdr.n_tensors
164}
165
166// For safetensors there is no index array, so "find" returns 1 when present and -1 when not;
167// callers only test the sign. GGUF keeps returning its real index.
168func nx_gw_find(gw: *i64, name: *u8, nl: i64) -> i64 {
169 if gw[3] == NX_GW_FMT_ST {
170 let inf: *i64 = sys_mmap(64) as *i64
171 if _gw_st_find(gw, name, nl, inf) != 0 { return 0 - 1 }
172 gw[6] = name as i64
173 gw[7] = nl
174 return 1
175 }
176 let hdr: *NxGgufHeader = gw[2] as *NxGgufHeader
177 return nx_gguf_find_tensor(hdr, name, nl as nx_int)
178}
179
180// safetensors lookups are by NAME, so these re-resolve from the last name passed to nx_gw_find.
181func _gw_st_last(gw: *i64, out: *i64) -> i64 {
182 return _gw_st_find(gw, gw[6] as *u8, gw[7], out)
183}
184
185func _gw_info(gw: *i64, idx: i64) -> *NxGgufTensorInfo {
186 let hdr: *NxGgufHeader = gw[2] as *NxGgufHeader
187 return nx_gguf_tensor_at(hdr, idx as nx_int)
188}
189
190// Pointer to the tensor's bytes inside the mapping, in whatever type the file stores.
191func nx_gw_data(gw: *i64, idx: i64) -> *u8 {
192 if gw[3] == NX_GW_FMT_ST {
193 let inf: *i64 = sys_mmap(64) as *i64
194 if _gw_st_last(gw, inf) != 0 { return 0 as *u8 }
195 return (gw[0] + gw[5] + inf[3]) as *u8
196 }
197 let hdr: *NxGgufHeader = gw[2] as *NxGgufHeader
198 let ti: *NxGgufTensorInfo = _gw_info(gw, idx)
199 return (gw[0] + hdr.data_off + ti.offset) as *u8
200}
201
202func nx_gw_type(gw: *i64, idx: i64) -> i64 {
203 if gw[3] == NX_GW_FMT_ST {
204 let inf: *i64 = sys_mmap(64) as *i64
205 if _gw_st_last(gw, inf) != 0 { return 0 - 1 }
206 return inf[0]
207 }
208 let ti: *NxGgufTensorInfo = _gw_info(gw, idx)
209 return ti.ggml_type
210}
211func nx_gw_dim0(gw: *i64, idx: i64) -> i64 {
212 if gw[3] == NX_GW_FMT_ST {
213 let inf: *i64 = sys_mmap(64) as *i64
214 if _gw_st_last(gw, inf) != 0 { return 0 - 1 }
215 return inf[1]
216 }
217 let ti: *NxGgufTensorInfo = _gw_info(gw, idx)
218 return ti.dim_0
219}
220func nx_gw_dim1(gw: *i64, idx: i64) -> i64 {
221 if gw[3] == NX_GW_FMT_ST {
222 let inf: *i64 = sys_mmap(64) as *i64
223 if _gw_st_last(gw, inf) != 0 { return 0 - 1 }
224 return inf[2]
225 }
226 let ti: *NxGgufTensorInfo = _gw_info(gw, idx)
227 return ti.dim_1
228}
229
230// Materialize a tensor as PACKED 4-byte f32 (the layout every kernel here consumes).
231//
232// Needed for the small tensors -- norms, biases -- which this engine wants in f32 even when the
233// file stores them quantized. sd.cpp does the same: RMSNorm::init_params always allocates F32
234// regardless of the file type, so reading the raw Q8_0 blocks for a norm and using them directly
235// would NOT match the engine.
236// Returns 0 on success, or a negative verdict for an unsupported type -- never a silent partial.
237func nx_gw_to_f32_packed(gw: *i64, idx: i64, out: *u8, n_values: i64) -> i64 {
238 let ty: i64 = nx_gw_type(gw, idx)
239 let src: *u8 = nx_gw_data(gw, idx)
240 if (src as i64) == 0 { return 0 - 5 }
241 if ty == NX_GW_ST_F32 {
242 var i: i64 = 0
243 while i < n_values { nx_le_write_u32(out, i * 4, nx_le_read_u32(src, i * 4)); i = i + 1 }
244 return 0
245 }
246 if ty == NX_GW_ST_F16 {
247 var i: i64 = 0
248 while i < n_values { nx_le_write_u32(out, i * 4, nx_f16_to_f32(nx_le_read_u16(src, i * 2))); i = i + 1 }
249 return 0
250 }
251 if ty == NX_GW_ST_BF16 {
252 // bf16 IS the high 16 bits of an f32 -- a shift, not a conversion.
253 var i: i64 = 0
254 while i < n_values { nx_le_write_u32(out, i * 4, nx_le_read_u16(src, i * 2) << 16); i = i + 1 }
255 return 0
256 }
257 if ty == NX_GW_ST_F8E4M3 {
258 var i: i64 = 0
259 while i < n_values { nx_le_write_u32(out, i * 4, nx_gw_f8e4m3_to_f32(src[i] & 0xFF)); i = i + 1 }
260 return 0
261 }
262 if ty == NX_GW_TYPE_F32 {
263 var i: i64 = 0
264 while i < n_values {
265 nx_le_write_u32(out, i * 4, nx_le_read_u32(src, i * 4))
266 i = i + 1
267 }
268 return 0
269 }
270 if ty == NX_GW_TYPE_F16 {
271 var i: i64 = 0
272 while i < n_values {
273 nx_le_write_u32(out, i * 4, nx_f16_to_f32(nx_le_read_u16(src, i * 2)))
274 i = i + 1
275 }
276 return 0
277 }
278 if ty == NX_GW_TYPE_Q8_0 {
279 let tmp: *i64 = sys_mmap(n_values * 8 + 64) as *i64
280 let v: nx_int = nx_q8_0_to_f32(src, 0, n_values, tmp)
281 if v != 0 { return 0 - 2 }
282 var i: i64 = 0
283 while i < n_values {
284 nx_le_write_u32(out, i * 4, tmp[i])
285 i = i + 1
286 }
287 return 0
288 }
289 return 0 - 1
290}