code wiki / (root) / nx_gelu.nx

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}