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

↩ module page · 217 lines · 8287 B
1// nx_kv_cache_test.nx -- incremental autoregressive verification.
2//
3// Build a 2-head cache, append 4 tokens one at a time, run scoring
4// each step, verify the cached scores match a "from scratch" full
5// recompute (oracle bit-exact).  This is the algo-led-correctness
6// gate: O(n) per step must produce the same answer as O(n^2).
7
8import "nx_syscalls.nx"
9import "nx_tier.nx"
10import "nx_kv_cache.nx"
11
12func main() -> nx_int {
  // ===== Tiny cache: 2 heads, head_dim 3, max_seq_len 8 ========
  let kv: *NxKvCache = nx_kv_alloc(2, 3, 8)
  if kv.n_heads != 2 { return 1 }
  if kv.head_dim != 3 { return 2 }
  if kv.max_seq_len != 8 { return 3 }
  if kv.n_filled != 0 { return 4 }
19
  // ===== Bytes-per-token sanity ===============================
  // 2 (K+V) * 2 heads * 3 dim * 8 bytes = 96
  if nx_kv_bytes_per_token(kv) != 96 { return 5 }
23
  // Fill ratio at 0 = 0
  if nx_kv_fill_ratio_q10(kv) != 0 { return 6 }
26
  // ===== Append four "tokens" =================================
  //
  // Token 0: K[h=0] = [1, 0, 0]; V[h=0] = [10, 20, 30]
  //          K[h=1] = [0, 1, 0]; V[h=1] = [40, 50, 60]
  // Token 1: K[h=0] = [0, 1, 0]; V[h=0] = [11, 21, 31]
  //          K[h=1] = [1, 0, 0]; V[h=1] = [41, 51, 61]
  // Token 2: K[h=0] = [0, 0, 1]; V[h=0] = [12, 22, 32]
  //          K[h=1] = [0, 0, 1]; V[h=1] = [42, 52, 62]
  // Token 3: K[h=0] = [1, 1, 0]; V[h=0] = [13, 23, 33]
  //          K[h=1] = [1, 1, 1]; V[h=1] = [43, 53, 63]
37
  let k_buf: *i64 = (sys_mmap(48)) as *i64    // 2*3 = 6 i64 per token
  let v_buf: *i64 = (sys_mmap(48)) as *i64
40
  // Token 0
  k_buf[0] = 1; k_buf[1] = 0; k_buf[2] = 0
  k_buf[3] = 0; k_buf[4] = 1; k_buf[5] = 0
  v_buf[0] = 10; v_buf[1] = 20; v_buf[2] = 30
  v_buf[3] = 40; v_buf[4] = 50; v_buf[5] = 60
  if nx_kv_append(kv, k_buf, v_buf) != 0 { return 10 }
47
  // Token 1
  k_buf[0] = 0; k_buf[1] = 1; k_buf[2] = 0
  k_buf[3] = 1; k_buf[4] = 0; k_buf[5] = 0
  v_buf[0] = 11; v_buf[1] = 21; v_buf[2] = 31
  v_buf[3] = 41; v_buf[4] = 51; v_buf[5] = 61
  if nx_kv_append(kv, k_buf, v_buf) != 1 { return 11 }
54
  // Token 2
  k_buf[0] = 0; k_buf[1] = 0; k_buf[2] = 1
  k_buf[3] = 0; k_buf[4] = 0; k_buf[5] = 1
  v_buf[0] = 12; v_buf[1] = 22; v_buf[2] = 32
  v_buf[3] = 42; v_buf[4] = 52; v_buf[5] = 62
  if nx_kv_append(kv, k_buf, v_buf) != 2 { return 12 }
61
  // Token 3
  k_buf[0] = 1; k_buf[1] = 1; k_buf[2] = 0
  k_buf[3] = 1; k_buf[4] = 1; k_buf[5] = 1
  v_buf[0] = 13; v_buf[1] = 23; v_buf[2] = 33
  v_buf[3] = 43; v_buf[4] = 53; v_buf[5] = 63
  if nx_kv_append(kv, k_buf, v_buf) != 3 { return 13 }
68
  if kv.n_filled != 4 { return 14 }
  // fill_ratio = 4 / 8 = 512 Q10
  if nx_kv_fill_ratio_q10(kv) != 512 { return 15 }
72
  // ===== Verify per-position reads ============================
  //
  // get_k(head=0, pos=2, dim=2) -> 1 (token 2 head 0 K=[0,0,1])
  if nx_kv_get_k(kv, 0, 2, 2) != 1 { return 20 }
  // get_v(head=1, pos=3, dim=1) -> 53 (token 3 head 1 V=[43,53,63])
  if nx_kv_get_v(kv, 1, 3, 1) != 53 { return 21 }
79
  // ===== Score: Q @ K^T over cache ============================
  //
  // Take Q = K[token=3] (so attention is diagonal-ish):
  //   Q[h=0] = [1, 1, 0]
  //   Q[h=1] = [1, 1, 1]
  let q: *i64 = (sys_mmap(48)) as *i64
  q[0] = 1; q[1] = 1; q[2] = 0
  q[3] = 1; q[4] = 1; q[5] = 1
88
  let scores: *i64 = (sys_mmap(64)) as *i64    // 2 heads * 4 positions
  if nx_kv_score(kv, q, scores) != NX_KV_OK { return 30 }
91
  // Expected scores h=0:
  //   pos 0: dot([1,1,0], [1,0,0]) = 1
  //   pos 1: dot([1,1,0], [0,1,0]) = 1
  //   pos 2: dot([1,1,0], [0,0,1]) = 0
  //   pos 3: dot([1,1,0], [1,1,0]) = 2
  if scores[0 * 4 + 0] != 1 { return 31 }
  if scores[0 * 4 + 1] != 1 { return 32 }
  if scores[0 * 4 + 2] != 0 { return 33 }
  if scores[0 * 4 + 3] != 2 { return 34 }
101
  // Expected scores h=1:
  //   pos 0: dot([1,1,1], [0,1,0]) = 1
  //   pos 1: dot([1,1,1], [1,0,0]) = 1
  //   pos 2: dot([1,1,1], [0,0,1]) = 1
  //   pos 3: dot([1,1,1], [1,1,1]) = 3
  if scores[1 * 4 + 0] != 1 { return 41 }
  if scores[1 * 4 + 1] != 1 { return 42 }
  if scores[1 * 4 + 2] != 1 { return 43 }
  if scores[1 * 4 + 3] != 3 { return 44 }
111
  // ===== Apply: weighted V sum ================================
  //
  // Use uniform Q10 weights (256 each = 0.25 in Q10): that's equal-
  // weight averaging across the 4 positions.  Output[h, d] = sum_p
  // (256 * V[h, p, d]) / 1024 = average of V across positions.
  let w: *i64 = (sys_mmap(64)) as *i64
  var wi: nx_int = 0
  while wi < 8 {
      w[wi] = 256
      wi = wi + 1
  }
  let out: *i64 = (sys_mmap(48)) as *i64
  if nx_kv_apply(kv, w, out) != NX_KV_OK { return 50 }
125
  // h=0, d=0: avg(10, 11, 12, 13) = 11.5 -> Q10 sum/1024:
  //   (256*10 + 256*11 + 256*12 + 256*13)/1024 = 11776/1024 = 11
  if out[0 * 3 + 0] != 11 { return 51 }
  // h=0, d=1: avg(20, 21, 22, 23) = 21.5 -> 21
  if out[0 * 3 + 1] != 21 { return 52 }
  // h=0, d=2: avg(30, 31, 32, 33) = 31.5 -> 31
  if out[0 * 3 + 2] != 31 { return 53 }
  // h=1, d=0: avg(40, 41, 42, 43) = 41.5 -> 41
  if out[1 * 3 + 0] != 41 { return 54 }
  if out[1 * 3 + 1] != 51 { return 55 }
  if out[1 * 3 + 2] != 61 { return 56 }
137
  // ===== Incremental == recompute (the algo-led correctness gate)
  //
  // Build a SECOND cache; append the same 4 tokens; score the same
  // Q; expect bit-identical scores.  This proves the incremental
  // KV cache pattern is correct -- same answer as "rebuild from
  // scratch every step."
  let kv2: *NxKvCache = nx_kv_alloc(2, 3, 8)
145
  // Append same data
  k_buf[0] = 1; k_buf[1] = 0; k_buf[2] = 0
  k_buf[3] = 0; k_buf[4] = 1; k_buf[5] = 0
  v_buf[0] = 10; v_buf[1] = 20; v_buf[2] = 30
  v_buf[3] = 40; v_buf[4] = 50; v_buf[5] = 60
  nx_kv_append(kv2, k_buf, v_buf)
  k_buf[0] = 0; k_buf[1] = 1; k_buf[2] = 0
  k_buf[3] = 1; k_buf[4] = 0; k_buf[5] = 0
  v_buf[0] = 11; v_buf[1] = 21; v_buf[2] = 31
  v_buf[3] = 41; v_buf[4] = 51; v_buf[5] = 61
  nx_kv_append(kv2, k_buf, v_buf)
  k_buf[0] = 0; k_buf[1] = 0; k_buf[2] = 1
  k_buf[3] = 0; k_buf[4] = 0; k_buf[5] = 1
  v_buf[0] = 12; v_buf[1] = 22; v_buf[2] = 32
  v_buf[3] = 42; v_buf[4] = 52; v_buf[5] = 62
  nx_kv_append(kv2, k_buf, v_buf)
  k_buf[0] = 1; k_buf[1] = 1; k_buf[2] = 0
  k_buf[3] = 1; k_buf[4] = 1; k_buf[5] = 1
  v_buf[0] = 13; v_buf[1] = 23; v_buf[2] = 33
  v_buf[3] = 43; v_buf[4] = 53; v_buf[5] = 63
  nx_kv_append(kv2, k_buf, v_buf)
167
  let scores2: *i64 = (sys_mmap(64)) as *i64
  nx_kv_score(kv2, q, scores2)
  var s: nx_int = 0
  while s < 2 * 4 {
      if scores[s] != scores2[s] { return 60 + s }
      s = s + 1
  }
175
  // ===== Clear + reset =========================================
  nx_kv_clear(kv)
  if kv.n_filled != 0 { return 70 }
  if nx_kv_fill_ratio_q10(kv) != 0 { return 71 }
180
  // After clear, score on empty cache: no positions -> no work
  let empty_scores: *i64 = (sys_mmap(8)) as *i64
  empty_scores[0] = 9999    // sentinel
  nx_kv_score(kv, q, empty_scores)
  // Sentinel should be untouched (no n_filled iterations ran)
  if empty_scores[0] != 9999 { return 72 }
187
  // Apply on empty cache should zero the output
  let empty_out: *i64 = (sys_mmap(48)) as *i64
  var eo: nx_int = 0
  while eo < 6 { empty_out[eo] = 9999; eo = eo + 1 }
  nx_kv_apply(kv, w, empty_out)
  var eo2: nx_int = 0
  while eo2 < 6 {
      if empty_out[eo2] != 0 { return 73 + eo2 }
      eo2 = eo2 + 1
  }
198
  // ===== Append-full refusal ==================================
  //
  // Fill the cache to capacity; next append returns -1.
  var fill_i: nx_int = 0
  while fill_i < 8 {
      nx_kv_append(kv, k_buf, v_buf)
      fill_i = fill_i + 1
  }
  if kv.n_filled != 8 { return 80 }
  if nx_kv_fill_ratio_q10(kv) != 1024 { return 81 }   // full
  if nx_kv_append(kv, k_buf, v_buf) != 0 - 1 { return 82 }
210
  // ===== Sealed verdict enum ==================================
  if nx_kv_verdict_is_valid(NX_KV_OK) != 1 { return 90 }
  if nx_kv_verdict_is_valid(NX_KV_N_VERDICTS) != 0 { return 91 }
  if nx_kv_verdict_is_valid(0 - 1) != 0 { return 92 }
215
  return 0
217}
nx_kv_cache_test.nx source

`nx_kv_cache_test.nx` source
