nxgguf.nx source
↩ module page · 535 lines · 19854 B
1// nxgguf.nx -- Nishi sovereign tensor weight format.
2//
3// Replaces HuggingFace safetensors / GGUF with content-addressed,
4// memory-mapped, dedupable weight blocks. Single-binary deploy of
5// any model; instant cold-load via mmap; LoRAs share base weights
6// at the storage layer.
7//
8// Format overview:
9//
10// [magic 4] = "NXGF"
11// [version 4] = u32 LE, currently 1
12// [tensor_count 8] = u64 LE, total tensors
13// [metadata_offset 8] = u64 LE, offset to metadata block
14// [metadata_length 8] = u64 LE, byte length of metadata
15// [block_table_offset 8] = u64 LE, offset to block dir
16// [block_table_length 8] = u64 LE, byte length of block dir
17// [data_offset 8] = u64 LE, where weight bytes start
18// [reserved 16] = zeros
19// = 64 bytes total header
20//
21// metadata block: protobuf-encoded TensorRegistry
22// repeated TensorEntry {
23// string name = 1; // "model.layers.0.attn.q.weight"
24// repeated int64 shape = 2; // [512, 512]
25// int32 dtype = 3; // 0=fp32 1=fp16 2=bf16 3=fp8 4=int8 5=int4
26// string block_hash = 4; // sha256 of the tensor's bytes
27// int64 byte_length = 5; // raw byte count
28// }
29//
30// block table: array of BlockDescriptor
31// bytes hash (32) // sha256
32// uint64 file_offset ( 8) // where in file
33// uint64 byte_length ( 8)
34// uint32 ref_count ( 4) // # tensors pointing here
35// uint32 reserved ( 4)
36// = 56 bytes per block descriptor
37//
38// data section: concatenated raw weight bytes; tensor reads
39// `block.byte_length` bytes starting at `block.file_offset`.
40//
41// Dedup mechanism: if two tensors have the same content, they share
42// one BlockDescriptor (block_hash matches). ref_count tracks usage.
43// LoRAs typically share 95%+ of base model weights this way.
44//
45// mmap path: caller mmaps the entire file, then nxgguf_open() walks
46// the header and metadata in-place. Tensor reads return pointers
47// into the mmap region -- zero-copy load.
48
49import "syscalls.nx"
50
51// === constants ===
52
53// "NXGF" stored little-endian: bytes 'N'(0x4E),'X'(0x58),'G'(0x47),'F'(0x46)
54// → u32 value with 'F' as MSB, 'N' as LSB = 0x4647_584E
55const NXGGUF_MAGIC: i64 = 0x4647584E
56const NXGGUF_VERSION: i64 = 1
57const NXGGUF_HEADER_BYTES: i64 = 64
58const NXGGUF_BLOCK_DESC_BYTES: i64 = 56
59
60// dtype enum (matches GGUF's dtype IDs where overlapping)
61const DTYPE_FP32: i64 = 0
62const DTYPE_FP16: i64 = 1
63const DTYPE_BF16: i64 = 2
64const DTYPE_FP8: i64 = 3 // E4M3
65const DTYPE_INT8: i64 = 4
66const DTYPE_INT4: i64 = 5 // group-quantized
67const DTYPE_INT2: i64 = 6 // for BitNet b1.58 + similar
68const DTYPE_TERN: i64 = 7 // ternary {-1, 0, 1} packed
69
70// element size in bits (for memory accounting)
71func dtype_bits(dt: i64) -> i64 {
72 if dt == DTYPE_FP32 { return 32 }
73 if dt == DTYPE_FP16 { return 16 }
74 if dt == DTYPE_BF16 { return 16 }
75 if dt == DTYPE_FP8 { return 8 }
76 if dt == DTYPE_INT8 { return 8 }
77 if dt == DTYPE_INT4 { return 4 }
78 if dt == DTYPE_INT2 { return 2 }
79 if dt == DTYPE_TERN { return 2 } // 1.58 bits effective, stored 2
80 return 32
81}
82
83// === in-memory representation ===
84
85struct NxgFile {
86 base: *u8, // mmap'd file start
87 length: i64, // mmap'd byte length
88 tensor_count: i64,
89 metadata_off: i64,
90 metadata_len: i64,
91 block_table_off: i64,
92 block_table_len: i64,
93 data_off: i64,
94}
95
96struct NxgTensor {
97 file: *NxgFile, // back-pointer to owner
98 name_off: i64, // offset into metadata block where name lives
99 name_len: i64, // bytes
100 rank: i64, // number of dims
101 shape: *i64, // pointer into metadata where dims live
102 dtype: i64,
103 block_hash_off: i64, // offset into metadata where 32-byte hash lives
104 byte_length: i64,
105 data_ptr: *u8, // resolved pointer into mmap region; null until resolved
106}
107
108struct NxgBlockDesc {
109 hash_off: i64, // offset into block table where 32-byte hash lives
110 file_offset: i64, // where the block's bytes are in the file
111 byte_length: i64,
112 ref_count: i64,
113}
114
115// === reader ===
116
117// Decode little-endian u64 at `buf[off..off+8]`.
118func read_u64_le(buf: *u8, off: i64) -> i64 {
119 var v: i64 = 0
120 var i: i64 = 7
121 while i >= 0 {
122 v = (v << 8) | buf[off + i]
123 i = i - 1
124 }
125 return v
126}
127
128func read_u32_le(buf: *u8, off: i64) -> i64 {
129 var v: i64 = 0
130 var i: i64 = 3
131 while i >= 0 {
132 v = (v << 8) | buf[off + i]
133 i = i - 1
134 }
135 return v
136}
137
138// Open an NXGF file from mmap'd bytes. Returns 0 on success;
139// negative on error. Validates magic + version + sanity-checks
140// offsets fit within the file.
141func nxgguf_open(file: *NxgFile, base: *u8, length: i64) -> i64 {
142 if length < NXGGUF_HEADER_BYTES { return -1 }
143 let magic: i64 = read_u32_le(base, 0)
144 if magic != NXGGUF_MAGIC { return -2 }
145 let version: i64 = read_u32_le(base, 4)
146 if version != NXGGUF_VERSION { return -3 }
147
148 file.base = base
149 file.length = length
150 file.tensor_count = read_u64_le(base, 8)
151 file.metadata_off = read_u64_le(base, 16)
152 file.metadata_len = read_u64_le(base, 24)
153 file.block_table_off = read_u64_le(base, 32)
154 file.block_table_len = read_u64_le(base, 40)
155 file.data_off = read_u64_le(base, 48)
156
157 // Bounds check.
158 if file.metadata_off + file.metadata_len > length { return -4 }
159 if file.block_table_off + file.block_table_len > length { return -5 }
160 if file.data_off > length { return -6 }
161 return 0
162}
163
164// Total weight bytes -- sum across all unique blocks. Used for VRAM
165// accounting before committing the model to GPU.
166func nxgguf_total_bytes(file: *NxgFile) -> i64 {
167 let n_blocks: i64 = file.block_table_len / NXGGUF_BLOCK_DESC_BYTES
168 var total: i64 = 0
169 var i: i64 = 0
170 while i < n_blocks {
171 let off: i64 = file.block_table_off + i * NXGGUF_BLOCK_DESC_BYTES
172 let len: i64 = read_u64_le(file.base, off + 40)
173 total = total + len
174 i = i + 1
175 }
176 return total
177}
178
179// Number of unique blocks (vs total tensors). Lower = better dedup.
180func nxgguf_unique_blocks(file: *NxgFile) -> i64 {
181 return file.block_table_len / NXGGUF_BLOCK_DESC_BYTES
182}
183
184// === content-addressed block hashing =================================
185//
186// sha256 of a tensor's raw bytes is the block_hash. Two tensors
187// with identical content -> identical hash -> share one
188// BlockDescriptor. This is what makes LoRAs near-free at the
189// storage layer (95%+ of base weights typically share with the
190// base model).
191
192import "sha256.nx"
193
194// Compute sha256(bytes[0..len]) into out (32 bytes).
195func nxgguf_hash_block(bytes: *u8, len: i64, out: *u8) -> i64 {
196 let ctx_raw: *u8 = sys_mmap(128)
197 let ctx: *Sha256 = ctx_raw as *Sha256
198 sha256_init(ctx)
199 sha256_update(ctx, bytes, len)
200 sha256_final(ctx, out)
201 return 0
202}
203
204// Compare two 32-byte sha256 hashes. Returns 0 if equal, 1 if
205// different. Constant-time (sha256 outputs aren't secret but
206// dedup speed matters).
207func nxgguf_hash_eq(a: *u8, b: *u8) -> i64 {
208 var i: i64 = 0
209 var diff: i64 = 0
210 while i < 32 {
211 if a[i] != b[i] { diff = 1 }
212 i = i + 1
213 }
214 if diff == 1 { return 1 }
215 return 0
216}
217
218// Look up a block by its hash in the file's block table. Returns
219// the block index (>= 0) on hit, or -1 on miss. O(N) scan; for
220// production use a hash-keyed index (later phase, ~80 LoC).
221func nxgguf_find_block(file: *NxgFile, target_hash: *u8) -> i64 {
222 let n_blocks: i64 = file.block_table_len / NXGGUF_BLOCK_DESC_BYTES
223 var i: i64 = 0
224 while i < n_blocks {
225 let desc_off: i64 = file.block_table_off + i * NXGGUF_BLOCK_DESC_BYTES
226 let block_hash_ptr: *u8 = (file.base as i64 + desc_off) as *u8
227 if nxgguf_hash_eq(block_hash_ptr, target_hash) == 0 {
228 return i
229 }
230 i = i + 1
231 }
232 return -1
233}
234
235// === write side ======================================================
236//
237// Build an nxgguf file in memory from an array of tensor entries.
238// Caller supplies tensor metadata + raw bytes; encoder hashes each
239// tensor's content, dedupes, packs the format.
240//
241// Usage:
242// let w: *NxgWriter = nxgguf_writer_new(out_buf, out_cap)
243// for each tensor:
244// nxgguf_writer_add(w, name, name_len, rank, shape, dtype,
245// data, byte_length)
246// let total: i64 = nxgguf_writer_finalize(w)
247// sys_write(fd, out_buf, total)
248//
249// The writer maintains an in-memory block table for dedup. Every
250// tensor's content is sha256'd; matching hashes share one
251// BlockDescriptor (ref_count++).
252
253const NXG_MAX_BLOCKS: i64 = 4096
254const NXG_MAX_TENSORS: i64 = 4096
255
256struct NxgWriter {
257 out: *u8,
258 out_cap: i64,
259 n_blocks: i64,
260 n_tensors: i64,
261 block_hashes: *u8,
262 block_lens: *i64,
263 block_data: *u8,
264 block_offs: *i64,
265 block_data_pos: i64,
266 tensor_names: *u8,
267 tensor_ranks: *i64,
268 tensor_shapes: *i64,
269 tensor_dtypes: *i64,
270 tensor_blocks: *i64,
271}
272
273func nxgguf_writer_new(out_buf: *u8, out_cap: i64) -> *NxgWriter {
274 let w_raw: *u8 = sys_mmap(128)
275 let w: *NxgWriter = w_raw as *NxgWriter
276 w.out = out_buf
277 w.out_cap = out_cap
278 w.n_blocks = 0
279 w.n_tensors = 0
280 w.block_hashes = sys_mmap(NXG_MAX_BLOCKS * 32 + 64)
281 w.block_lens = sys_mmap(NXG_MAX_BLOCKS * 8 + 64) as *i64
282 w.block_data = sys_mmap(out_cap)
283 w.block_offs = sys_mmap(NXG_MAX_BLOCKS * 8 + 64) as *i64
284 w.block_data_pos = 0
285 w.tensor_names = sys_mmap(NXG_MAX_TENSORS * 256 + 64)
286 w.tensor_ranks = sys_mmap(NXG_MAX_TENSORS * 8 + 64) as *i64
287 w.tensor_shapes = sys_mmap(NXG_MAX_TENSORS * 8 * 8 + 64) as *i64
288 w.tensor_dtypes = sys_mmap(NXG_MAX_TENSORS * 8 + 64) as *i64
289 w.tensor_blocks = sys_mmap(NXG_MAX_TENSORS * 8 + 64) as *i64
290 return w
291}
292
293// Look up a hash in the writer's block table. Returns block index
294// or -1 if not present.
295func nxgguf_writer_find_block(w: *NxgWriter, target_hash: *u8) -> i64 {
296 var i: i64 = 0
297 while i < w.n_blocks {
298 let h_off: i64 = i * 32
299 let bh: *u8 = (w.block_hashes as i64 + h_off) as *u8
300 if nxgguf_hash_eq(bh, target_hash) == 0 { return i }
301 i = i + 1
302 }
303 return -1
304}
305
306// Add a tensor to the writer. Hashes the data; dedupes against
307// existing blocks; appends a new BlockDescriptor + bytes if novel.
308// Returns 0 on success, -1 on capacity overflow.
309func nxgguf_writer_add(w: *NxgWriter,
310 name: *u8, name_len: i64,
311 rank: i64, shape: *i64, dtype: i64,
312 data: *u8, byte_length: i64) -> i64 {
313 if w.n_tensors >= NXG_MAX_TENSORS { return -1 }
314
315 let hash_buf: *u8 = sys_mmap(64)
316 nxgguf_hash_block(data, byte_length, hash_buf)
317
318 var block_idx: i64 = nxgguf_writer_find_block(w, hash_buf)
319 if block_idx < 0 {
320 if w.n_blocks >= NXG_MAX_BLOCKS { return -1 }
321 block_idx = w.n_blocks
322 let h_off: i64 = block_idx * 32
323 var k: i64 = 0
324 while k < 32 {
325 let dst: *u8 = (w.block_hashes as i64 + h_off + k) as *u8
326 *dst = hash_buf[k]
327 k = k + 1
328 }
329 w.block_lens[block_idx] = byte_length
330 var bi: i64 = 0
331 while bi < byte_length {
332 let dst: *u8 = (w.block_data as i64 + w.block_data_pos + bi) as *u8
333 *dst = data[bi]
334 bi = bi + 1
335 }
336 w.block_offs[block_idx] = w.block_data_pos
337 w.block_data_pos = w.block_data_pos + byte_length
338 w.n_blocks = w.n_blocks + 1
339 }
340
341 let t_idx: i64 = w.n_tensors
342 let name_off: i64 = t_idx * 256
343 var k: i64 = 0
344 while k < name_len {
345 let dst: *u8 = (w.tensor_names as i64 + name_off + k) as *u8
346 *dst = name[k]
347 k = k + 1
348 }
349 let term: *u8 = (w.tensor_names as i64 + name_off + name_len) as *u8
350 *term = 0
351 w.tensor_ranks[t_idx] = rank
352 let shape_off: i64 = t_idx * 8
353 var d: i64 = 0
354 while d < rank {
355 if d >= 8 { return -1 }
356 w.tensor_shapes[shape_off + d] = shape[d]
357 d = d + 1
358 }
359 w.tensor_dtypes[t_idx] = dtype
360 w.tensor_blocks[t_idx] = block_idx
361 w.n_tensors = t_idx + 1
362 return 0
363}
364
365// Encode a u64 little-endian into out[off..off+8].
366func write_u64_le(out: *u8, off: i64, v: i64) -> i64 {
367 out[off + 0] = v & 0xFF
368 out[off + 1] = (v >> 8) & 0xFF
369 out[off + 2] = (v >> 16) & 0xFF
370 out[off + 3] = (v >> 24) & 0xFF
371 out[off + 4] = (v >> 32) & 0xFF
372 out[off + 5] = (v >> 40) & 0xFF
373 out[off + 6] = (v >> 48) & 0xFF
374 out[off + 7] = (v >> 56) & 0xFF
375 return 0
376}
377
378// Finalize: emit header + block table + data section. Skips the
379// metadata block for v0.0.1 (protobuf encoder is a future commit).
380// Returns total bytes written.
381func nxgguf_writer_finalize(w: *NxgWriter) -> i64 {
382 let header_bytes: i64 = NXGGUF_HEADER_BYTES
383 let block_table_bytes: i64 = w.n_blocks * NXGGUF_BLOCK_DESC_BYTES
384 let data_off: i64 = header_bytes + block_table_bytes
385 let total: i64 = data_off + w.block_data_pos
386 if total > w.out_cap { return -1 }
387
388 w.out[0] = 0x4E // 'N'
389 w.out[1] = 0x58 // 'X'
390 w.out[2] = 0x47 // 'G'
391 w.out[3] = 0x46 // 'F'
392 w.out[4] = 1; w.out[5] = 0; w.out[6] = 0; w.out[7] = 0
393 write_u64_le(w.out, 8, w.n_tensors)
394 write_u64_le(w.out, 16, header_bytes) // metadata_off
395 write_u64_le(w.out, 24, 0) // metadata_len
396 write_u64_le(w.out, 32, header_bytes) // block_table_off
397 write_u64_le(w.out, 40, block_table_bytes)
398 write_u64_le(w.out, 48, data_off)
399 var i: i64 = 56
400 while i < 64 { w.out[i] = 0; i = i + 1 }
401
402 var b: i64 = 0
403 while b < w.n_blocks {
404 let desc_off: i64 = header_bytes + b * NXGGUF_BLOCK_DESC_BYTES
405 var k: i64 = 0
406 while k < 32 {
407 let src: *u8 = (w.block_hashes as i64 + b * 32 + k) as *u8
408 w.out[desc_off + k] = *src
409 k = k + 1
410 }
411 write_u64_le(w.out, desc_off + 32, data_off + w.block_offs[b])
412 write_u64_le(w.out, desc_off + 40, w.block_lens[b])
413 write_u64_le(w.out, desc_off + 48, 0)
414 b = b + 1
415 }
416
417 var di: i64 = 0
418 while di < w.block_data_pos {
419 let src: *u8 = (w.block_data as i64 + di) as *u8
420 w.out[data_off + di] = *src
421 di = di + 1
422 }
423 return total
424}
425
426// === self-test ===
427
428// Construct a 4-tensor fixture in memory + verify open/total-bytes
429// round-trip. No file I/O so this runs anywhere.
430func main() -> i64 {
431 let buf: *u8 = sys_mmap(8192)
432
433 // header: "NXGF" magic bytes
434 buf[0] = 0x4E // 'N'
435 buf[1] = 0x58 // 'X'
436 buf[2] = 0x47 // 'G'
437 buf[3] = 0x46 // 'F'
438 // version = 1, little-endian u32
439 buf[4] = 1; buf[5] = 0; buf[6] = 0; buf[7] = 0
440 // tensor_count = 4
441 buf[8] = 4; buf[9] = 0; buf[10] = 0; buf[11] = 0
442 buf[12] = 0; buf[13] = 0; buf[14] = 0; buf[15] = 0
443 // metadata_off = 64, metadata_len = 0 (no metadata in this fixture)
444 buf[16] = 64; buf[17] = 0; buf[18] = 0; buf[19] = 0
445 buf[20] = 0; buf[21] = 0; buf[22] = 0; buf[23] = 0
446 buf[24] = 0; buf[25] = 0; buf[26] = 0; buf[27] = 0
447 buf[28] = 0; buf[29] = 0; buf[30] = 0; buf[31] = 0
448 // block_table_off = 64, block_table_len = 56 (one block)
449 buf[32] = 64; buf[33] = 0; buf[34] = 0; buf[35] = 0
450 buf[36] = 0; buf[37] = 0; buf[38] = 0; buf[39] = 0
451 buf[40] = 56; buf[41] = 0; buf[42] = 0; buf[43] = 0
452 buf[44] = 0; buf[45] = 0; buf[46] = 0; buf[47] = 0
453 // data_off = 120
454 buf[48] = 120; buf[49] = 0; buf[50] = 0; buf[51] = 0
455 buf[52] = 0; buf[53] = 0; buf[54] = 0; buf[55] = 0
456 // reserved
457 buf[56] = 0; buf[57] = 0; buf[58] = 0; buf[59] = 0
458 buf[60] = 0; buf[61] = 0; buf[62] = 0; buf[63] = 0
459
460 // block descriptor at offset 64: byte_length = 1024 at offset 64+40 = 104
461 buf[104] = 0; buf[105] = 4; buf[106] = 0; buf[107] = 0 // 1024 LE
462 buf[108] = 0; buf[109] = 0; buf[110] = 0; buf[111] = 0
463
464 let file: *NxgFile = sys_mmap(64) as *NxgFile
465 let rc: i64 = nxgguf_open(file, buf, 8192)
466 if rc != 0 { return __syscall(93, 100 + (0 - rc), 0, 0, 0, 0, 0) }
467 if file.tensor_count != 4 { return __syscall(93, 50, 0, 0, 0, 0, 0) }
468 if file.data_off != 120 { return __syscall(93, 51, 0, 0, 0, 0, 0) }
469 if nxgguf_total_bytes(file) != 1024 {
470 return __syscall(93, 52, 0, 0, 0, 0, 0)
471 }
472 if nxgguf_unique_blocks(file) != 1 {
473 return __syscall(93, 53, 0, 0, 0, 0, 0)
474 }
475
476 // Hash-equality + dedup self-test:
477 // Two identical 32-byte buffers -> same sha256 -> hash_eq returns 0.
478 let hash_a: *u8 = sys_mmap(64)
479 let hash_b: *u8 = sys_mmap(64)
480 let payload: *u8 = sys_mmap(128)
481 var ki: i64 = 0
482 while ki < 64 { payload[ki] = ki & 0xFF; ki = ki + 1 }
483 nxgguf_hash_block(payload, 64, hash_a)
484 nxgguf_hash_block(payload, 64, hash_b)
485 if nxgguf_hash_eq(hash_a, hash_b) != 0 {
486 return __syscall(93, 60, 0, 0, 0, 0, 0)
487 }
488 // Modify one byte of payload, re-hash -> different.
489 payload[0] = payload[0] ^ 1
490 nxgguf_hash_block(payload, 64, hash_b)
491 if nxgguf_hash_eq(hash_a, hash_b) != 1 {
492 return __syscall(93, 61, 0, 0, 0, 0, 0)
493 }
494
495 // Writer round-trip: build a 3-tensor file where two tensors
496 // share content (dedup test), then re-read it.
497 let wbuf: *u8 = sys_mmap(65536)
498 let w: *NxgWriter = nxgguf_writer_new(wbuf, 65536)
499
500 // Tensor data: 64 bytes of distinct + 64 bytes that match
501 // tensor 0 + 32 bytes of distinct.
502 let t0_data: *u8 = sys_mmap(64); var ti: i64 = 0
503 while ti < 64 { t0_data[ti] = ti & 0xFF; ti = ti + 1 }
504 let t1_data: *u8 = sys_mmap(64); ti = 0
505 while ti < 64 { t1_data[ti] = ti & 0xFF; ti = ti + 1 } // identical to t0
506 let t2_data: *u8 = sys_mmap(32); ti = 0
507 while ti < 32 { t2_data[ti] = (ti + 100) & 0xFF; ti = ti + 1 }
508
509 let shape_raw: *u8 = sys_mmap(64); let shape_p: *i64 = shape_raw as *i64
510 shape_p[0] = 8; shape_p[1] = 8
511
512 nxgguf_writer_add(w, "tensor.0" as *u8, 8, 2, shape_p, DTYPE_INT8, t0_data, 64)
513 nxgguf_writer_add(w, "tensor.1" as *u8, 8, 2, shape_p, DTYPE_INT8, t1_data, 64)
514 let shape2_raw: *u8 = sys_mmap(64); let shape2_p: *i64 = shape2_raw as *i64
515 shape2_p[0] = 32
516 nxgguf_writer_add(w, "tensor.2" as *u8, 8, 1, shape2_p, DTYPE_INT8, t2_data, 32)
517
518 let total: i64 = nxgguf_writer_finalize(w)
519 if total <= 0 { return __syscall(93, 70, 0, 0, 0, 0, 0) }
520
521 // Round-trip verify: open + check counts. Dedup means 3
522 // tensors share 2 unique blocks.
523 let rd_file: *NxgFile = sys_mmap(64) as *NxgFile
524 let rrc: i64 = nxgguf_open(rd_file, wbuf, total)
525 if rrc != 0 { return __syscall(93, 71, 0, 0, 0, 0, 0) }
526 if rd_file.tensor_count != 3 { return __syscall(93, 72, 0, 0, 0, 0, 0) }
527 if nxgguf_unique_blocks(rd_file) != 2 {
528 return __syscall(93, 73, 0, 0, 0, 0, 0)
529 }
530 if nxgguf_total_bytes(rd_file) != 96 { // 64 + 32 = 96 (deduped)
531 return __syscall(93, 74, 0, 0, 0, 0, 0)
532 }
533
534 return __syscall(93, 42, 0, 0, 0, 0, 0)
535}