nx_speculative_verify_test.nx source

1// nx_speculative_verify_test.nx -- smoke for H3 speculative decoding. 2 3import "nx_syscalls.nx" 4import "nx_speculative_verify.nx" 5 6func main() -> i64 { 7 // ----- 1. Session construction ----- 8 let s: *NxSpeculativeSession = nx_speculative_session_new(2024, 8) 9 if nx_spec_session_is_valid(s) != 1 { return 1 } 10 if nx_spec_session_request_id(s) != 2024 { return 2 } 11 if nx_spec_session_max_k(s) != 8 { return 3 } 12 if nx_spec_total_drafted(s) != 0 { return 4 } 13 if nx_spec_total_accepted(s) != 0 { return 5 } 14 if nx_spec_total_committed(s) != 0 { return 6 } 15 if nx_spec_total_rounds(s) != 0 { return 7 } 16 if nx_spec_acceptance_rate_q16(s) != 0 { return 8 } 17 18 // Workspace buffers. 19 let draft: *i64 = (sys_mmap(128)) as *i64 20 let flags: *i64 = (sys_mmap(128)) as *i64 21 let out: *i64 = (sys_mmap(128)) as *i64 22 23 // ----- 2. All-accept round (8/8) ----- 24 var i: i64 = 0 25 while i < 8 { 26 draft[i] = 100 + i 27 flags[i] = 1 28 i = i + 1 29 } 30 let committed1: i64 = nx_speculative_verify(s, draft, flags, 8, 999, out) 31 // Accepted 8 + 1 bonus = 9 committed 32 if committed1 != 9 { return 9 } 33 // Check committed sequence: 100..107 then 999 34 if out[0] != 100 { return 10 } 35 if out[7] != 107 { return 11 } 36 if out[8] != 999 { return 12 } 37 if nx_spec_total_drafted(s) != 8 { return 13 } 38 if nx_spec_total_accepted(s) != 8 { return 14 } 39 if nx_spec_total_committed(s) != 9 { return 15 } 40 if nx_spec_total_rounds(s) != 1 { return 16 } 41 // Acceptance rate Q16 = 8/8 * 65536 = 65536 42 if nx_spec_acceptance_rate_q16(s) != 65536 { return 17 } 43 44 // ----- 3. Partial-accept round (3/5) ----- 45 let s2: *NxSpeculativeSession = nx_speculative_session_new(7, 8) 46 draft[0] = 200 47 draft[1] = 201 48 draft[2] = 202 49 draft[3] = 203 50 draft[4] = 204 51 flags[0] = 1 52 flags[1] = 1 53 flags[2] = 1 54 flags[3] = 0 // reject at index 3 55 flags[4] = 1 // -- ignored after first reject 56 let committed2: i64 = nx_speculative_verify(s2, draft, flags, 5, 777, out) 57 // Accepted 3 + 1 bonus = 4 committed 58 if committed2 != 4 { return 18 } 59 if out[0] != 200 { return 19 } 60 if out[1] != 201 { return 20 } 61 if out[2] != 202 { return 21 } 62 if out[3] != 777 { return 22 } 63 if nx_spec_total_drafted(s2) != 5 { return 23 } 64 if nx_spec_total_accepted(s2) != 3 { return 24 } 65 if nx_spec_total_committed(s2) != 4 { return 25 } 66 67 // ----- 4. All-reject round (0/4) ----- 68 flags[0] = 0 69 flags[1] = 0 70 flags[2] = 0 71 flags[3] = 0 72 let committed3: i64 = nx_speculative_verify(s2, draft, flags, 4, 555, out) 73 // Accepted 0 + 1 bonus = 1 committed 74 if committed3 != 1 { return 26 } 75 if out[0] != 555 { return 27 } 76 if nx_spec_total_drafted(s2) != 9 { return 28 } 77 if nx_spec_total_accepted(s2) != 3 { return 29 } 78 if nx_spec_total_committed(s2) != 5 { return 30 } 79 if nx_spec_total_rounds(s2) != 2 { return 31 } 80 81 // ----- 5. Acceptance rate Q16 = 3/9 * 65536 = 21845 ----- 82 if nx_spec_acceptance_rate_q16(s2) != 21845 { return 32 } 83 84 // ----- 6. Single-token round (k=1) ----- 85 let s3: *NxSpeculativeSession = nx_speculative_session_new(8, 4) 86 draft[0] = 42 87 flags[0] = 1 88 let committed4: i64 = nx_speculative_verify(s3, draft, flags, 1, 88, out) 89 if committed4 != 2 { return 33 } 90 if out[0] != 42 { return 34 } 91 if out[1] != 88 { return 35 } 92 93 // ----- 7. Bad k_actual ----- 94 if nx_speculative_verify(s, draft, flags, 0, 0, out) != (0 - NX_SPEC_BAD_K) { return 36 } 95 if nx_speculative_verify(s, draft, flags, -1, 0, out) != (0 - NX_SPEC_BAD_K) { return 37 } 96 // session max_k is 8; k=9 is over the per-session ceiling. 97 if nx_speculative_verify(s, draft, flags, 9, 0, out) != (0 - NX_SPEC_BAD_K) { return 38 } 98 // Even if session.max_k allowed it, NX_SPEC_MAX_K caps at 16. 99 let s_big: *NxSpeculativeSession = nx_speculative_session_new(9, NX_SPEC_MAX_K) 100 if nx_speculative_verify(s_big, draft, flags, NX_SPEC_MAX_K + 1, 0, out) != (0 - NX_SPEC_BAD_K) { return 39 } 101 102 // ----- 8. Null pointers rejected ----- 103 let null_p: *i64 = (0 as i64) as *i64 104 if nx_speculative_verify(s, null_p, flags, 4, 0, out) != (0 - NX_SPEC_BAD_INPUT) { return 40 } 105 if nx_speculative_verify(s, draft, null_p, 4, 0, out) != (0 - NX_SPEC_BAD_INPUT) { return 41 } 106 if nx_speculative_verify(s, draft, flags, 4, 0, null_p) != (0 - NX_SPEC_BAD_INPUT) { return 42 } 107 108 // ----- 9. Bad session_new inputs ----- 109 if (nx_speculative_session_new(1, 0) as i64) != 0 { return 43 } 110 if (nx_speculative_session_new(1, -1) as i64) != 0 { return 44 } 111 if (nx_speculative_session_new(1, NX_SPEC_MAX_K + 1) as i64) != 0 { return 45 } 112 113 // ----- 10. Tamper detection ----- 114 let tamper_s: *NxSpeculativeSession = nx_speculative_session_new(11, 4) 115 tamper_s.canary_post = 0xDEADBEEF 116 if nx_spec_session_is_valid(tamper_s) != 0 { return 46 } 117 if nx_speculative_verify(tamper_s, draft, flags, 4, 0, out) != (0 - NX_SPEC_TAMPER) { return 47 } 118 if nx_spec_total_drafted(tamper_s) != -1 { return 48 } 119 if nx_spec_total_accepted(tamper_s) != -1 { return 49 } 120 if nx_spec_acceptance_rate_q16(tamper_s) != -1 { return 50 } 121 122 // ----- 11. Sealed-enum gates ----- 123 if nx_spec_verdict_is_valid(NX_SPEC_OK) != 1 { return 51 } 124 if nx_spec_verdict_is_valid(NX_SPEC_TAMPER) != 1 { return 52 } 125 if nx_spec_verdict_is_valid(-1) != 0 { return 53 } 126 if nx_spec_verdict_is_valid(NX_SPEC_N_VERDICTS) != 0 { return 54 } 127 128 // ----- 12. Monotonic stats invariant ----- 129 // After several mixed rounds, accepted <= drafted, rounds <= committed. 130 if s2.total_accepted > s2.total_drafted { return 55 } 131 if s2.total_rounds > s2.total_committed { return 56 } 132 // committed == accepted + rounds (one bonus per round). 133 if s2.total_committed != (s2.total_accepted + s2.total_rounds) { return 57 } 134 135 // ----- 13. Reject-at-first round (0 accepted but valid) ----- 136 let s4: *NxSpeculativeSession = nx_speculative_session_new(12, 4) 137 flags[0] = 0 138 flags[1] = 1 139 flags[2] = 1 140 flags[3] = 1 141 draft[0] = 50 142 draft[1] = 51 143 draft[2] = 52 144 draft[3] = 53 145 let committed5: i64 = nx_speculative_verify(s4, draft, flags, 4, 999, out) 146 if committed5 != 1 { return 58 } 147 if out[0] != 999 { return 59 } 148 if nx_spec_total_accepted(s4) != 0 { return 60 } 149 if nx_spec_acceptance_rate_q16(s4) != 0 { return 61 } // 0/4 * 65536 = 0 150 151 return 0 152}