nx_q4k_linear_hp.nx source
↩ module page · 139 lines · 5283 B
1// nx_q4k_linear_hp.nx -- higher-precision integer linear (Q14 activation -> Q24 output).
2//
3// sd-server -> Nishi migration (precision upgrade for the finding: Q10 activation quant is amplified by
4// SwiGLU -> FFN worst 3.3%). Uses a Q14 activation (4 more fractional bits) and keeps the output at Q24
5// (dot is Q24-weight x Q14-act = Q38 -> >>14 -> Q24) instead of collapsing to Q10. Test: on real
6// blk.0.attn_q, compare the Q14/Q24 path vs the Q10 path against the FULL-precision-activation f32 dot ->
7// the hp path's total error must be smaller (the precision knob works).
8// license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_tier.nx"
11import "nx_le.nx"
12import "nx_strconv.nx"
13import "nx_tensor.nx"
14import "nx_gguf.nx"
15import "nx_gguf_load.nx"
16import "nx_gguf_meta.nx"
17import "nx_placement.nx"
18import "nx_gguf_load_lazy.nx"
19import "nx_q4k_matmul.nx"
20import "nx_dequant_iter.nx"
21import "nx_q4k_linear.nx"
22import "nx_q4k_to_f32.nx"
23import "nx_f32.nx"
24import "nx_f32_cvt.nx"
25import "nx_f32_div.nx"
26
27func hp_q24_to_f32(q24: i64) -> i64 {
28 return nx_f32_div(nx_i32_to_f32(q24), nx_i32_to_f32(16777216)) // / 2^24
29}
30
31// Q14 activation x Q24 weight = Q38 dot -> Q24 output (round-half >>14).
32func nx_q4k_linear_hp(buf: *u8, w_off: i64, out_dim: i64, in_dim: i64,
33 act_q14: *i64, n_tokens: i64, it: *NxQ4KBlockIter, out_q24: *i64) -> i64 {
34 let n_blocks: i64 = in_dim / 256
35 let rstride: i64 = n_blocks * 144
36 var t: i64 = 0
37 while t < n_tokens {
38 let arow: *i64 = ((act_q14 as i64) + t * in_dim * 8) as *i64
39 var o: i64 = 0
40 while o < out_dim {
41 let dot_q38: i64 = nx_q4k_dot_row_col(buf, w_off + o * rstride, n_blocks, arow, it)
42 out_q24[t * out_dim + o] = (dot_q38 + 8192) / 16384
43 o = o + 1
44 }
45 t = t + 1
46 }
47 return 0
48}
49
50func hp_emit(fd: i64, key: *u8, kl: i64, v: i64) -> i64 {
51 let line: *u8 = sys_mmap(80)
52 var lo: i64 = 0
53 var ki: i64 = 0
54 while ki < kl { line[lo] = key[ki]; lo = lo + 1; ki = ki + 1 }
55 line[lo] = 0x3D; lo = lo + 1
56 let dec: *u8 = sys_mmap(32)
57 let nd: i64 = nx_strconv_format_i64(v, dec)
58 var k: i64 = 0
59 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
60 line[lo] = 0x0A; lo = lo + 1
61 return sys_write(fd, line, lo)
62}
63
64func main() -> i64 {
65 let path: *u8 = "/mnt/c/Users/elder/elder-ai-platform/models/unified/text_encoder/Huihui-Qwen3-4B-Instruct-2507-abliterated-Q4_K_M.gguf" as *u8
66 let fd: i64 = sys_openat_rd(path)
67 if fd < 0 { return 30 }
68 let CAP: i64 = 1153433600
69 let buf: *u8 = sys_mmap(CAP)
70 var total: i64 = 0
71 var go: i64 = 1
72 while go == 1 {
73 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
74 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
75 }
76 sys_close(fd)
77 if total < 100000000 { return 31 }
78 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
79 if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
80 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
81 if qi < 0 { return 60 }
82 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
83 if ti.ggml_type != 12 { return 61 }
84 let IN: i64 = ti.dim_0
85 let w_off: i64 = hdr.data_off + ti.offset
86 let M: i64 = 8
87
88 // activation f32 + its Q10 and Q14 quantizations
89 let act_f32: *i64 = sys_mmap(IN * 8) as *i64
90 let act_q10: *i64 = sys_mmap(IN * 8) as *i64
91 let act_q14: *i64 = sys_mmap(IN * 8) as *i64
92 var i: i64 = 0
93 while i < IN {
94 let a: i64 = nx_q10_to_f32(700 + (i - (i / 9) * 9) * 111)
95 act_f32[i] = a
96 act_q10[i] = _gguf_f32_to_q10(a)
97 act_q14[i] = _gguf_f32_to_q14(a)
98 i = i + 1
99 }
100
101 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
102 let o_q10: *i64 = sys_mmap(M * 8) as *i64
103 let o_q24: *i64 = sys_mmap(M * 8) as *i64
104 nx_q4k_linear(buf, w_off, M, IN, act_q10, 1, it, o_q10) // Q10 path
105 nx_q4k_linear_hp(buf, w_off, M, IN, act_q14, 1, it, o_q24) // Q14/Q24 path
106
107 let wf32: *i64 = sys_mmap(M * IN * 8) as *i64
108 nx_q4k_to_f32(buf, w_off, M * IN, wf32)
109
110 var tot_q10: i64 = 0
111 var tot_hp: i64 = 0
112 let ofd: i64 = sys_openat_wr("/tmp/zimg_hp.txt" as *u8, 0x1a4)
113 var o: i64 = 0
114 while o < M {
115 var ref: i64 = 0
116 i = 0
117 while i < IN { ref = nx_f32_add(ref, nx_f32_mul(wf32[o * IN + i], act_f32[i])); i = i + 1 }
118 let q10f: i64 = nx_q10_to_f32(o_q10[o])
119 let hpf: i64 = hp_q24_to_f32(o_q24[o])
120 let eq: i64 = nx_f32_sub(q10f, ref) & 0x7FFFFFFF
121 let eh: i64 = nx_f32_sub(hpf, ref) & 0x7FFFFFFF
122 tot_q10 = nx_f32_add(tot_q10, eq)
123 tot_hp = nx_f32_add(tot_hp, eh)
124 if ofd >= 0 {
125 hp_emit(ofd, "ref" as *u8, 3, ref)
126 hp_emit(ofd, "err_q10" as *u8, 7, eq)
127 hp_emit(ofd, "err_hp" as *u8, 6, eh)
128 }
129 o = o + 1
130 }
131 if ofd >= 0 {
132 hp_emit(ofd, "tot_q10_bits" as *u8, 12, tot_q10)
133 hp_emit(ofd, "tot_hp_bits" as *u8, 11, tot_hp)
134 sys_close(ofd)
135 }
136 // the Q14/Q24 path must have strictly smaller total error than the Q10 path
137 if (tot_hp & 0x7FFFFFFF) >= (tot_q10 & 0x7FFFFFFF) { return 80 }
138 return 0
139}