code wiki / (root) / nx_llm_run_v2_test.nx
`nx_llm_run_v2_test.nx` source

↩ module page · 243 lines · 9450 B
1// nx_llm_run_v2_test.nx -- end-to-end smoke for nx_llm_run_v2.nx.
2//
3// Builds the FULL Llama-shape GGUF (12 tensors) for n_layers=1 with
4// all-zero F32 weights, builds a tiny BPE vocab, and runs
5// nx_llm_generate_one_v2.  With all-zero weights the forward yields
6// uniform logits and the sampler returns a token in [0, vocab_size).
7//
8// This is the integration proof: every loader + composer brick
9// shipped this session walks end-to-end from text prompt to emitted
10// token ID via pure NishiLang.
11
12import "nx_syscalls.nx"
13import "nx_tier.nx"
14import "nx_le.nx"
15import "nx_tensor.nx"
16import "nx_bpe.nx"
17import "nx_model_spec.nx"
18import "nx_prng.nx"
19import "nx_gguf.nx"
20import "nx_gguf_load.nx"
21import "nx_llm_run_v2.nx"
22
23func _write_ti(buf: *u8, off: i64,
             name: *u8, name_len: i64,
             n_dims: i64, dim_0: i64, dim_1: i64,
             ggml_type: i64, data_offset: i64) -> i64 {
  nx_le_write_u64(buf, off, name_len)
  var o: i64 = off + 8
  var i: i64 = 0
  while i < name_len { buf[o + i] = name[i]; i = i + 1 }
  o = o + name_len
  nx_le_write_u32(buf, o, n_dims); o = o + 4
  nx_le_write_u64(buf, o, dim_0);  o = o + 8
  if n_dims >= 2 { nx_le_write_u64(buf, o, dim_1); o = o + 8 }
  nx_le_write_u32(buf, o, ggml_type); o = o + 4
  nx_le_write_u64(buf, o, data_offset); o = o + 8
  return o
38}
39
40func _mk_blk_name(suffix: *u8, suffix_len: nx_int, out: *u8) -> nx_int {
  out[0]=0x62; out[1]=0x6c; out[2]=0x6b; out[3]=0x2e
  out[4]=0x30   // layer 0
  out[5]=0x2e
  var i: nx_int = 0
  while i < suffix_len {
      out[6 + i] = suffix[i]
      i = i + 1
  }
  return 6 + suffix_len
50}
51
52func main() -> i64 {
  // ----- Verdict gate -----
  var vi: nx_int = 0
  while vi < NX_LR2_N_VERDICTS {
      if nx_lr2_verdict_is_valid(vi) != 1 { return 5 + vi }
      vi = vi + 1
  }
59
  let buf: *u8 = sys_mmap(4096)
  buf[0]=0x47; buf[1]=0x47; buf[2]=0x55; buf[3]=0x46
  buf[4]=3
  nx_le_write_u64(buf, 8, 12)        // 12 tensors total
  nx_le_write_u64(buf, 16, 0)
65
  // ----- Top-level tensor names -----
  let n_te: *u8 = sys_mmap(17)
  n_te[0]=0x74; n_te[1]=0x6f; n_te[2]=0x6b; n_te[3]=0x65
  n_te[4]=0x6e; n_te[5]=0x5f; n_te[6]=0x65; n_te[7]=0x6d
  n_te[8]=0x62; n_te[9]=0x64; n_te[10]=0x2e; n_te[11]=0x77
  n_te[12]=0x65; n_te[13]=0x69; n_te[14]=0x67; n_te[15]=0x68
  n_te[16]=0x74
73
  let n_on: *u8 = sys_mmap(18)
  n_on[0]=0x6f; n_on[1]=0x75; n_on[2]=0x74; n_on[3]=0x70
  n_on[4]=0x75; n_on[5]=0x74; n_on[6]=0x5f; n_on[7]=0x6e
  n_on[8]=0x6f; n_on[9]=0x72; n_on[10]=0x6d; n_on[11]=0x2e
  n_on[12]=0x77; n_on[13]=0x65; n_on[14]=0x69; n_on[15]=0x67
  n_on[16]=0x68; n_on[17]=0x74
80
  let n_ow: *u8 = sys_mmap(13)
  n_ow[0]=0x6f; n_ow[1]=0x75; n_ow[2]=0x74; n_ow[3]=0x70
  n_ow[4]=0x75; n_ow[5]=0x74; n_ow[6]=0x2e; n_ow[7]=0x77
  n_ow[8]=0x65; n_ow[9]=0x69; n_ow[10]=0x67; n_ow[11]=0x68
  n_ow[12]=0x74
86
  // ----- Per-layer tensor suffixes -----
  let s_an: *u8 = sys_mmap(16)
  s_an[0]=0x61; s_an[1]=0x74; s_an[2]=0x74; s_an[3]=0x6e
  s_an[4]=0x5f; s_an[5]=0x6e; s_an[6]=0x6f; s_an[7]=0x72
  s_an[8]=0x6d; s_an[9]=0x2e; s_an[10]=0x77; s_an[11]=0x65
  s_an[12]=0x69; s_an[13]=0x67; s_an[14]=0x68; s_an[15]=0x74
93
  let s_q: *u8 = sys_mmap(13)
  s_q[0]=0x61; s_q[1]=0x74; s_q[2]=0x74; s_q[3]=0x6e
  s_q[4]=0x5f; s_q[5]=0x71; s_q[6]=0x2e; s_q[7]=0x77
  s_q[8]=0x65; s_q[9]=0x69; s_q[10]=0x67; s_q[11]=0x68; s_q[12]=0x74
98
  let s_k: *u8 = sys_mmap(13)
  s_k[0]=0x61; s_k[1]=0x74; s_k[2]=0x74; s_k[3]=0x6e
  s_k[4]=0x5f; s_k[5]=0x6b; s_k[6]=0x2e; s_k[7]=0x77
  s_k[8]=0x65; s_k[9]=0x69; s_k[10]=0x67; s_k[11]=0x68; s_k[12]=0x74
103
  let s_v: *u8 = sys_mmap(13)
  s_v[0]=0x61; s_v[1]=0x74; s_v[2]=0x74; s_v[3]=0x6e
  s_v[4]=0x5f; s_v[5]=0x76; s_v[6]=0x2e; s_v[7]=0x77
  s_v[8]=0x65; s_v[9]=0x69; s_v[10]=0x67; s_v[11]=0x68; s_v[12]=0x74
108
  let s_o: *u8 = sys_mmap(18)
  s_o[0]=0x61; s_o[1]=0x74; s_o[2]=0x74; s_o[3]=0x6e
  s_o[4]=0x5f; s_o[5]=0x6f; s_o[6]=0x75; s_o[7]=0x74
  s_o[8]=0x70; s_o[9]=0x75; s_o[10]=0x74; s_o[11]=0x2e
  s_o[12]=0x77; s_o[13]=0x65; s_o[14]=0x69; s_o[15]=0x67
  s_o[16]=0x68; s_o[17]=0x74
115
  let s_fn: *u8 = sys_mmap(15)
  s_fn[0]=0x66; s_fn[1]=0x66; s_fn[2]=0x6e; s_fn[3]=0x5f
  s_fn[4]=0x6e; s_fn[5]=0x6f; s_fn[6]=0x72; s_fn[7]=0x6d
  s_fn[8]=0x2e; s_fn[9]=0x77; s_fn[10]=0x65; s_fn[11]=0x69
  s_fn[12]=0x67; s_fn[13]=0x68; s_fn[14]=0x74
121
  let s_fg: *u8 = sys_mmap(15)
  s_fg[0]=0x66; s_fg[1]=0x66; s_fg[2]=0x6e; s_fg[3]=0x5f
  s_fg[4]=0x67; s_fg[5]=0x61; s_fg[6]=0x74; s_fg[7]=0x65
  s_fg[8]=0x2e; s_fg[9]=0x77; s_fg[10]=0x65; s_fg[11]=0x69
  s_fg[12]=0x67; s_fg[13]=0x68; s_fg[14]=0x74
127
  let s_fu: *u8 = sys_mmap(13)
  s_fu[0]=0x66; s_fu[1]=0x66; s_fu[2]=0x6e; s_fu[3]=0x5f
  s_fu[4]=0x75; s_fu[5]=0x70; s_fu[6]=0x2e; s_fu[7]=0x77
  s_fu[8]=0x65; s_fu[9]=0x69; s_fu[10]=0x67; s_fu[11]=0x68; s_fu[12]=0x74
132
  let s_fd: *u8 = sys_mmap(15)
  s_fd[0]=0x66; s_fd[1]=0x66; s_fd[2]=0x6e; s_fd[3]=0x5f
  s_fd[4]=0x64; s_fd[5]=0x6f; s_fd[6]=0x77; s_fd[7]=0x6e
  s_fd[8]=0x2e; s_fd[9]=0x77; s_fd[10]=0x65; s_fd[11]=0x69
  s_fd[12]=0x67; s_fd[13]=0x68; s_fd[14]=0x74
138
  let n_an: *u8 = sys_mmap(32); let l_an: nx_int = _mk_blk_name(s_an, 16, n_an)
  let n_q:  *u8 = sys_mmap(32); let l_q:  nx_int = _mk_blk_name(s_q,  13, n_q)
  let n_k:  *u8 = sys_mmap(32); let l_k:  nx_int = _mk_blk_name(s_k,  13, n_k)
  let n_v:  *u8 = sys_mmap(32); let l_v:  nx_int = _mk_blk_name(s_v,  13, n_v)
  let n_o:  *u8 = sys_mmap(32); let l_o:  nx_int = _mk_blk_name(s_o,  18, n_o)
  let n_fn: *u8 = sys_mmap(32); let l_fn: nx_int = _mk_blk_name(s_fn, 15, n_fn)
  let n_fg: *u8 = sys_mmap(32); let l_fg: nx_int = _mk_blk_name(s_fg, 15, n_fg)
  let n_fu: *u8 = sys_mmap(32); let l_fu: nx_int = _mk_blk_name(s_fu, 13, n_fu)
  let n_fd: *u8 = sys_mmap(32); let l_fd: nx_int = _mk_blk_name(s_fd, 15, n_fd)
148
  // Cumulative byte offsets per F32 tensor (each at byte offset
  // relative to data_off):
  //   token_embd  [4,2]  = 32   (off 0)
  //   output_norm [2]    = 8    (off 32)
  //   output      [2,4]  = 32   (off 40)
  //   attn_norm   [2]    = 8    (off 72)
  //   attn_q      [2,2]  = 16   (off 80)
  //   attn_k      [2,2]  = 16   (off 96)
  //   attn_v      [2,2]  = 16   (off 112)
  //   attn_output [2,2]  = 16   (off 128)
  //   ffn_norm    [2]    = 8    (off 144)
  //   ffn_gate    [2,4]  = 32   (off 152)
  //   ffn_up      [2,4]  = 32   (off 184)
  //   ffn_down    [4,2]  = 32   (off 216)
  //   total                     248 bytes
164
  var p: i64 = 24
  p = _write_ti(buf, p, n_te, 17, 2, 4, 2, NX_GGML_TYPE_F32, 0)
  p = _write_ti(buf, p, n_on, 18, 1, 2, 1, NX_GGML_TYPE_F32, 32)
  p = _write_ti(buf, p, n_ow, 13, 2, 2, 4, NX_GGML_TYPE_F32, 40)
  p = _write_ti(buf, p, n_an, l_an, 1, 2, 1, NX_GGML_TYPE_F32, 72)
  p = _write_ti(buf, p, n_q,  l_q,  2, 2, 2, NX_GGML_TYPE_F32, 80)
  p = _write_ti(buf, p, n_k,  l_k,  2, 2, 2, NX_GGML_TYPE_F32, 96)
  p = _write_ti(buf, p, n_v,  l_v,  2, 2, 2, NX_GGML_TYPE_F32, 112)
  p = _write_ti(buf, p, n_o,  l_o,  2, 2, 2, NX_GGML_TYPE_F32, 128)
  p = _write_ti(buf, p, n_fn, l_fn, 1, 2, 1, NX_GGML_TYPE_F32, 144)
  p = _write_ti(buf, p, n_fg, l_fg, 2, 2, 4, NX_GGML_TYPE_F32, 152)
  p = _write_ti(buf, p, n_fu, l_fu, 2, 2, 4, NX_GGML_TYPE_F32, 184)
  p = _write_ti(buf, p, n_fd, l_fd, 2, 4, 2, NX_GGML_TYPE_F32, 216)
178
  let data_off: i64 = (p + 31) / 32 * 32
180
  // Zero out the 248-byte data section.  sys_mmap may not zero by
  // default depending on substrate config, so write explicit zeros.
  var zi: nx_int = 0
  while zi < 248 {
      buf[data_off + zi] = 0
      zi = zi + 1
  }
188
  // ----- Parse the GGUF -----
  let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
  let vp: nx_int = nx_gguf_parse(buf, 4096, hdr)
  if vp != NX_GGUF_OK { return 20 + vp }
  if hdr.n_tensors != 12 { return 30 }
194
  // ----- Model spec -----
  let spec: *NxModelSpec = nx_model_spec_new()
  spec.n_layers    = 1
  spec.hidden_dim  = 2
  spec.n_heads     = 1
  spec.head_dim    = 2
  spec.n_kv_heads  = 1
  spec.ffn_dim     = 4
  spec.vocab_size  = 4
  spec.max_seq_len = 8
  if nx_model_spec_validate(spec) != NX_MS_OK { return 40 }
206
  // ----- Tiny BPE vocab: 4 single-byte tokens 'a' 'b' 'c' 'd' -----
  let bpe: *NxBpeVocab = nx_bpe_vocab_new(256, 16, 8)
  let ba: *u8 = sys_mmap(1); ba[0] = 0x61
  let bb: *u8 = sys_mmap(1); bb[0] = 0x62
  let bc: *u8 = sys_mmap(1); bc[0] = 0x63
  let bd: *u8 = sys_mmap(1); bd[0] = 0x64
  nx_bpe_add_token(bpe, ba, 1)
  nx_bpe_add_token(bpe, bb, 1)
  nx_bpe_add_token(bpe, bc, 1)
  nx_bpe_add_token(bpe, bd, 1)
217
  // Prompt = "a" (1 byte)
  let prompt: *u8 = sys_mmap(1); prompt[0] = 0x61
220
  // PRNG
  let prng: *i64 = sys_mmap(8) as *i64
  nx_prng_init(prng, 0xcafebabe)
224
  // ----- Run! -----
  //   temperature_q10 = 1024 (1.0)
  //   top_k = 4 (entire vocab)
  //   rope_base = 10000
  //   attn_scale_q10 = 724 (~ 1/sqrt(2))
  let next_token: nx_int = nx_llm_generate_one_v2(
      spec, buf, hdr, bpe,
      prompt, 1,
      1024, 4,
      prng,
      10000, 724)
236
  // Negative return = error verdict
  if next_token < 0 { return 100 - next_token }
  // Token must be in [0, vocab_size)
  if next_token >= spec.vocab_size { return 200 }
241
  return 0
243}
nx_llm_run_v2_test.nx source

`nx_llm_run_v2_test.nx` source
