nx_q8_0_from_q5_0.nx source

1// nx_q8_0_from_q5_0.nx -- re-quantize a Q5_0 tensor to Q8_0 IN MEMORY. 2// 3// THE KEY MOVE (2026-07-08): Q5_0's nibble+qh unpack is what made a per-token 4// Q5_0 SIMD dequant-dot COMPUTE-bound (2x slower). But Q5_0 -> Q8_0 is EXACT 5// (the 5-bit values -16..15 fit in int8, same per-block f16 scale d), so we 6// repack ONCE AT LOAD -- amortizing the unpack -- and the forward then runs 7// the TRIVIAL Q8_0 SIMD dequant-dot (__f32_i8dot32, PROVEN 10.2x) per token. 8// The unpack uses the blessed __q5_unpack32 intrinsic itself. 9// 10// Q5_0 block (22B): d(f16) + qh(u32) + qs(16B). Q8_0 block (34B): d(f16) + 11// 32 int8. Repack: copy d; __q5_unpack32(qh+qs) -> the 32 int8 = q5-16. 12// lineage_id: q8_0_from_q5_0_v1 13import "nx_syscalls.nx" 14import "nx_le.nx" 15import "nx_q5_0_to_f32.nx" 16import "nx_q8_0_to_f32.nx" 17 18static g_q5repack_consts: i64 19func _q5repack_consts() -> *u8 { 20 if g_q5repack_consts == 0 { 21 let cc: *u8 = sys_mmap(80) 22 var i: i64 = 0 23 while i < 16 { cc[i] = 0x0F as u8; i = i + 1 } // c_0F 24 i = 0 25 while i < 8 { cc[16+i] = 0 as u8; i = i + 1 } 26 while i < 16 { cc[16+i] = 1 as u8; i = i + 1 } // pshuf_lo 27 i = 0 28 while i < 8 { cc[32+i] = 2 as u8; i = i + 1 } 29 while i < 16 { cc[32+i] = 3 as u8; i = i + 1 } // pshuf_hi 30 i = 0 31 while i < 8 { cc[48+i] = (1 << i) as u8; cc[48+8+i] = (1 << i) as u8; i = i + 1 } // bitmask 32 i = 0 33 while i < 16 { cc[64+i] = 0x10 as u8; i = i + 1 } // c_10 34 g_q5repack_consts = cc as i64 35 } 36 return g_q5repack_consts as *u8 37} 38 39// Repack n_values (must be a multiple of 32) of Q5_0 at q5buf+q5off into 40// contiguous Q8_0 blocks at out_q8. Returns 0. 41func nx_q8_0_from_q5_0(q5buf: *u8, q5off: i64, n_values: i64, out_q8: *u8) -> i64 { 42 let consts: *u8 = _q5repack_consts() 43 let nblk: i64 = (n_values + NX_Q5_0_VPB - 1) / NX_Q5_0_VPB 44 let q5pb: i64 = q5buf as i64 45 let q8pb: i64 = out_q8 as i64 46 var b: i64 = 0 47 while b < nblk { 48 let q5boff: i64 = q5off + b * NX_Q5_0_BPB // 22 B/block 49 let q8boff: i64 = b * NX_Q8_0_BPB // 34 B/block 50 out_q8[q8boff + 0] = q5buf[q5boff + 0] // copy d (f16) lo 51 out_q8[q8boff + 1] = q5buf[q5boff + 1] // copy d (f16) hi 52 // qh at q5boff+2, qs at q5boff+6 are contiguous -> qhqs = q5boff+2. 53 // int8 output at q8boff+2 (the 32 quants). 54 __q5_unpack32((q5pb + q5boff + 2) as *u8, (q8pb + q8boff + 2) as *u8, consts) 55 b = b + 1 56 } 57 return 0 58}