nx_gelu.nx source
↩ module page · 201 lines · 6726 B
1// nx_gelu.nx -- Gaussian Error Linear Unit activation.
2//
3// The activation used by every non-Llama transformer family:
4//
5// GPT-2 / GPT-3 / GPT-Neo / GPT-J
6// BERT / RoBERTa / DistilBERT / ALBERT
7// T5 / mT5 / ByT5
8// ViT / DeiT / Swin / CLIP / SigLIP
9// PaLM, Gemma (uses GeGLU variant = GELU * linear gate)
10//
11// Llama / Mistral / Qwen use SiLU instead (closely related;
12// nx_silu.nx already ships). Substrate now ships both -- model
13// loader picks per-config.
14//
15// ===== Math =======================================================
16//
17// Exact GELU (Hendrycks & Gimpel 2016 eq. 1):
18// GELU(x) = x * Phi(x)
19// = x * 0.5 * (1 + erf(x / sqrt(2)))
20//
21// Exact requires erf, which we don't ship. Two practical
22// approximations:
23//
24// 1. Tanh form (most common):
25// GELU(x) ~= 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715*x^3)))
26//
27// 2. Sigmoid form (Hendrycks 2016 alt):
28// GELU(x) ~= x * sigmoid(1.702 * x)
29//
30// We use form #2 because nx_sigmoid_q10 already ships in nx_silu.nx
31// and it's a one-multiply + sigmoid composition. Max relative error
32// vs exact GELU is ~0.5% across [-5, 5] -- well within substrate
33// quantization budget.
34//
35// Tanh-form upgrade queued -- needs nx_tanh.nx primitive. For
36// q4_K-quant transformer inference the sigmoid form is more than
37// sufficient.
38//
39// Q-format: x in Q10, output in Q10. Constant 1.702 in Q10 = 1743.
40//
41// Bits-up composition (pure):
42// nx_silu.nx_sigmoid_q10 (L3, sigmoid, shipped c478b0d1)
43// NxTensor (L1)
44// nx_loop (control)
45//
46// genealogy_id: hendrycks_gimpel_2016_gelu + radford_2019_gpt2_gelu +
47// devlin_2018_bert_gelu
48// lineage_id: substrate_gelu_v1_sigmoid_approx
49
50// nx_safety_envelope:
51// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
52// sil_target: SIL1
53// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
54// verdict: NOT_YET_EVALUATED
55
56import "nx_syscalls.nx"
57import "nx_tier.nx"
58import "nx_loop.nx"
59import "nx_tensor.nx"
60import "nx_silu.nx"
61const NX_MAGIC_1024: i64 = 1024
62const NX_MAGIC_10240: i64 = 10240
63const NX_MAGIC_2048: i64 = 2048
64
65const NX_GELU_Q10: nx_int = 1024
66const NX_GELU_K_1_702_Q10: nx_int = 1743 // 1.702 in Q10 = 1.702 * 1024 = 1742.85
67
68// ===== Sealed-enum: GeluVerdict ===================================
69
70const NX_GELU_OK: nx_int = 0
71const NX_GELU_ERR_BAD_DTYPE: nx_int = 1
72const NX_GELU_ERR_SHAPE_MISMATCH: nx_int = 2
73const NX_GELU_ERR_NOT_CONTIGUOUS: nx_int = 3
74const NX_GELU_N_VERDICTS: nx_int = 4
75
76func nx_gelu_verdict_is_valid(v: nx_int) -> nx_int {
77 if v < 0 { return 0 }
78 if v >= NX_GELU_N_VERDICTS { return 0 }
79 return 1
80}
81
82// ===== Scalar GELU ===============================================
83//
84// gelu(x) = x * sigmoid(1.702 * x)
85//
86// 1.702 * x is in Q10 -> Q20 after multiply; divide by Q10 to
87// bring back to Q10 before passing to sigmoid.
88
89func nx_gelu_q10(x_q10: nx_int) -> nx_int {
90 let arg: nx_int = (x_q10 * NX_GELU_K_1_702_Q10) / NX_GELU_Q10
91 let s: nx_int = nx_sigmoid_q10(arg)
92 return (x_q10 * s) / NX_GELU_Q10
93}
94
95// ===== Tensor forward ============================================
96//
97// Apply GELU elementwise. In-place supported.
98
99func nx_gelu_forward(x: *NxTensor, out: *NxTensor) -> nx_int {
100 if x.dtype != NX_DT_I64 { return NX_GELU_ERR_BAD_DTYPE }
101 if out.dtype != NX_DT_I64 { return NX_GELU_ERR_BAD_DTYPE }
102 if x.numel != out.numel { return NX_GELU_ERR_SHAPE_MISMATCH }
103 if nx_t_is_contiguous(x) == 0 { return NX_GELU_ERR_NOT_CONTIGUOUS }
104 if nx_t_is_contiguous(out) == 0 { return NX_GELU_ERR_NOT_CONTIGUOUS }
105
106 let px: *i64 = x.storage as *i64
107 let po: *i64 = out.storage as *i64
108 let n: nx_int = x.numel
109
110 var i: nx_int = 0
111 var iter: nx_int = 0
112 var verdict: nx_int = NX_LOOP_RUNNING
113 let BUDGET: nx_int = n
114 while verdict == NX_LOOP_RUNNING && iter < BUDGET {
115 po[i] = nx_gelu_q10(px[i])
116 i = i + 1
117 iter = iter + 1
118 }
119 return NX_GELU_OK
120}
121
122// ===== Self-test ==================================================
123//
124// Reference values (float exact GELU):
125// x = 0.0: gelu = 0.0
126// x = 1.0: gelu = 0.8413 (~862 Q10)
127// x = -1.0: gelu = -0.1587 (~-163 Q10)
128// x = 2.0: gelu = 1.9545 (~2002 Q10)
129// x = -2.0: gelu = -0.0455 (~-47 Q10)
130// x = 10.0: gelu ~= x (saturates to linear regime)
131// x = -10.0: gelu ~= 0 (saturates)
132//
133// Sigmoid-form approximation has ~0.5% relative error, so absolute
134// drift in Q10 is ~5-10 for x in [-3, 3] range. Tighter bounds at
135// origin (drift ~ 0), looser bounds far from origin (asymptote
136// matches exact).
137//
138// Closed-form invariants:
139// (a) gelu(0) = 0
140// (b) gelu(1.0) in [840, 880] (842 expected from exact * 1024)
141// (c) gelu(-1.0) in [-170, -150]
142// (d) gelu(large positive) approaches x
143// (e) gelu(large negative) approaches 0
144// (f) Tensor forward applies elementwise
145// (g) Verdict gate
146
147func main() -> i64 {
148 if nx_gelu_q10(0) != 0 { return 10 }
149
150 let g1: nx_int = nx_gelu_q10(NX_MAGIC_1024)
151 if g1 < 840 { return 20 }
152 if g1 > 880 { return 21 }
153
154 let gn1: nx_int = nx_gelu_q10(-NX_MAGIC_1024)
155 if gn1 < -170 { return 30 }
156 if gn1 > -150 { return 31 }
157
158 // Large positive: gelu(x) -> x. At x = 10 Q10 = 10240:
159 let g_big: nx_int = nx_gelu_q10(NX_MAGIC_10240)
160 let drift_big: nx_int = g_big - NX_MAGIC_10240
161 if drift_big > 0 { return 40 } // gelu(x) <= x
162 if drift_big < -100 { return 41 } // close to x
163
164 // Large negative: gelu(x) -> 0.
165 let g_nbig: nx_int = nx_gelu_q10(-NX_MAGIC_10240)
166 if g_nbig > 0 { return 50 }
167 if g_nbig < -200 { return 51 }
168
169 // Tensor forward.
170 let sh: *nx_int = sys_mmap(8) as *nx_int
171 sh[0] = 5
172 let err: *nx_int = sys_mmap(8) as *nx_int
173 err[0] = 0
174 let xt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 1, err)
175 let yt: *NxTensor = nx_t_alloc(NX_DT_I64, sh, 1, err)
176 if err[0] != 0 { return 60 }
177
178 let px: *i64 = xt.storage as *i64
179 px[0] = -NX_MAGIC_2048
180 px[1] = -NX_MAGIC_1024
181 px[2] = 0
182 px[3] = NX_MAGIC_1024
183 px[4] = NX_MAGIC_2048
184
185 let v: nx_int = nx_gelu_forward(xt, yt)
186 if v != NX_GELU_OK { return 70 + v }
187
188 let py: *i64 = yt.storage as *i64
189 if py[2] != 0 { return 80 } // gelu(0) = 0
190 if py[0] >= 0 { return 81 } // gelu(-2) < 0
191 if py[4] <= 0 { return 82 } // gelu(+2) > 0
192
193 // --- Verdict gate ---
194 var vi: nx_int = 0
195 while vi < NX_GELU_N_VERDICTS {
196 if nx_gelu_verdict_is_valid(vi) != 1 { return 90 + vi }
197 vi = vi + 1
198 }
199
200 return 0
201}