nx_f32_sampler_test.nx source

1// nx_f32_sampler_test.nx -- smoke for nx_f32_sampler.nx. 2 3import "nx_syscalls.nx" 4import "nx_tier.nx" 5import "nx_prng.nx" 6import "nx_f32.nx" 7import "nx_f32_sampler.nx" 8 9func main() -> i64 { 10 var vi: nx_int = 0 11 while vi < NX_FSP_N_VERDICTS { 12 if nx_fsp_verdict_is_valid(vi) != 1 { return 5 + vi } 13 vi = vi + 1 14 } 15 16 // Logits: [1.0, 2.0, 3.0, 4.0] 17 let logits: *i64 = sys_mmap(4 * 8) as *i64 18 logits[0] = 0x3F800000 19 logits[1] = 0x40000000 20 logits[2] = 0x40400000 21 logits[3] = 0x40800000 22 23 // ===== Argmax ===== 24 let amax: nx_int = nx_f32_sampler_argmax(logits, 4) 25 if amax != 3 { return 10 } 26 27 // Negative logits. 28 let neg_logits: *i64 = sys_mmap(4 * 8) as *i64 29 neg_logits[0] = 0xC0800000 // -4 30 neg_logits[1] = 0xC0400000 // -3 31 neg_logits[2] = 0xC0000000 // -2 32 neg_logits[3] = 0xBF800000 // -1 33 let amax2: nx_int = nx_f32_sampler_argmax(neg_logits, 4) 34 if amax2 != 3 { return 11 } 35 36 // ===== Temperature sample ===== 37 let prng: *i64 = sys_mmap(8) as *i64 38 nx_prng_init(prng, 42 as i64) 39 40 // temp=1.0, inv_temp=1.0 41 let one: i64 = 0x3F800000 42 var ti: nx_int = 0 43 while ti < 20 { 44 let id: nx_int = nx_f32_sampler_sample_temp(logits, 4, one, prng) 45 if id < 0 { return 20 + ti } 46 if id >= 4 { return 50 + ti } 47 ti = ti + 1 48 } 49 50 // High inv_temp (low temperature) -> argmax collapse. inv_temp = 100.0 = 0x42C80000. 51 let inv_temp_100: i64 = 0x42C80000 52 var lt_count: nx_int = 0 53 var t2: nx_int = 0 54 while t2 < 20 { 55 let id: nx_int = nx_f32_sampler_sample_temp(logits, 4, inv_temp_100, prng) 56 if id == 3 { lt_count = lt_count + 1 } 57 t2 = t2 + 1 58 } 59 // With inv_temp=100 the softmax distribution heavily peaks at the argmax. 60 // We expect >= 18/20 hits on id=3 (well above pure-chance 5/20 = 25%). 61 if lt_count < 15 { return 80 } 62 63 // ===== Top-k ===== 64 // top-k=1 with low temp -> deterministic argmax. 65 let id_k1: nx_int = nx_f32_sampler_sample_top_k(logits, 4, 1, one, prng) 66 if id_k1 != 3 { return 90 } 67 68 // top-k=2 with low temp -> id in {2, 3}. 69 var hi: nx_int = 0 70 while hi < 10 { 71 let id: nx_int = nx_f32_sampler_sample_top_k(logits, 4, 2, one, prng) 72 if id != 2 { 73 if id != 3 { return 100 + hi } 74 } 75 hi = hi + 1 76 } 77 78 // top-k >= vocab_size -> behaves like sample_temp. 79 let id_kbig: nx_int = nx_f32_sampler_sample_top_k(logits, 4, 10, one, prng) 80 if id_kbig < 0 { return 120 } 81 if id_kbig >= 4 { return 121 } 82 83 // top-k <= 0 -> falls through to argmax. 84 let id_k0: nx_int = nx_f32_sampler_sample_top_k(logits, 4, 0, one, prng) 85 if id_k0 != 3 { return 130 } 86 87 // ===== Top-p ===== 88 // p_q14 = 0 -> argmax. 89 let id_p0: nx_int = nx_f32_sampler_sample_top_p(logits, 4, 0, one, prng) 90 if id_p0 != 3 { return 140 } 91 92 // p_q14 >= 16384 (= prob >= 1.0) -> no truncation; should sample 93 // from full distribution. Just verify validity. 94 let id_p1: nx_int = nx_f32_sampler_sample_top_p(logits, 4, 16384, one, prng) 95 if id_p1 < 0 { return 150 } 96 if id_p1 >= 4 { return 151 } 97 98 // p_q14 = 4915 (= ~0.3). Logits [1,2,3,4] with temp=1.0 produce probs 99 // ~[0.032, 0.087, 0.237, 0.643]. Cumsum-descending: 0.643 -> stops 100 // at one element (top-1 = id=3, the argmax). Sampling collapses to 3. 101 let id_p_low: nx_int = nx_f32_sampler_sample_top_p(logits, 4, 4915, one, prng) 102 if id_p_low != 3 { return 160 } 103 104 // p_q14 = 13107 (= 0.8). Cumsum hits 0.8 at top-2 (0.643 + 0.237 = 0.880). 105 // Sample should be in {2, 3}. 106 var pp: nx_int = 0 107 while pp < 10 { 108 let id: nx_int = nx_f32_sampler_sample_top_p(logits, 4, 13107, one, prng) 109 if id != 2 { 110 if id != 3 { return 170 + pp } 111 } 112 pp = pp + 1 113 } 114 115 // ===== Repetition penalty ===== 116 // Logits [1, 2, 3, 4]; argmax = 3. Apply penalty=2.0 to recent=[3, 1]: 117 // logits[3] = 4 / 2 = 2.0 118 // logits[1] = 2 / 2 = 1.0 119 // After: logits = [1, 1, 3, 2] -> argmax = 2. 120 let logits_rep: *i64 = sys_mmap(4 * 8) as *i64 121 logits_rep[0] = 0x3F800000 122 logits_rep[1] = 0x40000000 123 logits_rep[2] = 0x40400000 124 logits_rep[3] = 0x40800000 125 let recent: *i64 = sys_mmap(2 * 8) as *i64 126 recent[0] = 3 127 recent[1] = 1 128 let penalty: i64 = 0x40000000 // 2.0 129 nx_f32_sampler_apply_repetition_penalty(logits_rep, 4, recent, 2, penalty) 130 if logits_rep[0] != 0x3F800000 { return 200 } // unchanged 131 if logits_rep[1] != 0x3F800000 { return 201 } // 2.0 / 2 = 1.0 132 if logits_rep[2] != 0x40400000 { return 202 } // unchanged 133 if logits_rep[3] != 0x40000000 { return 203 } // 4.0 / 2 = 2.0 134 let amax_rep: nx_int = nx_f32_sampler_argmax(logits_rep, 4) 135 if amax_rep != 2 { return 210 } // argmax shifted from 3 to 2 136 137 // Negative-logit case: l < 0 -> multiply (more negative -> less likely). 138 let logits_neg: *i64 = sys_mmap(2 * 8) as *i64 139 logits_neg[0] = 0xBF800000 // -1.0 140 logits_neg[1] = 0x3F800000 // 1.0 141 let recent_neg: *i64 = sys_mmap(1 * 8) as *i64 142 recent_neg[0] = 0 143 nx_f32_sampler_apply_repetition_penalty(logits_neg, 2, recent_neg, 1, penalty) 144 if logits_neg[0] != 0xC0000000 { return 220 } // -1.0 * 2 = -2.0 145 if logits_neg[1] != 0x3F800000 { return 221 } // unchanged 146 147 // Out-of-range token id silently skipped. 148 let logits_oor: *i64 = sys_mmap(2 * 8) as *i64 149 logits_oor[0] = 0x3F800000 150 logits_oor[1] = 0x40000000 151 let recent_oor: *i64 = sys_mmap(2 * 8) as *i64 152 recent_oor[0] = 99 153 recent_oor[1] = -1 154 nx_f32_sampler_apply_repetition_penalty(logits_oor, 2, recent_oor, 2, penalty) 155 if logits_oor[0] != 0x3F800000 { return 230 } 156 if logits_oor[1] != 0x40000000 { return 231 } 157 158 // ===== Min-p ===== 159 // Logits [1, 2, 3, 4]; softmax probs ~[0.032, 0.087, 0.237, 0.643] 160 // max_prob ~ 0.643. 161 // min_p=0.5 (q14=8192): threshold = 0.5 * 0.643 = 0.321 162 // Tokens with prob >= 0.321: only id=3 (0.643). k=1, argmax = 3. 163 let id_mp_high: nx_int = nx_f32_sampler_sample_min_p(logits, 4, 8192, one, prng) 164 if id_mp_high != 3 { return 240 } 165 166 // min_p=0.1 (q14=1638): threshold = 0.1 * 0.643 = 0.064 167 // Tokens with prob >= 0.064: ids 1, 2, 3 (0.087, 0.237, 0.643). k=3. 168 var mp: nx_int = 0 169 while mp < 10 { 170 let id: nx_int = nx_f32_sampler_sample_min_p(logits, 4, 1638, one, prng) 171 if id != 1 { 172 if id != 2 { 173 if id != 3 { return 250 + mp } 174 } 175 } 176 mp = mp + 1 177 } 178 179 // min_p=0 -> falls through to plain sample_temp. 180 let id_mp_zero: nx_int = nx_f32_sampler_sample_min_p(logits, 4, 0, one, prng) 181 if id_mp_zero < 0 { return 270 } 182 if id_mp_zero >= 4 { return 271 } 183 184 return 0 185}