nx_f32_llm_read_dims_test.nx source

1// nx_f32_llm_read_dims_test.nx -- smoke for nx_f32_llm_read_dims.nx. 2// 3// Builds a synthetic GGUF with: 4// general.architecture STR "qwen2" 5// qwen2.block_count U32 2 6// qwen2.embedding_length U32 8 7// qwen2.attention.head_count U32 4 8// qwen2.attention.head_count_kv U32 2 (GQA: 2 KV heads, 4 Q heads -> group=2) 9// qwen2.feed_forward_length U32 16 10// (no head_count_kv variant tested below too) 11// + 1 token_embd.weight tensor with dim_0=12 (vocab_size=12) 12// 13// Verifies all dims are read correctly into NxF32LlamaModel. 14 15import "nx_syscalls.nx" 16import "nx_tier.nx" 17import "nx_le.nx" 18import "nx_gguf.nx" 19import "nx_gguf_load.nx" 20import "nx_gguf_meta.nx" 21import "nx_f32_llm.nx" 22import "nx_f32_llm_read_dims.nx" 23 24func _write_meta_u32(buf: *u8, p_in: i64, 25 key: *u8, key_len: i64, value: i64) -> i64 { 26 var p: i64 = p_in 27 nx_le_write_u64(buf, p, key_len); p = p + 8 28 var i: i64 = 0 29 while i < key_len { buf[p + i] = key[i]; i = i + 1 } 30 p = p + key_len 31 nx_le_write_u32(buf, p, NX_GGUF_TYPE_UINT32); p = p + 4 32 nx_le_write_u32(buf, p, value); p = p + 4 33 return p 34} 35 36func _write_meta_string(buf: *u8, p_in: i64, 37 key: *u8, key_len: i64, 38 sval: *u8, sval_len: i64) -> i64 { 39 var p: i64 = p_in 40 nx_le_write_u64(buf, p, key_len); p = p + 8 41 var i: i64 = 0 42 while i < key_len { buf[p + i] = key[i]; i = i + 1 } 43 p = p + key_len 44 nx_le_write_u32(buf, p, NX_GGUF_TYPE_STRING); p = p + 4 45 nx_le_write_u64(buf, p, sval_len); p = p + 8 46 var j: i64 = 0 47 while j < sval_len { buf[p + j] = sval[j]; j = j + 1 } 48 p = p + sval_len 49 return p 50} 51 52func main() -> i64 { 53 var vi: nx_int = 0 54 while vi < NX_FLD_N_VERDICTS { 55 if nx_fld_verdict_is_valid(vi) != 1 { return 5 + vi } 56 vi = vi + 1 57 } 58 59 let buf: *u8 = sys_mmap(1024) 60 buf[0]=0x47 as u8; buf[1]=0x47 as u8; buf[2]=0x55 as u8; buf[3]=0x46 as u8 61 nx_le_write_u32(buf, 4, 3) // version 62 nx_le_write_u64(buf, 8, 1) // 1 tensor (token_embd.weight) 63 nx_le_write_u64(buf, 16, 6) // 6 metadata entries 64 65 var p: i64 = 24 66 67 // general.architecture = "qwen2" 68 let k_arch: *u8 = sys_mmap(20) 69 k_arch[0]=0x67 as u8; k_arch[1]=0x65 as u8; k_arch[2]=0x6e as u8; k_arch[3]=0x65 as u8 70 k_arch[4]=0x72 as u8; k_arch[5]=0x61 as u8; k_arch[6]=0x6c as u8; k_arch[7]=0x2e as u8 71 k_arch[8]=0x61 as u8; k_arch[9]=0x72 as u8; k_arch[10]=0x63 as u8; k_arch[11]=0x68 as u8 72 k_arch[12]=0x69 as u8; k_arch[13]=0x74 as u8; k_arch[14]=0x65 as u8; k_arch[15]=0x63 as u8 73 k_arch[16]=0x74 as u8; k_arch[17]=0x75 as u8; k_arch[18]=0x72 as u8; k_arch[19]=0x65 as u8 74 let v_qwen2: *u8 = sys_mmap(5) 75 v_qwen2[0]=0x71 as u8; v_qwen2[1]=0x77 as u8; v_qwen2[2]=0x65 as u8; v_qwen2[3]=0x6e as u8 76 v_qwen2[4]=0x32 as u8 77 p = _write_meta_string(buf, p, k_arch, 20, v_qwen2, 5) 78 79 // qwen2.block_count = 2 80 let k_bc: *u8 = sys_mmap(17) 81 k_bc[0]=0x71 as u8; k_bc[1]=0x77 as u8; k_bc[2]=0x65 as u8; k_bc[3]=0x6e as u8 82 k_bc[4]=0x32 as u8; k_bc[5]=0x2e as u8; k_bc[6]=0x62 as u8; k_bc[7]=0x6c as u8 83 k_bc[8]=0x6f as u8; k_bc[9]=0x63 as u8; k_bc[10]=0x6b as u8; k_bc[11]=0x5f as u8 84 k_bc[12]=0x63 as u8; k_bc[13]=0x6f as u8; k_bc[14]=0x75 as u8; k_bc[15]=0x6e as u8 85 k_bc[16]=0x74 as u8 86 p = _write_meta_u32(buf, p, k_bc, 17, 2) 87 88 // qwen2.embedding_length = 8 89 let k_el: *u8 = sys_mmap(22) 90 k_el[0]=0x71 as u8; k_el[1]=0x77 as u8; k_el[2]=0x65 as u8; k_el[3]=0x6e as u8 91 k_el[4]=0x32 as u8; k_el[5]=0x2e as u8; k_el[6]=0x65 as u8; k_el[7]=0x6d as u8 92 k_el[8]=0x62 as u8; k_el[9]=0x65 as u8; k_el[10]=0x64 as u8; k_el[11]=0x64 as u8 93 k_el[12]=0x69 as u8; k_el[13]=0x6e as u8; k_el[14]=0x67 as u8; k_el[15]=0x5f as u8 94 k_el[16]=0x6c as u8; k_el[17]=0x65 as u8; k_el[18]=0x6e as u8; k_el[19]=0x67 as u8 95 k_el[20]=0x74 as u8; k_el[21]=0x68 as u8 96 p = _write_meta_u32(buf, p, k_el, 22, 8) 97 98 // qwen2.attention.head_count = 4 99 let k_hc: *u8 = sys_mmap(26) 100 k_hc[0]=0x71 as u8; k_hc[1]=0x77 as u8; k_hc[2]=0x65 as u8; k_hc[3]=0x6e as u8 101 k_hc[4]=0x32 as u8; k_hc[5]=0x2e as u8; k_hc[6]=0x61 as u8; k_hc[7]=0x74 as u8 102 k_hc[8]=0x74 as u8; k_hc[9]=0x65 as u8; k_hc[10]=0x6e as u8; k_hc[11]=0x74 as u8 103 k_hc[12]=0x69 as u8; k_hc[13]=0x6f as u8; k_hc[14]=0x6e as u8; k_hc[15]=0x2e as u8 104 k_hc[16]=0x68 as u8; k_hc[17]=0x65 as u8; k_hc[18]=0x61 as u8; k_hc[19]=0x64 as u8 105 k_hc[20]=0x5f as u8; k_hc[21]=0x63 as u8; k_hc[22]=0x6f as u8; k_hc[23]=0x75 as u8 106 k_hc[24]=0x6e as u8; k_hc[25]=0x74 as u8 107 p = _write_meta_u32(buf, p, k_hc, 26, 4) 108 109 // qwen2.attention.head_count_kv = 2 110 let k_hckv: *u8 = sys_mmap(29) 111 k_hckv[0]=0x71 as u8; k_hckv[1]=0x77 as u8; k_hckv[2]=0x65 as u8; k_hckv[3]=0x6e as u8 112 k_hckv[4]=0x32 as u8; k_hckv[5]=0x2e as u8; k_hckv[6]=0x61 as u8; k_hckv[7]=0x74 as u8 113 k_hckv[8]=0x74 as u8; k_hckv[9]=0x65 as u8; k_hckv[10]=0x6e as u8; k_hckv[11]=0x74 as u8 114 k_hckv[12]=0x69 as u8; k_hckv[13]=0x6f as u8; k_hckv[14]=0x6e as u8; k_hckv[15]=0x2e as u8 115 k_hckv[16]=0x68 as u8; k_hckv[17]=0x65 as u8; k_hckv[18]=0x61 as u8; k_hckv[19]=0x64 as u8 116 k_hckv[20]=0x5f as u8; k_hckv[21]=0x63 as u8; k_hckv[22]=0x6f as u8; k_hckv[23]=0x75 as u8 117 k_hckv[24]=0x6e as u8; k_hckv[25]=0x74 as u8; k_hckv[26]=0x5f as u8; k_hckv[27]=0x6b as u8 118 k_hckv[28]=0x76 as u8 119 p = _write_meta_u32(buf, p, k_hckv, 29, 2) 120 121 // qwen2.feed_forward_length = 16 122 let k_ffl: *u8 = sys_mmap(25) 123 k_ffl[0]=0x71 as u8; k_ffl[1]=0x77 as u8; k_ffl[2]=0x65 as u8; k_ffl[3]=0x6e as u8 124 k_ffl[4]=0x32 as u8; k_ffl[5]=0x2e as u8; k_ffl[6]=0x66 as u8; k_ffl[7]=0x65 as u8 125 k_ffl[8]=0x65 as u8; k_ffl[9]=0x64 as u8; k_ffl[10]=0x5f as u8; k_ffl[11]=0x66 as u8 126 k_ffl[12]=0x6f as u8; k_ffl[13]=0x72 as u8; k_ffl[14]=0x77 as u8; k_ffl[15]=0x61 as u8 127 k_ffl[16]=0x72 as u8; k_ffl[17]=0x64 as u8; k_ffl[18]=0x5f as u8; k_ffl[19]=0x6c as u8 128 k_ffl[20]=0x65 as u8; k_ffl[21]=0x6e as u8; k_ffl[22]=0x67 as u8; k_ffl[23]=0x74 as u8 129 k_ffl[24]=0x68 as u8 130 p = _write_meta_u32(buf, p, k_ffl, 25, 16) 131 132 // token_embd.weight tensor with dim_0=12, dim_1=8 (vocab=12, hidden=8). 133 let n_te: *u8 = sys_mmap(17) 134 n_te[0]=0x74 as u8; n_te[1]=0x6f as u8; n_te[2]=0x6b as u8; n_te[3]=0x65 as u8 135 n_te[4]=0x6e as u8; n_te[5]=0x5f as u8; n_te[6]=0x65 as u8; n_te[7]=0x6d as u8 136 n_te[8]=0x62 as u8; n_te[9]=0x64 as u8; n_te[10]=0x2e as u8; n_te[11]=0x77 as u8 137 n_te[12]=0x65 as u8; n_te[13]=0x69 as u8; n_te[14]=0x67 as u8; n_te[15]=0x68 as u8 138 n_te[16]=0x74 as u8 139 nx_le_write_u64(buf, p, 17); p = p + 8 140 var ti: i64 = 0 141 while ti < 17 { buf[p + ti] = n_te[ti]; ti = ti + 1 } 142 p = p + 17 143 nx_le_write_u32(buf, p, 2); p = p + 4 // n_dims 144 // ggml convention: token_embd.weight [dim_0=hidden, dim_1=vocab] 145 nx_le_write_u64(buf, p, 8); p = p + 8 // dim_0 = hidden_dim 146 nx_le_write_u64(buf, p, 12); p = p + 8 // dim_1 = vocab_size 147 nx_le_write_u32(buf, p, 0); p = p + 4 // ggml_type F32 148 nx_le_write_u64(buf, p, 0); p = p + 8 // data offset 149 150 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader 151 let v_p: nx_int = nx_gguf_parse(buf, 1024, hdr) 152 if v_p != NX_GGUF_OK { return 20 + v_p } 153 154 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc() 155 let out_err: *i64 = sys_mmap(8) as *i64 156 157 let v: nx_int = nx_f32_llm_read_dims_from_gguf(buf, 1024, hdr, model, out_err) 158 if v != NX_FLD_OK { return 30 + v } 159 160 if model.n_layers != 2 { return 50 } 161 if model.hidden_dim != 8 { return 51 } 162 if model.n_heads != 4 { return 52 } 163 if model.n_kv_heads != 2 { return 53 } 164 if model.head_dim != 2 { return 54 } // 8/4 165 if model.ffn_dim != 16 { return 55 } 166 if model.vocab_size != 12 { return 56 } 167 168 return 0 169}