nx_attention_test.nx source

1// nx_attention_test.nx -- attention kernel composes the stack. 2 3import "nx_syscalls.nx" 4import "nx_tier.nx" 5import "nx_tensor.nx" 6import "nx_attention.nx" 7import "nx_numeric_oracle.nx" 8 9func main() -> nx_int { 10 let err: *i64 = (sys_mmap(8)) as *i64 11 let idx: *i64 = (sys_mmap(64)) as *i64 12 13 // ===== Build a 2-token query against a 3-token key/value cache = 14 // 15 // n_tokens = 2 (the queries we want to attend FROM) 16 // n_kv = 3 (the past tokens we attend TO) 17 // head_dim = 4 18 19 let sh_q: *i64 = (sys_mmap(16)) as *i64 20 sh_q[0] = 2; sh_q[1] = 4 21 let q: *NxTensor = nx_t_alloc(NX_DT_I64, sh_q, 2, err) 22 nx_t_fill_zero(q) 23 let pq: *i64 = q.storage as *i64 24 // Q[0] = [1, 0, 0, 0] 25 pq[0] = 1 26 // Q[1] = [0, 1, 0, 0] 27 pq[5] = 1 28 29 let sh_k: *i64 = (sys_mmap(16)) as *i64 30 sh_k[0] = 3; sh_k[1] = 4 31 let k: *NxTensor = nx_t_alloc(NX_DT_I64, sh_k, 2, err) 32 nx_t_fill_zero(k) 33 let pk: *i64 = k.storage as *i64 34 // K[0] = [4, 0, 0, 0] -> Q[0] @ K[0]^T = 4 35 pk[0] = 4 36 // K[1] = [0, 8, 0, 0] -> Q[1] @ K[1]^T = 8 37 pk[5] = 8 38 // K[2] = [2, 2, 0, 0] -> Q[0] @ K[2]^T = 2; Q[1] @ K[2]^T = 2 39 pk[8] = 2; pk[9] = 2 40 41 let sh_scores: *i64 = (sys_mmap(16)) as *i64 42 sh_scores[0] = 2; sh_scores[1] = 3 43 let scores: *NxTensor = nx_t_alloc(NX_DT_I64, sh_scores, 2, err) 44 45 // ===== Step 1: score matrix Q @ K^T ========================== 46 if nx_attn_score_matrix(q, k, scores) != NX_ATTN_OK { return 1 } 47 let ps: *i64 = scores.storage as *i64 48 if ps[0 * 3 + 0] != 4 { return 2 } // Q[0].K[0] 49 if ps[0 * 3 + 1] != 0 { return 3 } // Q[0].K[1] 50 if ps[0 * 3 + 2] != 2 { return 4 } // Q[0].K[2] 51 if ps[1 * 3 + 0] != 0 { return 5 } // Q[1].K[0] 52 if ps[1 * 3 + 1] != 8 { return 6 } // Q[1].K[1] 53 if ps[1 * 3 + 2] != 2 { return 7 } // Q[1].K[2] 54 55 // ===== Step 2: scale (use Q10 = 512 = 0.5) =================== 56 nx_attn_scale_q10(scores, 512) 57 if ps[0 * 3 + 0] != 2 { return 10 } // 4 * 512 / 1024 = 2 58 if ps[1 * 3 + 1] != 4 { return 11 } // 8 * 512 / 1024 = 4 59 60 // ===== Step 3: softmax per row (Q10 probabilities) ========== 61 nx_attn_softmax_row_q10(scores) 62 // Each row should sum to ~Q10 (1024) within rounding tolerance 63 var row0_sum: nx_int = ps[0] + ps[1] + ps[2] 64 var row1_sum: nx_int = ps[3] + ps[4] + ps[5] 65 // Allow small rounding error (the per-row rounding pads by n_cols) 66 if row0_sum < 1020 { return 20 } 67 if row0_sum > 1028 { return 21 } 68 if row1_sum < 1020 { return 22 } 69 if row1_sum > 1028 { return 23 } 70 71 // The cell with the highest score gets the highest probability -- but only down to the 72 // softmax's Q10 RESOLUTION. Logit gaps of 1-2 Q10 (score 2 vs 0, 1 vs 0, 4 vs 1 = 73 // ~0.001-0.004) sit BELOW resolution: the 16-entry exp LUT distinguishes them (1024 vs 74 // 1023) yet the Q10 normalize (* Q10 / sum, integer floor) rounds both to equal (~341). 75 // So near-tied positions assert MONOTONICITY (>=, via `<`); the one resolvable gap below 76 // (ps[4] score 4 vs ps[3] score 0 -> 341 vs 340) keeps the STRICT check. A real inversion 77 // (higher score -> strictly lower prob) still fails either form. 78 // Row 0: scores [2, 0, 1] -> positions 0 and 2 (>0) tie with position 1 at this resolution. 79 if ps[0] < ps[1] { return 24 } 80 if ps[2] < ps[1] { return 25 } 81 // Row 1: scores [0, 4, 1] -> Q[1].K[1] (highest) dominates; STRICT vs the 0, >= vs the 1. 82 if ps[4] <= ps[3] { return 26 } 83 if ps[4] < ps[5] { return 27 } 84 85 // ===== Step 4: apply to V ==================================== 86 // 87 // V[0] = [10, 20, 30, 40] 88 // V[1] = [50, 60, 70, 80] 89 // V[2] = [90, 100, 110, 120] 90 let v: *NxTensor = nx_t_alloc(NX_DT_I64, sh_k, 2, err) 91 let pv: *i64 = v.storage as *i64 92 pv[0] = 10; pv[1] = 20; pv[2] = 30; pv[3] = 40 93 pv[4] = 50; pv[5] = 60; pv[6] = 70; pv[7] = 80 94 pv[8] = 90; pv[9] = 100; pv[10] = 110; pv[11] = 120 95 96 let sh_out: *i64 = (sys_mmap(16)) as *i64 97 sh_out[0] = 2; sh_out[1] = 4 98 let out: *NxTensor = nx_t_alloc(NX_DT_I64, sh_out, 2, err) 99 if nx_attn_apply_to_v(scores, v, out) != NX_ATTN_OK { return 30 } 100 let po: *i64 = out.storage as *i64 101 102 // Row 0: weights ~= [w0, w1, w2] with w0+w2 ~= 1024-w1 dominating 103 // (since K[0] and K[2] tied at score 2) 104 // Output is roughly an average of V[0] and V[2] weighted by w0, w2. 105 // V[0][0]=10, V[2][0]=90 -> expected out[0][0] is somewhere 106 // between 10 and 90, closer to the midpoint (~50). 107 if po[0 * 4 + 0] < 20 { return 40 } 108 if po[0 * 4 + 0] > 80 { return 41 } 109 110 // Row 1: weights mostly on K[1] (largest score), so output should 111 // be close to V[1] = [50, 60, 70, 80]. 112 if po[1 * 4 + 0] < 30 { return 42 } 113 if po[1 * 4 + 0] > 60 { return 43 } // dominated by V[1] 114 if po[1 * 4 + 3] < 50 { return 44 } 115 if po[1 * 4 + 3] > 90 { return 45 } 116 117 // ===== Forward pass: same answer as step-by-step ============== 118 // 119 // Re-create scores tensor since the step-by-step path consumed 120 // it in place; rebuild from fresh Q/K/V. 121 let q2: *NxTensor = nx_t_alloc(NX_DT_I64, sh_q, 2, err) 122 let pq2: *i64 = q2.storage as *i64 123 pq2[0] = 1; pq2[5] = 1 124 let k2: *NxTensor = nx_t_alloc(NX_DT_I64, sh_k, 2, err) 125 let pk2: *i64 = k2.storage as *i64 126 pk2[0] = 4; pk2[5] = 8; pk2[8] = 2; pk2[9] = 2 127 let v2: *NxTensor = nx_t_alloc(NX_DT_I64, sh_k, 2, err) 128 let pv2: *i64 = v2.storage as *i64 129 pv2[0] = 10; pv2[1] = 20; pv2[2] = 30; pv2[3] = 40 130 pv2[4] = 50; pv2[5] = 60; pv2[6] = 70; pv2[7] = 80 131 pv2[8] = 90; pv2[9] = 100; pv2[10] = 110; pv2[11] = 120 132 133 let out_fw: *NxTensor = nx_t_alloc(NX_DT_I64, sh_out, 2, err) 134 if nx_attn_forward(q2, k2, v2, out_fw, 512) != NX_ATTN_OK { return 50 } 135 136 // Verify via oracle: step-by-step output == forward output (bit-exact) 137 let witness: *i64 = (sys_mmap(NX_NO_WITNESS_FIELDS * 8)) as *i64 138 let v_fw: nx_int = nx_no_check_bit_exact_i64(out, out_fw, witness) 139 if v_fw != NX_NO_VERDICT_EQUAL { return 60 + v_fw } 140 141 // ===== exp(x) sanity ======================================== 142 if _attn_exp_q10(0) != 1024 { return 70 } // exp(0) = 1 143 let e_neg2: nx_int = _attn_exp_q10(0 - 2048) // exp(-2.0) 144 if e_neg2 < 130 { return 71 } 145 if e_neg2 > 145 { return 72 } 146 if _attn_exp_q10(0 - 10240) != 0 { return 73 } // far below -> 0 147 if _attn_exp_q10(99999) != 1024 { return 74 } // x>0 clamped to 1 148 149 // ===== Sealed verdict band coverage ========================= 150 if nx_attn_verdict_is_valid(NX_ATTN_OK) != 1 { return 80 } 151 if nx_attn_verdict_is_valid(NX_ATTN_N_VERDICTS) != 0 { return 81 } 152 if nx_attn_verdict_is_valid(0 - 1) != 0 { return 82 } 153 154 return 0 155}