nx_zimage_real_qkv.nx source
↩ module page · 114 lines · 4612 B
1// nx_zimage_real_qkv.nx -- REAL Qwen layer-0 Q/K/V projections on REAL weights, sovereignly.
2//
3// sd-server -> Nishi migration: extends the real forward from Q to the FULL attention input -- computes
4// Q, K and V projections of token-0's real embedding against blk.0.attn_{q,k,v}.weight (all Q6_K, ~335MB
5// into the file) with our own f32 matmul, and records the out-dims (confirming GQA: q_out=4096 > kv_out).
6// Real Qwen weights -> our sovereign matmuls -> the real attention inputs. Bounded 384MB prefix; only the
7// slices we need are dequantized.
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_q6_k_to_f32.nx"
20import "nx_f32.nx"
21import "nx_f32_cvt.nx"
22
23func qkv_emit(fd: i64, key: *u8, key_len: i64, value: i64) -> i64 {
24 let line: *u8 = sys_mmap(64)
25 var lo: i64 = 0
26 var ki: i64 = 0
27 while ki < key_len { line[lo] = key[ki]; lo = lo + 1; ki = ki + 1 }
28 line[lo] = 0x3D; lo = lo + 1
29 let dec: *u8 = sys_mmap(32)
30 let nd: i64 = nx_strconv_format_i64(value, dec)
31 var k: i64 = 0
32 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
33 line[lo] = 0x0A; lo = lo + 1
34 return sys_write(fd, line, lo)
35}
36
37// dequant the first 4 output rows of a weight at w_off and project: out[o] = sum_i embed[i] * W[o][i].
38// returns 0 ok, 1 if any output is Inf/NaN.
39func qkv_proj4(buf: *u8, w_off: i64, HID: i64, embed: *i64, out: *i64) -> i64 {
40 let wq: *i64 = sys_mmap(4 * HID * 8) as *i64
41 nx_q6_k_to_f32(buf, w_off, 4 * HID, wq)
42 var o: i64 = 0
43 while o < 4 {
44 var acc: i64 = 0
45 var i: i64 = 0
46 while i < HID { acc = nx_f32_add(acc, nx_f32_mul(embed[i], wq[o * HID + i])); i = i + 1 }
47 if (acc & 0x7F800000) == 0x7F800000 { return 1 }
48 out[o] = acc
49 o = o + 1
50 }
51 return 0
52}
53
54func main() -> i64 {
55 let path: *u8 = "/mnt/c/Users/elder/elder-ai-platform/models/unified/text_encoder/Z-Image_Qwen_3_4b-Q6_K.gguf" as *u8
56 let fd: i64 = sys_openat_rd(path)
57 if fd < 0 { return 30 }
58 let CAP: i64 = 402653184 // 384 MB
59 let buf: *u8 = sys_mmap(CAP)
60 var total: i64 = 0
61 var go: i64 = 1
62 while go == 1 {
63 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
64 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
65 }
66 sys_close(fd)
67 if total < 100000000 { return 31 }
68
69 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
70 if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
71
72 let ei: nx_int = nx_gguf_find_tensor(hdr, "token_embd.weight" as *u8, 17)
73 if ei < 0 { return 60 }
74 let eti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, ei)
75 let HID: i64 = eti.dim_0
76 let embed: *i64 = sys_mmap(HID * 8) as *i64
77 nx_q6_k_to_f32(buf, hdr.data_off + eti.offset, HID, embed)
78
79 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
80 let ki: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_k.weight" as *u8, 19)
81 let vi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_v.weight" as *u8, 19)
82 if qi < 0 { return 61 }
83 if ki < 0 { return 62 }
84 if vi < 0 { return 63 }
85 let qti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
86 let kti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, ki)
87 let vti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, vi)
88 let q_off: i64 = hdr.data_off + qti.offset
89 let k_off: i64 = hdr.data_off + kti.offset
90 let v_off: i64 = hdr.data_off + vti.offset
91 if q_off + 262144 > total { return 64 }
92 if k_off + 262144 > total { return 65 }
93 if v_off + 262144 > total { return 66 }
94
95 let Q: *i64 = sys_mmap(4 * 8) as *i64
96 let K: *i64 = sys_mmap(4 * 8) as *i64
97 let V: *i64 = sys_mmap(4 * 8) as *i64
98 if qkv_proj4(buf, q_off, HID, embed, Q) != 0 { return 80 }
99 if qkv_proj4(buf, k_off, HID, embed, K) != 0 { return 81 }
100 if qkv_proj4(buf, v_off, HID, embed, V) != 0 { return 82 }
101
102 let ofd: i64 = sys_openat_wr("/tmp/zimg_qkv.txt" as *u8, 0x1a4)
103 if ofd >= 0 {
104 qkv_emit(ofd, "HID" as *u8, 3, HID)
105 qkv_emit(ofd, "q_out" as *u8, 5, qti.dim_1)
106 qkv_emit(ofd, "k_out" as *u8, 5, kti.dim_1)
107 qkv_emit(ofd, "v_out" as *u8, 5, vti.dim_1)
108 qkv_emit(ofd, "Q0_bits" as *u8, 7, Q[0])
109 qkv_emit(ofd, "K0_bits" as *u8, 7, K[0])
110 qkv_emit(ofd, "V0_bits" as *u8, 7, V[0])
111 sys_close(ofd)
112 }
113 return 0
114}