nx_rope.nx source

1// nx_rope.nx -- Rotary Position Embeddings (Su et al. 2021). 2// 3// Position-encoding primitive for modern transformers. Replaces 4// sinusoidal-absolute (Vaswani 2017) and learned-absolute encoding 5// with a multiplicative rotation that injects RELATIVE position into 6// attention scores. RoPE has been adopted by every modern decoder- 7// only LLM: 8// 9// Llama 1/2/3, Mistral, Mixtral, Qwen, Phi-3, Gemma 10// Z-Image / Flux / Stable Diffusion 3 attention 11// Mamba state-space hybrid layers 12// 13// Per the no-skipping cardinal: composes the just-shipped 14// `nx_trig.nx` for cos/sin and `nx_root.nx` for inverse-frequency 15// precompute. No skipping, no inlining trig; full bits-up. 16// 17// ===== Math ======================================================= 18// 19// For a query/key vector q of dimension d (d must be even), pair 20// adjacent dimensions (2i, 2i+1). Each pair has its own frequency 21// theta_i = 1 / base^(2i/d) where base = 10000 (Su's published 22// recommendation). 23// 24// For position m, the rotation angle for pair i is angle_{m,i} = m * 25// theta_i radians. The 2D rotation: 26// 27// q'[2i] = q[2i] * cos(angle_{m,i}) - q[2i+1] * sin(angle_{m,i}) 28// q'[2i+1] = q[2i] * sin(angle_{m,i}) + q[2i+1] * cos(angle_{m,i}) 29// 30// This is a 2x2 orthogonal rotation -- norm-preserving by 31// construction. When Q'K^T is computed downstream, the dot product 32// gains a position-dependent term that depends ONLY on the relative 33// position (m_q - m_k), not the absolute positions. That's the key 34// RoPE property. 35// 36// ===== Inverse-frequency precompute ============================== 37// 38// theta_i = base^(-2i/d) for i in [0, d/2). 39// 40// Computing base^(-2i/d) in i64 + Q10: 41// 42// theta_step = base^(-2/d) = (1/base)^(2/d) 43// = nx_nth_root_q14(Q14^2 / base_q14, d/2) 44// ^^^^^^^^^^^^^ 45// represents (1/base) in Q14 46// 47// then theta_i = theta_step^i in Q14 via repeated multiplication. 48// 49// Caller-friendly: nx_rope_compute_inv_freq fills a length-d/2 50// array with theta_i values in Q10 (radians/position). Caller 51// stores this once at model load; runtime apply uses the cached 52// array. 53// 54// Bits-up composition: 55// nx_trig.nx_cos_rad_q10 / nx_sin_rad_q10 (L2) 56// nx_root.nx_nth_root_q14 + nx_pow_q14 (L2) 57// nx_loop.LoopVerdict (control) 58// 59// genealogy_id: su_2021_roformer_rope + vaswani_2017_sinusoidal_positional + 60// touvron_2023_llama_rope_adoption 61// lineage_id: substrate_rope_v1 62 63// nx_safety_envelope: 64// intended_use: AUTO_APPLIED -- primitive-specific tuning queued 65// sil_target: SIL1 66// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail] 67// verdict: NOT_YET_EVALUATED 68 69import "nx_syscalls.nx" 70import "nx_tier.nx" 71import "nx_loop.nx" 72import "nx_root.nx" 73import "nx_trig.nx" 74 75const NX_ROPE_Q10: nx_int = 1024 76const NX_ROPE_Q14_ONE: nx_int = 16384 77const NX_ROPE_BASE: nx_int = 10000 78 79// ===== Sealed-enum: RopeVerdict =================================== 80 81const NX_ROPE_OK: nx_int = 0 82const NX_ROPE_ERR_ODD_DIM: nx_int = 1 83const NX_ROPE_ERR_BAD_DIM: nx_int = 2 84const NX_ROPE_ERR_BAD_BASE: nx_int = 3 85const NX_ROPE_N_VERDICTS: nx_int = 4 86 87func nx_rope_verdict_is_valid(v: nx_int) -> nx_int { 88 if v < 0 { return 0 } 89 if v >= NX_ROPE_N_VERDICTS { return 0 } 90 return 1 91} 92 93// ===== Inverse-frequency precompute =============================== 94// 95// out_inv_freq: [d/2] i64 Q10 (radians per position, descending) 96// 97// theta_i = base^(-2i/d). Computed by: 98// step_q14 = nth_root_q14(Q14^2 / base_q14, d/2) -- (1/base)^(2/d) in Q14 99// theta_q14[i] = step_q14^i (via repeated mul) 100// theta_q10[i] = theta_q14[i] / 16 -- shift Q14 -> Q10 101// 102// For d=128 base=10000 these values descend from theta_0 = 1.0 to 103// theta_63 = 0.0001 over the d/2 dimensions. 104 105func nx_rope_compute_inv_freq(d: nx_int, base: nx_int, out_inv_freq: *i64) -> nx_int { 106 if d <= 0 { return NX_ROPE_ERR_BAD_DIM } 107 let half: nx_int = d / 2 108 if half * 2 != d { return NX_ROPE_ERR_ODD_DIM } 109 if base <= 1 { return NX_ROPE_ERR_BAD_BASE } 110 111 // (1/base) in Q14: Q14 * Q14 / (base * Q14) = Q14 / base 112 let one_over_base_q14: nx_int = NX_ROPE_Q14_ONE / base 113 // step_q14 = (1/base)^(2/d) = (1/base)^(1/(d/2)) = (d/2)-th root of (1/base) 114 let step_q14: nx_int = nx_nth_root_q14(one_over_base_q14, half) 115 116 // theta_q14[0] = 1.0 = Q14. theta_q14[i+1] = theta_q14[i] * step / Q14. 117 var theta_q14: nx_int = NX_ROPE_Q14_ONE 118 var i: nx_int = 0 119 var iter: nx_int = 0 120 var verdict: nx_int = NX_LOOP_RUNNING 121 let BUDGET: nx_int = half 122 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 123 // Convert Q14 -> Q10 by /16. theta values are in radians/position. 124 out_inv_freq[i] = theta_q14 / 16 125 theta_q14 = (theta_q14 * step_q14) / NX_ROPE_Q14_ONE 126 i = i + 1 127 iter = iter + 1 128 } 129 return NX_ROPE_OK 130} 131 132// ===== Apply rotation to one vector ============================== 133// 134// x: [d] i64 Q10 input 135// position: nx_int token position m 136// inv_freq: [d/2] i64 Q10 precomputed theta values 137// out: [d] i64 Q10 output (in-place supported) 138 139func nx_rope_apply_vector(x: *i64, d: nx_int, position: nx_int, 140 inv_freq: *i64, out: *i64) -> nx_int { 141 let half: nx_int = d / 2 142 if half * 2 != d { return NX_ROPE_ERR_ODD_DIM } 143 144 var i: nx_int = 0 145 var iter: nx_int = 0 146 var verdict: nx_int = NX_LOOP_RUNNING 147 let BUDGET: nx_int = half 148 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 149 let angle_q10: nx_int = position * inv_freq[i] 150 let c: nx_int = nx_cos_rad_q10(angle_q10) 151 let s: nx_int = nx_sin_rad_q10(angle_q10) 152 153 let x_even: i64 = x[2 * i] 154 let x_odd: i64 = x[2 * i + 1] 155 156 out[2 * i] = (x_even * c - x_odd * s) / NX_ROPE_Q10 157 out[2 * i + 1] = (x_even * s + x_odd * c) / NX_ROPE_Q10 158 159 i = i + 1 160 iter = iter + 1 161 } 162 return NX_ROPE_OK 163} 164 165// ===== Apply to a batch of tokens ================================= 166// 167// x: [n_tokens, d] Q10 168// positions: [n_tokens] raw position indices 169// inv_freq: [d/2] Q10 170// out: [n_tokens, d] Q10 (in-place supported) 171 172func nx_rope_apply_batch(x: *i64, n_tokens: nx_int, d: nx_int, 173 positions: *i64, inv_freq: *i64, out: *i64) -> nx_int { 174 var t: nx_int = 0 175 var iter: nx_int = 0 176 var verdict: nx_int = NX_LOOP_RUNNING 177 let BUDGET: nx_int = n_tokens 178 while verdict == NX_LOOP_RUNNING && iter < BUDGET { 179 let row_base: nx_int = t * d 180 let v: nx_int = nx_rope_apply_vector( 181 (x as i64 + row_base * 8) as *i64, 182 d, positions[t], inv_freq, 183 (out as i64 + row_base * 8) as *i64) 184 if v != NX_ROPE_OK { verdict = NX_LOOP_ABORTED } 185 t = t + 1 186 iter = iter + 1 187 } 188 if verdict == NX_LOOP_ABORTED { return NX_ROPE_ERR_BAD_DIM } 189 return NX_ROPE_OK 190} 191 192// ===== Self-test ================================================== 193// 194// Closed-form invariants: 195// 196// (a) Position 0 -> identity (cos 0 = 1, sin 0 = 0; no rotation). 197// (b) Norm preservation: ||x'||^2 = ||x||^2 for any position. 198// (c) inv_freq is monotonically decreasing (theta_i drops as i grows). 199// (d) inv_freq[0] is the slowest-decaying = closest to 1.0 (radians/pos). 200// (e) Apply at +m followed by apply at -m returns to original 201// (within rounding -- rotation is its own inverse). 202 203func main() -> i64 { 204 let d: nx_int = 16 205 let half: nx_int = d / 2 206 207 let inv_freq: *i64 = sys_mmap(half * 8) as *i64 208 let v: nx_int = nx_rope_compute_inv_freq(d, NX_ROPE_BASE, inv_freq) 209 if v != NX_ROPE_OK { return 10 + v } 210 211 // --- (c) Monotone descending --- 212 var ci: nx_int = 1 213 while ci < half { 214 if inv_freq[ci] > inv_freq[ci - 1] { return 20 } 215 ci = ci + 1 216 } 217 // --- (d) inv_freq[0] near Q10 = 1024 (theta_0 = 1.0 radian/pos) --- 218 if inv_freq[0] < 1000 { return 30 } 219 if inv_freq[0] > 1024 { return 31 } 220 221 // --- (a) Identity at position 0 --- 222 let x: *i64 = sys_mmap(d * 8) as *i64 223 let y: *i64 = sys_mmap(d * 8) as *i64 224 var k: nx_int = 0 225 while k < d { x[k] = (k + 1) * 100; k = k + 1 } 226 nx_rope_apply_vector(x, d, 0, inv_freq, y) 227 var c0: nx_int = 0 228 while c0 < d { 229 if y[c0] != x[c0] { return 40 } 230 c0 = c0 + 1 231 } 232 233 // --- (b) Norm preservation at position 7 --- 234 nx_rope_apply_vector(x, d, 7, inv_freq, y) 235 var sum_x_sq: i64 = 0 236 var sum_y_sq: i64 = 0 237 var cn: nx_int = 0 238 while cn < d { 239 sum_x_sq = sum_x_sq + x[cn] * x[cn] 240 sum_y_sq = sum_y_sq + y[cn] * y[cn] 241 cn = cn + 1 242 } 243 // Allow ~1% drift due to Q10 trig precision. 244 let drift: i64 = sum_y_sq - sum_x_sq 245 let bound: i64 = sum_x_sq / 50 246 if drift > bound { return 50 } 247 if drift < -bound { return 51 } 248 249 // --- (e) Inverse rotation: apply +m then -m == identity (approx) --- 250 let z: *i64 = sys_mmap(d * 8) as *i64 251 nx_rope_apply_vector(x, d, 3, inv_freq, y) 252 nx_rope_apply_vector(y, d, -3, inv_freq, z) 253 var ce: nx_int = 0 254 while ce < d { 255 let r: i64 = z[ce] - x[ce] 256 if r > 8 { return 60 } 257 if r < -8 { return 61 } 258 ce = ce + 1 259 } 260 261 // --- (f) Verdict gate --- 262 var vi: nx_int = 0 263 while vi < NX_ROPE_N_VERDICTS { 264 if nx_rope_verdict_is_valid(vi) != 1 { return 70 + vi } 265 vi = vi + 1 266 } 267 268 return 0 269}