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

↩ module page · 207 lines · 8025 B
1// nx_transformer_stack_test.nx -- smoke for nx_transformer_stack.nx.
2//
3// Builds a 1-layer synthetic Llama-shape GGUF with ALL-ZERO weights
4// (so the forward pass becomes an identity), feeds x_in through the
5// stack, and verifies x_in is unchanged.
6//
7// Test spec: hidden=2, head_dim=2, ffn_dim=4, n_tokens=1.
8// Weights = 0 means:
9//   - RMSNorm(x, 0) = 0           (gamma=0 zeros out the norm output)
10//   - matmul(0, W) = 0            (zero input/weight = zero output)
11//   - attention(0, 0, 0) = 0      (softmax of all-zero scores * zero V)
12//   - residual: x + 0 = x          (UNCHANGED)
13//   - FFN: same identity-via-zero pattern
14// So output x == input x after the entire forward pass.
15
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_le.nx"
19import "nx_tensor.nx"
20import "nx_gguf.nx"
21import "nx_gguf_load.nx"
22import "nx_transformer_stack.nx"
23
24// Write tensor_info entry; return new offset.
25func _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
40}
41
42// Build a layer-0 tensor name "blk.0.{suffix}" into buf-of-size 32.
43// Returns name length.
44func _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   // '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
54}
55
56func main() -> i64 {
  // ----- Verdict gate -----
  var vi: nx_int = 0
  while vi < NX_TS_N_VERDICTS {
      if nx_ts_verdict_is_valid(vi) != 1 { return 5 + vi }
      vi = vi + 1
  }
63
  let buf: *u8 = sys_mmap(2048)
  buf[0]=0x47; buf[1]=0x47; buf[2]=0x55; buf[3]=0x46
  buf[4]=3
  nx_le_write_u64(buf, 8, 9)         // 9 tensors
  nx_le_write_u64(buf, 16, 0)
69
  // ----- Suffix constants -----
  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
76
  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
82
  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
88
  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
94
  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
101
  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
107
  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
113
  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
119
  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
125
  // Full names "blk.0.{suffix}"
  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)
136
  // ----- Tensor sizes (F32, in bytes) -----
  //   attn_norm   [2]   = 8
  //   attn_q/k/v  [2,2] = 16 each
  //   attn_output [2,2] = 16
  //   ffn_norm    [2]   = 8
  //   ffn_gate    [2,4] = 32
  //   ffn_up      [2,4] = 32
  //   ffn_down    [4,2] = 32
  // Cumulative offsets in data section: 0, 8, 24, 40, 56, 72, 80, 112, 144
146
  // Write tensor_info entries (each F32).
  var p: i64 = 24
  p = _write_ti(buf, p, n_an, l_an, 1, 2, 1, NX_GGML_TYPE_F32, 0)
  p = _write_ti(buf, p, n_q,  l_q,  2, 2, 2, NX_GGML_TYPE_F32, 8)
  p = _write_ti(buf, p, n_k,  l_k,  2, 2, 2, NX_GGML_TYPE_F32, 24)
  p = _write_ti(buf, p, n_v,  l_v,  2, 2, 2, NX_GGML_TYPE_F32, 40)
  p = _write_ti(buf, p, n_o,  l_o,  2, 2, 2, NX_GGML_TYPE_F32, 56)
  p = _write_ti(buf, p, n_fn, l_fn, 1, 2, 1, NX_GGML_TYPE_F32, 72)
  p = _write_ti(buf, p, n_fg, l_fg, 2, 2, 4, NX_GGML_TYPE_F32, 80)
  p = _write_ti(buf, p, n_fu, l_fu, 2, 2, 4, NX_GGML_TYPE_F32, 112)
  p = _write_ti(buf, p, n_fd, l_fd, 2, 4, 2, NX_GGML_TYPE_F32, 144)
158
  // data_off = align_up(p, 32).
  let data_off: i64 = (p + 31) / 32 * 32
161
  // Tensor data: all zeros (sys_mmap already returns zeroed memory
  // on the substrate's mmap; but we write zeros explicitly for the
  // 176 bytes that hold weight values).
  var zi: nx_int = 0
  while zi < 176 {
      buf[data_off + zi] = 0
      zi = zi + 1
  }
170
  let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
  let vp: nx_int = nx_gguf_parse(buf, 2048, hdr)
  if vp != NX_GGUF_OK { return 20 + vp }
  if hdr.n_tensors != 9 { return 30 }
175
  // ----- Build x_in [1, 2] = (42, 17) Q10 -----
  let sh: *i64 = sys_mmap(2 * 8) as *i64
  sh[0] = 1; sh[1] = 2
  let err: *i64 = sys_mmap(8) as *i64
  err[0] = 0
  let x: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 2, err)
  if err[0] != 0 { return 40 }
  let px: *i64 = x.storage as *i64
  px[0] = 42 * 1024     // 42 in Q10
  px[1] = 17 * 1024     // 17 in Q10
186
  // positions [1] = (0)
  let pos: *i64 = sys_mmap(8) as *i64
  pos[0] = 0
190
  // ----- Call the stack forward -----
  //   n_layers = 1, hidden = 2, head_dim = 2, ffn_dim = 4
  //   rope_base = 10000, attn_scale_q10 = 724 (~ 1/sqrt(2))
  //   final_norm_gamma = null (skip final norm; we want pure identity)
  let v: nx_int = nx_transformer_stack_forward(
      x, pos, buf, hdr,
      1, 2, 2, 4,
      10000, 724,
      0 as *i64)
  if v != NX_TS_OK { return 100 + v }
201
  // ----- Verify x is unchanged (all-zero weights = identity) -----
  if px[0] != 42 * 1024 { return 200 }
  if px[1] != 17 * 1024 { return 201 }
205
  return 0
207}
nx_transformer_stack_test.nx source

`nx_transformer_stack_test.nx` source
