nx_f32_llm_dim_dump_test.nx source
↩ module page · 101 lines · 4878 B
1// nx_f32_llm_dim_dump_test.nx -- dump attn_k.weight tensor dimensions
2// to /tmp/nx_gguf_dims.tsv to verify ggml's storage convention.
3
4import "nx_syscalls.nx"
5import "nx_tier.nx"
6import "nx_le.nx"
7import "nx_strconv.nx"
8import "nx_gguf.nx"
9import "nx_gguf_load.nx"
10
11func _emit_kv(fd: i64, key: *u8, key_len: nx_int, value: i64) -> i64 {
12 let line: *u8 = sys_mmap(128)
13 var lo: i64 = 0
14 var ki: nx_int = 0
15 while ki < key_len { line[lo] = key[ki]; lo = lo + 1; ki = ki + 1 }
16 line[lo] = 0x09 as u8; lo = lo + 1
17 let dec: *u8 = sys_mmap(32)
18 let n_dec: i64 = nx_strconv_format_i64(value, dec)
19 var k: i64 = 0
20 while k < n_dec { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
21 line[lo] = 0x0A as u8; lo = lo + 1
22 return sys_write(fd, line, lo)
23}
24
25func main() -> i64 {
26 let path: *u8 = sys_mmap(64)
27 path[0]=0x2F as u8; path[1]=0x74 as u8; path[2]=0x6D as u8; path[3]=0x70 as u8
28 path[4]=0x2F as u8; path[5]=0x6E as u8; path[6]=0x78 as u8; path[7]=0x5F as u8
29 path[8]=0x72 as u8; path[9]=0x65 as u8; path[10]=0x61 as u8; path[11]=0x6C as u8
30 path[12]=0x5F as u8; path[13]=0x6D as u8; path[14]=0x6F as u8; path[15]=0x64 as u8
31 path[16]=0x65 as u8; path[17]=0x6C as u8; path[18]=0x2E as u8
32 path[19]=0x67 as u8; path[20]=0x67 as u8; path[21]=0x75 as u8; path[22]=0x66 as u8
33 path[23]=0 as u8
34
35 let len_out: *i64 = sys_mmap(8) as *i64
36 let buf: *u8 = sys_read_file(path, len_out)
37 if buf == (0 as *u8) { return 10 }
38
39 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
40 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
41
42 // Open TSV.
43 let tsv_path: *u8 = sys_mmap(32)
44 tsv_path[0]=0x2F as u8; tsv_path[1]=0x74 as u8; tsv_path[2]=0x6D as u8; tsv_path[3]=0x70 as u8
45 tsv_path[4]=0x2F as u8; tsv_path[5]=0x6E as u8; tsv_path[6]=0x78 as u8; tsv_path[7]=0x5F as u8
46 tsv_path[8]=0x67 as u8; tsv_path[9]=0x67 as u8; tsv_path[10]=0x75 as u8; tsv_path[11]=0x66 as u8
47 tsv_path[12]=0x5F as u8; tsv_path[13]=0x64 as u8; tsv_path[14]=0x69 as u8; tsv_path[15]=0x6D as u8
48 tsv_path[16]=0x73 as u8; tsv_path[17]=0x2E as u8
49 tsv_path[18]=0x74 as u8; tsv_path[19]=0x73 as u8; tsv_path[20]=0x76 as u8
50 tsv_path[21]=0 as u8
51 let fd: i64 = sys_openat_wr(tsv_path, 0x1A4)
52 if fd < 0 { return 30 }
53
54 // Lookup blk.0.attn_k.weight (kv_dim = 128, hidden = 896 -- ASYMMETRIC).
55 let n_k: *u8 = sys_mmap(24)
56 n_k[0]=0x62 as u8; n_k[1]=0x6c as u8; n_k[2]=0x6b as u8; n_k[3]=0x2e as u8
57 n_k[4]=0x30 as u8; n_k[5]=0x2e as u8; n_k[6]=0x61 as u8; n_k[7]=0x74 as u8
58 n_k[8]=0x74 as u8; n_k[9]=0x6e as u8; n_k[10]=0x5f as u8; n_k[11]=0x6b as u8
59 n_k[12]=0x2e as u8; n_k[13]=0x77 as u8; n_k[14]=0x65 as u8; n_k[15]=0x69 as u8
60 n_k[16]=0x67 as u8; n_k[17]=0x68 as u8; n_k[18]=0x74 as u8
61 let idx_k: nx_int = nx_gguf_find_tensor(hdr, n_k, 19)
62 if idx_k >= 0 {
63 let ti_k: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx_k)
64 let k_d0: *u8 = sys_mmap(11)
65 k_d0[0]=0x61 as u8; k_d0[1]=0x74 as u8; k_d0[2]=0x74 as u8; k_d0[3]=0x6e as u8
66 k_d0[4]=0x5f as u8; k_d0[5]=0x6b as u8; k_d0[6]=0x5f as u8
67 k_d0[7]=0x64 as u8; k_d0[8]=0x69 as u8; k_d0[9]=0x6d as u8; k_d0[10]=0x30 as u8
68 _emit_kv(fd, k_d0, 11, ti_k.dim_0)
69 let k_d1: *u8 = sys_mmap(11)
70 k_d1[0]=0x61 as u8; k_d1[1]=0x74 as u8; k_d1[2]=0x74 as u8; k_d1[3]=0x6e as u8
71 k_d1[4]=0x5f as u8; k_d1[5]=0x6b as u8; k_d1[6]=0x5f as u8
72 k_d1[7]=0x64 as u8; k_d1[8]=0x69 as u8; k_d1[9]=0x6d as u8; k_d1[10]=0x31 as u8
73 _emit_kv(fd, k_d1, 11, ti_k.dim_1)
74 }
75
76 // Also lookup blk.0.ffn_gate.weight (ffn_dim=4864, hidden=896).
77 let n_g: *u8 = sys_mmap(24)
78 n_g[0]=0x62 as u8; n_g[1]=0x6c as u8; n_g[2]=0x6b as u8; n_g[3]=0x2e as u8
79 n_g[4]=0x30 as u8; n_g[5]=0x2e as u8; n_g[6]=0x66 as u8; n_g[7]=0x66 as u8
80 n_g[8]=0x6e as u8; n_g[9]=0x5f as u8; n_g[10]=0x67 as u8; n_g[11]=0x61 as u8
81 n_g[12]=0x74 as u8; n_g[13]=0x65 as u8; n_g[14]=0x2e as u8; n_g[15]=0x77 as u8
82 n_g[16]=0x65 as u8; n_g[17]=0x69 as u8; n_g[18]=0x67 as u8; n_g[19]=0x68 as u8
83 n_g[20]=0x74 as u8
84 let idx_g: nx_int = nx_gguf_find_tensor(hdr, n_g, 21)
85 if idx_g >= 0 {
86 let ti_g: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx_g)
87 let k_g0: *u8 = sys_mmap(11)
88 k_g0[0]=0x66 as u8; k_g0[1]=0x66 as u8; k_g0[2]=0x6e as u8; k_g0[3]=0x67 as u8
89 k_g0[4]=0x5f as u8; k_g0[5]=0x64 as u8; k_g0[6]=0x69 as u8; k_g0[7]=0x6d as u8
90 k_g0[8]=0x30 as u8
91 _emit_kv(fd, k_g0, 9, ti_g.dim_0)
92 let k_g1: *u8 = sys_mmap(11)
93 k_g1[0]=0x66 as u8; k_g1[1]=0x66 as u8; k_g1[2]=0x6e as u8; k_g1[3]=0x67 as u8
94 k_g1[4]=0x5f as u8; k_g1[5]=0x64 as u8; k_g1[6]=0x69 as u8; k_g1[7]=0x6d as u8
95 k_g1[8]=0x31 as u8
96 _emit_kv(fd, k_g1, 9, ti_g.dim_1)
97 }
98
99 sys_close(fd)
100 return 0
101}