nx_q4k_linear.nx source
↩ module page · 121 lines · 5052 B
1// nx_q4k_linear.nx -- the integer GEMM linear: real Q4_K weights x Q10 activation -> Q10 output.
2//
3// sd-server -> Nishi migration (CPU speed path). The integer analogue of qw_linear (nx_f32_qwen_layer):
4// out[t][o] = sum_i act[t][i] * W[o][i], with W a real Q4_K tensor (ggml-correct dequant, ~26x faster than
5// emulated f32) and the activation in Q10 fixed-point. Each output neuron o is one fused nx_q4k_dot_row_col
6// over weight row o (byte stride = n_blocks*144). This is the reusable matmul brick for an integer Qwen
7// layer. Gate proves it matches the f32 path across MULTIPLE real output neurons on blk.0.attn_q.
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_to_f32.nx"
22import "nx_f32.nx"
23import "nx_f32_cvt.nx"
24import "nx_f32_div.nx"
25
26// out_q10[t*out_dim + o] = sum_i act_q10[t*in_dim + i] * W[o][i] (W real Q4_K at w_off, row stride n_blocks*144)
27func nx_q4k_linear(buf: *u8, w_off: i64, out_dim: i64, in_dim: i64,
28 act_q10: *i64, n_tokens: i64, it: *NxQ4KBlockIter, out_q10: *i64) -> i64 {
29 let n_blocks: i64 = in_dim / 256
30 let row_stride: i64 = n_blocks * 144
31 var t: i64 = 0
32 while t < n_tokens {
33 let act_row: *i64 = ((act_q10 as i64) + t * in_dim * 8) as *i64
34 var o: i64 = 0
35 while o < out_dim {
36 let dot_q34: i64 = nx_q4k_dot_row_col(buf, w_off + o * row_stride, n_blocks, act_row, it)
37 out_q10[t * out_dim + o] = nx_q4km_q20_to_q10(dot_q34)
38 o = o + 1
39 }
40 t = t + 1
41 }
42 return 0
43}
44
45func lin_emit(fd: i64, key: *u8, key_len: i64, value: i64) -> i64 {
46 let line: *u8 = sys_mmap(80)
47 var lo: i64 = 0
48 var ki: i64 = 0
49 while ki < key_len { line[lo] = key[ki]; lo = lo + 1; ki = ki + 1 }
50 line[lo] = 0x3D; lo = lo + 1
51 let dec: *u8 = sys_mmap(32)
52 let nd: i64 = nx_strconv_format_i64(value, dec)
53 var k: i64 = 0
54 while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
55 line[lo] = 0x0A; lo = lo + 1
56 return sys_write(fd, line, lo)
57}
58
59func main() -> i64 {
60 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
61 let fd: i64 = sys_openat_rd(path)
62 if fd < 0 { return 30 }
63 let CAP: i64 = 1153433600
64 let buf: *u8 = sys_mmap(CAP)
65 var total: i64 = 0
66 var go: i64 = 1
67 while go == 1 {
68 let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
69 if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
70 }
71 sys_close(fd)
72 if total < 100000000 { return 31 }
73 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
74 if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
75 let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
76 if qi < 0 { return 60 }
77 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
78 if ti.ggml_type != 12 { return 61 }
79 let HID: i64 = ti.dim_0
80 let w_off: i64 = hdr.data_off + ti.offset
81 let n_blocks: i64 = HID / 256
82 let M: i64 = 8 // verify the first 8 output neurons
83 if w_off + M * n_blocks * 144 > total { return 63 }
84
85 // Q10 activation (1 token, HID values)
86 let act: *i64 = sys_mmap(HID * 8) as *i64
87 let act_f32: *i64 = sys_mmap(HID * 8) as *i64
88 var i: i64 = 0
89 while i < HID { act[i] = 1024 + (i - (i / 5) * 5) * 256; act_f32[i] = nx_q10_to_f32(act[i]); i = i + 1 }
90
91 // integer linear -> first M outputs
92 let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
93 let out_int: *i64 = sys_mmap(M * 8) as *i64
94 nx_q4k_linear(buf, w_off, M, HID, act, 1, it, out_int)
95
96 // f32 reference: dequant first M rows (contiguous) + f32 dot
97 let wf32: *i64 = sys_mmap(M * HID * 8) as *i64
98 nx_q4k_to_f32(buf, w_off, M * HID, wf32)
99
100 let tolf: i64 = nx_f32_div(nx_i32_to_f32(2), nx_i32_to_f32(100))
101 let ofd: i64 = sys_openat_wr("/tmp/zimg_q4k_linear.txt" as *u8, 0x1a4)
102 if ofd >= 0 { lin_emit(ofd, "HID" as *u8, 3, HID); lin_emit(ofd, "M" as *u8, 1, M) }
103
104 var o: i64 = 0
105 while o < M {
106 var accf: i64 = 0
107 i = 0
108 while i < HID { accf = nx_f32_add(accf, nx_f32_mul(wf32[o * HID + i], act_f32[i])); i = i + 1 }
109 let int_f32: i64 = nx_q10_to_f32(out_int[o])
110 let absdiff: i64 = nx_f32_sub(int_f32, accf) & 0x7FFFFFFF
111 let thresh: i64 = nx_f32_mul(tolf, accf & 0x7FFFFFFF)
112 if ofd >= 0 {
113 lin_emit(ofd, "int_o_bits" as *u8, 10, int_f32)
114 lin_emit(ofd, "f32_o_bits" as *u8, 10, accf)
115 }
116 if absdiff >= thresh { sys_close(ofd); return 80 + o } // output o mismatch
117 o = o + 1
118 }
119 if ofd >= 0 { lin_emit(ofd, "all_match" as *u8, 9, M); sys_close(ofd) }
120 return 0
121}