code wiki / (root) / nxgguf.nx

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}