nx_qwen_hybrid_qkv.nx source
↩ module page · 129 lines · 5708 B
1// nx_qwen_hybrid_qkv.nx -- the hybrid mixed-precision Q/K/V projection stage on real Qwen weights.
2//
3// sd-server -> Nishi migration (integer layer, task#19). The real Q4_K_M blk.0 mixes quant types, so the
4// integer layer is a hybrid: attn_q/attn_k are Q4_K (fast integer path nx_q4k_linear -> Q10), attn_v is Q5_K
5// (f32 path nx_q5_k_to_f32 + f32 dot). This proves those compose correctly on real weights: integer Q/K
6// (converted Q10->f32) match the f32 reference within 2%, and f32 V (via the ggml-verified Q5_K dequant)
7// produces finite values. The Q10<->f32 boundary is the hybrid layer's key seam.
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_q5_k_to_f32.nx"
24import "nx_f32.nx"
25import "nx_f32_cvt.nx"
26import "nx_f32_div.nx"
27
28func hq_emit(fd: i64, key: *u8, key_len: i64, value: i64) -> i64 {
29 let line: *u8 = sys_mmap(80)
30 var lo: i64 = 0
31 var ki: i64 = 0
32 while ki < key_len { line[lo] = key[ki]; lo = lo + 1; ki = ki + 1 }
33 line[lo] = 0x3D; lo = lo + 1
34 let dec: *u8 = sys_mmap(32)
35 let nd: i64 = nx_strconv_format_i64(value, dec)
36 var k: i64 = 0
37 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
38 line[lo] = 0x0A; lo = lo + 1
39 return sys_write(fd, line, lo)
40}
41
42// f32 dot of one weight row (f32) with the f32 activation
43func hq_f32dot(wrow: *i64, act_f32: *i64, in_dim: i64) -> i64 {
44 var acc: i64 = 0
45 var i: i64 = 0
46 while i < in_dim { acc = nx_f32_add(acc, nx_f32_mul(wrow[i], act_f32[i])); i = i + 1 }
47 return acc
48}
49
50func main() -> i64 {
51 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
52 let fd: i64 = sys_openat_rd(path)
53 if fd < 0 { return 30 }
54 let CAP: i64 = 1153433600
55 let buf: *u8 = sys_mmap(CAP)
56 var total: i64 = 0
57 var go: i64 = 1
58 while go == 1 {
59 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
60 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
61 }
62 sys_close(fd)
63 if total < 100000000 { return 31 }
64 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
65 if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
66
67 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
68 let ki: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_k.weight" as *u8, 19)
69 let vi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_v.weight" as *u8, 19)
70 if qi < 0 { return 60 }
71 if ki < 0 { return 61 }
72 if vi < 0 { return 62 }
73 let qti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
74 let kti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, ki)
75 let vti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, vi)
76 if qti.ggml_type != 12 { return 63 } // Q4_K
77 if kti.ggml_type != 12 { return 64 } // Q4_K
78 if vti.ggml_type != 13 { return 65 } // Q5_K
79 let IN: i64 = qti.dim_0
80 let q_off: i64 = hdr.data_off + qti.offset
81 let k_off: i64 = hdr.data_off + kti.offset
82 let v_off: i64 = hdr.data_off + vti.offset
83 let M: i64 = 4
84 if v_off + M * (IN / 256) * 176 > total { return 66 }
85
86 let act: *i64 = sys_mmap(IN * 8) as *i64
87 let act_f32: *i64 = sys_mmap(IN * 8) as *i64
88 var i: i64 = 0
89 while i < IN { act[i] = 1024 + (i - (i / 5) * 5) * 256; act_f32[i] = nx_q10_to_f32(act[i]); i = i + 1 }
90
91 // integer Q/K linears (Q4_K) -> Q10
92 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
93 let q_int: *i64 = sys_mmap(M * 8) as *i64
94 let k_int: *i64 = sys_mmap(M * 8) as *i64
95 nx_q4k_linear(buf, q_off, M, IN, act, 1, it, q_int)
96 nx_q4k_linear(buf, k_off, M, IN, act, 1, it, k_int)
97
98 // f32 references: dequant first M rows of each (contiguous) + f32 dot
99 let q_wf32: *i64 = sys_mmap(M * IN * 8) as *i64
100 let k_wf32: *i64 = sys_mmap(M * IN * 8) as *i64
101 let v_wf32: *i64 = sys_mmap(M * IN * 8) as *i64
102 nx_q4k_to_f32(buf, q_off, M * IN, q_wf32)
103 nx_q4k_to_f32(buf, k_off, M * IN, k_wf32)
104 nx_q5_k_to_f32(buf, v_off, M * IN, v_wf32) // Q5_K f32 path (the V branch)
105
106 let tolf: i64 = nx_f32_div(nx_i32_to_f32(2), nx_i32_to_f32(100))
107 let ofd: i64 = sys_openat_wr("/tmp/zimg_hybrid_qkv.txt" as *u8, 0x1a4)
108 if ofd >= 0 { hq_emit(ofd, "IN" as *u8, 2, IN); hq_emit(ofd, "M" as *u8, 1, M) }
109
110 var o: i64 = 0
111 while o < M {
112 let qf: i64 = hq_f32dot(((q_wf32 as i64) + o * IN * 8) as *i64, act_f32, IN)
113 let kf: i64 = hq_f32dot(((k_wf32 as i64) + o * IN * 8) as *i64, act_f32, IN)
114 let vf: i64 = hq_f32dot(((v_wf32 as i64) + o * IN * 8) as *i64, act_f32, IN)
115 let qint_f: i64 = nx_q10_to_f32(q_int[o])
116 let kint_f: i64 = nx_q10_to_f32(k_int[o])
117 if (nx_f32_sub(qint_f, qf) & 0x7FFFFFFF) >= nx_f32_mul(tolf, qf & 0x7FFFFFFF) { sys_close(ofd); return 70 + o }
118 if (nx_f32_sub(kint_f, kf) & 0x7FFFFFFF) >= nx_f32_mul(tolf, kf & 0x7FFFFFFF) { sys_close(ofd); return 74 + o }
119 if (vf & 0x7F800000) == 0x7F800000 { sys_close(ofd); return 78 } // V Inf/NaN
120 if ofd >= 0 {
121 hq_emit(ofd, "Qint_bits" as *u8, 9, qint_f)
122 hq_emit(ofd, "Qf32_bits" as *u8, 9, qf)
123 hq_emit(ofd, "Vf32_bits" as *u8, 9, vf)
124 }
125 o = o + 1
126 }
127 if ofd >= 0 { hq_emit(ofd, "hybrid_ok" as *u8, 9, M); sys_close(ofd) }
128 return 0
129}