nx_q5_k_to_f32.nx source
↩ module page · 143 lines · 5869 B
1// nx_q5_k_to_f32.nx -- bits-up Q5_K dequantization producing IEEE 754 binary32 outputs.
2//
3// Q5_K = the (ggml-correct) Q4_K structure + a 5th bit per value from qh. Real Q4_K_M models store
4// attn_v / attn_output / ffn_down as Q5_K, so this is required to load them sovereignly.
5//
6// ggml block_q5_K (176 bytes / 256 values):
7// [0..2) d f16 super-scale
8// [2..4) dmin f16 super-min
9// [4..16) scales 12 bytes = 8x(6-bit sc) + 8x(6-bit m) (get_scale_min_k4, same as Q4_K)
10// [16..48) qh 32 bytes = 1 high bit per value
11// [48..176) qs 128 bytes = 4 low bits per value
12// ggml dequantize_row_q5_K layout: 4 GROUPS of 64; group g feeds sub-block 2g (low nibble, 32 outputs)
13// and 2g+1 (high nibble, 32 outputs); the 5th bit is qh[l] tested with u1=1<<(2g) (low) / u2=1<<(2g+1)
14// (high). value = d*sc*(q4 + (bit?16:0)) - dmin*m (q5 in 0..31, min subtracted -- same form as Q4_K).
15// license_tier: ORIGINAL
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_le.nx"
19import "nx_gguf.nx"
20import "nx_gguf_load.nx"
21import "nx_f32.nx"
22import "nx_f32_cvt.nx"
23
24const NX_Q5K_OK: nx_int = 0
25const NX_Q5_K_VPB: i64 = 256
26const NX_Q5_K_BPB: i64 = 176
27
28func nx_q5k_block_to_f32(buf: *u8, super_off: i64, n_values: i64, out_f32: *i64) -> nx_int {
29 if buf == (0 as *u8) { return 1 }
30
31 let d_raw: i64 = nx_le_read_u16(buf, super_off)
32 let dmin_raw: i64 = nx_le_read_u16(buf, super_off + 2)
33 let d_f32: i64 = nx_f16_to_f32(d_raw)
34 let dmin_f32: i64 = nx_f16_to_f32(dmin_raw)
35
36 let scales_off: i64 = super_off + 4
37 let qh_off: i64 = super_off + 16
38 let qs_off: i64 = super_off + 48
39
40 var g: nx_int = 0
41 while g < 4 {
42 let is0: nx_int = g + g // 2g
43 let is1: nx_int = is0 + 1 // 2g+1
44
45 // 6-bit (sc,m) unpack -- identical to the ggml-correct Q4_K path.
46 var sc0: i64 = 0
47 var m0: i64 = 0
48 var sc1: i64 = 0
49 var m1s: i64 = 0
50 if is0 < 4 {
51 sc0 = nx_le_read_u8(buf, scales_off + is0) & 0x3F
52 m0 = nx_le_read_u8(buf, scales_off + is0 + 4) & 0x3F
53 } else {
54 let k0: nx_int = is0 - 4
55 let b_k: i64 = nx_le_read_u8(buf, scales_off + k0)
56 let b_k4: i64 = nx_le_read_u8(buf, scales_off + 4 + k0)
57 let b_8k: i64 = nx_le_read_u8(buf, scales_off + 8 + k0)
58 sc0 = ((b_k >> 6) << 4) | (b_8k & 0x0F)
59 m0 = ((b_k4 >> 6) << 4) | (b_8k >> 4)
60 }
61 if is1 < 4 {
62 sc1 = nx_le_read_u8(buf, scales_off + is1) & 0x3F
63 m1s = nx_le_read_u8(buf, scales_off + is1 + 4) & 0x3F
64 } else {
65 let k1: nx_int = is1 - 4
66 let c_k: i64 = nx_le_read_u8(buf, scales_off + k1)
67 let c_k4: i64 = nx_le_read_u8(buf, scales_off + 4 + k1)
68 let c_8k: i64 = nx_le_read_u8(buf, scales_off + 8 + k1)
69 sc1 = ((c_k >> 6) << 4) | (c_8k & 0x0F)
70 m1s = ((c_k4 >> 6) << 4) | (c_8k >> 4)
71 }
72
73 let d1_f32: i64 = __f32_mul(d_f32, nx_i32_to_f32(sc0))
74 let neg_m0: i64 = __f32_mul(__f32_mul(dmin_f32, nx_i32_to_f32(m0)), __f32_from_i64(0 - 1))
75 let d2_f32: i64 = __f32_mul(d_f32, nx_i32_to_f32(sc1))
76 let neg_m1: i64 = __f32_mul(__f32_mul(dmin_f32, nx_i32_to_f32(m1s)), __f32_from_i64(0 - 1))
77
78 let grp_off: i64 = qs_off + g * 32
79 let u1: i64 = 1 << (g + g) // low-nibble high bit mask for this group
80 let u2: i64 = 1 << (g + g + 1) // high-nibble high bit mask
81
82 var l: nx_int = 0
83 while l < 32 {
84 let byte_v: i64 = nx_le_read_u8(buf, grp_off + l)
85 let qh_l: i64 = nx_le_read_u8(buf, qh_off + l)
86 var q_lo: i64 = byte_v & 0x0F
87 var q_hi: i64 = byte_v >> 4
88 if (qh_l & u1) != 0 { q_lo = q_lo + 16 }
89 if (qh_l & u2) != 0 { q_hi = q_hi + 16 }
90 let out_lo: i64 = is0 * 32 + l
91 let out_hi: i64 = is1 * 32 + l
92 if out_lo < n_values {
93 out_f32[out_lo] = __f32_add(__f32_mul(d1_f32, nx_i32_to_f32(q_lo)), neg_m0)
94 }
95 if out_hi < n_values {
96 out_f32[out_hi] = __f32_add(__f32_mul(d2_f32, nx_i32_to_f32(q_hi)), neg_m1)
97 }
98 l = l + 1
99 }
100 g = g + 1
101 }
102 return NX_Q5K_OK
103}
104
105func nx_q5_k_to_f32(buf: *u8, base_off: i64, n_values: i64, out_f32: *i64) -> nx_int {
106 let n_super: i64 = (n_values + NX_Q5_K_VPB - 1) / NX_Q5_K_VPB
107 var sb: i64 = 0
108 while sb < n_super {
109 let super_off: i64 = base_off + sb * NX_Q5_K_BPB
110 let block_n: i64 = n_values - sb * NX_Q5_K_VPB
111 var take: i64 = block_n
112 if take > NX_Q5_K_VPB { take = NX_Q5_K_VPB }
113 nx_q5k_block_to_f32(buf, super_off, take,
114 (out_f32 as i64 + sb * NX_Q5_K_VPB * 8) as *i64)
115 sb = sb + 1
116 }
117 return NX_Q5K_OK
118}
119
120// Build-only sanity: dequant a tiny synthetic block and confirm the value formula/layout run.
121// d=1.0 (0x3C00), dmin=0, scales all sc=1/m=0 region set so element0 = d*sc*q5. qh bit sets +16.
122func main() -> i64 {
123 let buf: *u8 = sys_mmap(NX_Q5_K_BPB)
124 var i: i64 = 0
125 while i < NX_Q5_K_BPB { buf[i] = 0; i = i + 1 }
126 // d = 1.0 (f16 0x3C00)
127 buf[0] = 0x00; buf[1] = 0x3C
128 // dmin = 0
129 // scales: sub-block 0 sc = 1 (scales[0]&0x3F = 1), m = 0
130 buf[4] = 0x01
131 // qs[0] low nibble = 7 (qs at offset 48)
132 buf[48] = 0x07
133 // qh[0] bit0 set -> low-nibble +16 -> q5 = 7 + 16 = 23 (qh at offset 16)
134 buf[16] = 0x01
135
136 let out: *i64 = sys_mmap(256 * 8) as *i64
137 nx_q5k_block_to_f32(buf, 0, 256, out)
138 // element0 = d(1.0) * sc(1) * q5(23) - 0 = 23.0
139 if out[0] != nx_i32_to_f32(23) { return 20 }
140 // without the qh bit (element1, qs[1]=0) -> 0
141 if out[1] != nx_i32_to_f32(0) { return 21 }
142 return 0
143}