nx_f32_rope.nx source
↩ module page · 186 lines · 7719 B
1// nx_f32_rope.nx -- bits-up f32 RoPE (Rotary Position Embedding).
2//
3// L7 composition brick. Composes L6 sin + cos + exp + L4 mul/add/
4// sub/div/cvt. No libm.
5//
6// Reference (Su 2021, "RoFormer"):
7// For each adjacent dim-pair (x_{2i}, x_{2i+1}) at position m:
8// theta_i = base^(-2i/d)
9// [x'_{2i}, x'_{2i+1}] = R(m*theta_i) @ [x_{2i}, x_{2i+1}]
10// where R(phi) is the 2x2 rotation matrix [[cos,-sin],[sin,cos]].
11//
12// Concretely:
13// x'_{2i} = x_{2i} * cos(m*theta_i) - x_{2i+1} * sin(m*theta_i)
14// x'_{2i+1} = x_{2i} * sin(m*theta_i) + x_{2i+1} * cos(m*theta_i)
15//
16// base is typically 10000 (Llama-2) or 500000 (Llama-3 / Qwen2).
17//
18// Caller passes log(base) as f32 so it's computed once across all
19// positions. log(10000) ~= 9.2103404 = 0x4113411D in f32 nearest.
20//
21// genealogy_id: su_2021_rope + standard_rotation_compose
22// lineage_id: substrate_f32_rope_v1
23
24import "nx_syscalls.nx"
25import "nx_tier.nx"
26import "nx_f32.nx"
27import "nx_f32_div.nx"
28import "nx_f32_cvt.nx"
29import "nx_f32_exp.nx"
30import "nx_f32_sincos.nx"
31
32const NX_F32_ROPE_OK: nx_int = 0
33const NX_F32_ROPE_ERR_BAD_DIM: nx_int = 1
34const NX_F32_ROPE_ERR_ODD_DIM: nx_int = 2
35const NX_F32_ROPE_N_VERDICTS: nx_int = 3
36
37func nx_f32_rope_verdict_is_valid(v: nx_int) -> nx_int {
38 if v < 0 { return 0 }
39 if v >= NX_F32_ROPE_N_VERDICTS { return 0 }
40 return 1
41}
42
43// Apply RoPE in-place to a single head vector x[head_dim] at the
44// given integer position. head_dim must be even (RoPE pairs up
45// adjacent dims).
46//
47// log_base_f32 is the precomputed natural log of the rope base
48// (e.g. for base=10000, log_base ~= 9.2103404 -> 0x4113411D).
49// Caller computes once via nx_f32_log(base_f32).
50
51func nx_f32_rope_apply_vector(x: *i64, head_dim: nx_int,
52 position: nx_int,
53 log_base_f32: i64) -> nx_int {
54 if head_dim <= 0 { return NX_F32_ROPE_ERR_BAD_DIM }
55 let half: nx_int = head_dim / 2
56 if half * 2 != head_dim { return NX_F32_ROPE_ERR_ODD_DIM }
57
58 // Pre-compute d as f32 (used in theta exponent denominator).
59 let d_f32: i64 = nx_i32_to_f32(head_dim)
60 let position_f32: i64 = nx_i32_to_f32(position)
61
62 var i: nx_int = 0
63 while i < half {
64 // theta_i = base^(-2i/d) = exp(-2i/d * log(base))
65 let two_i: i64 = nx_i32_to_f32(i + i)
66 let neg_two_i: i64 = nx_f32_neg(two_i)
67 let neg_two_i_div_d: i64 = nx_f32_div(neg_two_i, d_f32)
68 let exponent: i64 = nx_f32_mul(neg_two_i_div_d, log_base_f32)
69 let theta: i64 = nx_f32_exp(exponent)
70
71 // m_theta = position * theta
72 let m_theta: i64 = nx_f32_mul(position_f32, theta)
73
74 let c: i64 = nx_f32_cos(m_theta)
75 let s: i64 = nx_f32_sin(m_theta)
76
77 let pair_idx: nx_int = i + i
78 let x_lo: i64 = x[pair_idx]
79 let x_hi: i64 = x[pair_idx + 1]
80
81 // x'_lo = x_lo * c - x_hi * s
82 // x'_hi = x_lo * s + x_hi * c
83 let new_lo: i64 = nx_f32_sub(nx_f32_mul(x_lo, c), nx_f32_mul(x_hi, s))
84 let new_hi: i64 = nx_f32_add(nx_f32_mul(x_lo, s), nx_f32_mul(x_hi, c))
85 x[pair_idx] = new_lo
86 x[pair_idx + 1] = new_hi
87
88 i = i + 1
89 }
90 return NX_F32_ROPE_OK
91}
92
93// ===== Precomputed (cos,sin) split (2026-07-08) ===================
94// The trig here (theta_i = exp(...), then cos/sin of m*theta_i) depends
95// ONLY on (position, i) -- NOT on the head or its data. A transformer
96// block applies RoPE to ~16 heads at the SAME position, so calling
97// apply_vector_neox per head recomputed the identical exp+cos+sin ~16x
98// (measured: RoPE = 123ms/token, mostly these transcendentals). Split
99// it: BUILD the (c,s) table once per position, APPLY it to every head
100// with pure mul/add (no transcendentals). Bit-exact vs apply_vector_neox
101// (same c,s, same rotation).
102//
103// cs layout: cs[2*i] = cos(m*theta_i), cs[2*i+1] = sin(m*theta_i), for
104// i in [0, head_dim/2). Caller sizes cs to head_dim i64 slots.
105
106func nx_f32_rope_build_cs(cs_out: *i64, head_dim: nx_int,
107 position: nx_int, log_base_f32: i64) -> nx_int {
108 if head_dim <= 0 { return NX_F32_ROPE_ERR_BAD_DIM }
109 let half: nx_int = head_dim / 2
110 if half * 2 != head_dim { return NX_F32_ROPE_ERR_ODD_DIM }
111 let d_f32: i64 = nx_i32_to_f32(head_dim)
112 let position_f32: i64 = nx_i32_to_f32(position)
113 var i: nx_int = 0
114 while i < half {
115 let two_i: i64 = nx_i32_to_f32(i + i)
116 let neg_two_i: i64 = nx_f32_neg(two_i)
117 let neg_two_i_div_d: i64 = nx_f32_div(neg_two_i, d_f32)
118 let exponent: i64 = nx_f32_mul(neg_two_i_div_d, log_base_f32)
119 let theta: i64 = nx_f32_exp(exponent)
120 let m_theta: i64 = nx_f32_mul(position_f32, theta)
121 cs_out[i + i] = nx_f32_cos(m_theta)
122 cs_out[i + i + 1] = nx_f32_sin(m_theta)
123 i = i + 1
124 }
125 return NX_F32_ROPE_OK
126}
127
128// Apply a precomputed (c,s) table (NEOX pairing) to one head vector,
129// in place. No transcendentals -- pure mul/add/sub.
130func nx_f32_rope_apply_cs_neox(x: *i64, head_dim: nx_int, cs: *i64) -> nx_int {
131 if head_dim <= 0 { return NX_F32_ROPE_ERR_BAD_DIM }
132 let half: nx_int = head_dim / 2
133 if half * 2 != head_dim { return NX_F32_ROPE_ERR_ODD_DIM }
134 // HARDWARE __f32_mul/add (2026-07-10): IEEE, bit-identical to the
135 // software path (same op order; a-b done as a + (-b) via sign-flip since
136 // there's no __f32_sub). The per-head rotation, hot every decode.
137 var i: nx_int = 0
138 while i < half {
139 let c: i64 = cs[i + i]
140 let s: i64 = cs[i + i + 1]
141 let x_lo: i64 = x[i]
142 let x_hi: i64 = x[i + half]
143 x[i] = __f32_add(__f32_mul(x_lo, c), __f32_mul(x_hi, s) ^ 0x80000000)
144 x[i + half] = __f32_add(__f32_mul(x_lo, s), __f32_mul(x_hi, c))
145 i = i + 1
146 }
147 return NX_F32_ROPE_OK
148}
149
150// NEOX-style RoPE (GPT-NeoX / Qwen2 / Llama-in-llama.cpp convention). Same theta_i = base^(-2i/d), but pairs
151// the two HALVES of the head vector -- (x[i], x[i+half]) -- instead of adjacent dims (x[2i], x[2i+1]). Qwen2.5
152// GGUF is exported for NEOX rope (llama.cpp LLAMA_ROPE_TYPE_NEOX), so the adjacent variant above rotates the
153// WRONG pairs at every position>0 -> incoherent multi-token output even though position 0 (rope=identity) looks
154// fine. This is the pairing the real weights expect.
155func nx_f32_rope_apply_vector_neox(x: *i64, head_dim: nx_int,
156 position: nx_int,
157 log_base_f32: i64) -> nx_int {
158 if head_dim <= 0 { return NX_F32_ROPE_ERR_BAD_DIM }
159 let half: nx_int = head_dim / 2
160 if half * 2 != head_dim { return NX_F32_ROPE_ERR_ODD_DIM }
161
162 let d_f32: i64 = nx_i32_to_f32(head_dim)
163 let position_f32: i64 = nx_i32_to_f32(position)
164
165 var i: nx_int = 0
166 while i < half {
167 let two_i: i64 = nx_i32_to_f32(i + i)
168 let neg_two_i: i64 = nx_f32_neg(two_i)
169 let neg_two_i_div_d: i64 = nx_f32_div(neg_two_i, d_f32)
170 let exponent: i64 = nx_f32_mul(neg_two_i_div_d, log_base_f32)
171 let theta: i64 = nx_f32_exp(exponent)
172 let m_theta: i64 = nx_f32_mul(position_f32, theta)
173 let c: i64 = nx_f32_cos(m_theta)
174 let s: i64 = nx_f32_sin(m_theta)
175
176 let x_lo: i64 = x[i]
177 let x_hi: i64 = x[i + half]
178 let new_lo: i64 = nx_f32_sub(nx_f32_mul(x_lo, c), nx_f32_mul(x_hi, s))
179 let new_hi: i64 = nx_f32_add(nx_f32_mul(x_lo, s), nx_f32_mul(x_hi, c))
180 x[i] = new_lo
181 x[i + half] = new_hi
182
183 i = i + 1
184 }
185 return NX_F32_ROPE_OK
186}