Skip to content

Commit ea0561b

Browse files
gbenrosciencegbenroscience
committed
tightening gelu loop
1 parent bd8c0c8 commit ea0561b

4 files changed

Lines changed: 138 additions & 25 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,8 @@ Add the following to your `pom.xml`:
148148
* **Javadoc:** [Latest API Documentation](https://www.google.com/search?q=https://javadoc.io/doc/com.github.gbenroscience/parser-ng/2.0.5)
149149

150150

151+
[jenkins-report-generic-chart-column](https://github.com/jenkinsci/report-generic-chart-column-plugin/) by Jiri Vanek, makes use of `ExpandingParser`, and his contributions are appreciated.
152+
151153

152154
## 📚 More Documentation
153155

‎parser-ng-simd/src/main/java/com/github/gbenroscience/simd/turbo/tools/FlatMatrixF.java‎

Lines changed: 133 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -252,31 +252,90 @@ private static void scalarMicrokernel(FlatMatrixF A, FlatMatrixF B, FlatMatrixF
252252
}
253253

254254
// ====================== ELEMENTWISE ======================
255-
public void geluInPlace() {
256-
int len = rows * cols, idx = 0;
255+
/**
256+
* Slower than Bert, but MORE accurate GELU in-place using the Sigmoid-based
257+
* exact identity math. Matches standard Tanh approximation curves
258+
* perfectly.
259+
*/
260+
public void geluInPlaceSigmoid() {
261+
int len = rows * cols;
262+
int idx = 0;
263+
257264
if (HAS_VECTOR && isContiguous()) {
265+
// 2 * 0.79788456f = 1.59576912f
266+
final FloatVector V_C1 = FloatVector.broadcast(F_SPECIES, 1.59576912f);
267+
final FloatVector V_C2 = FloatVector.broadcast(F_SPECIES, 0.044715f);
268+
final FloatVector V_ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
269+
258270
for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
259-
var v = FloatVector.fromArray(F_SPECIES, data, offset + idx);
260-
var x2 = v.mul(v);
261-
var inner = v.add(x2.mul(v).mul(V_GELU_C2)).mul(V_GELU_C1);
262-
var t = inner.lanewise(VectorOperators.TANH);
263-
v.mul(V_HALF).mul(t.add(V_ONE)).intoArray(data, offset + idx);
271+
var x = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
272+
273+
var x3 = x.mul(x).mul(x);
274+
var z = x.add(x3.mul(V_C2)).mul(V_C1);
275+
276+
// Compute the denominator: 1 + e^(-z)
277+
var denominator = V_ONE.add(fastExp(z.neg()));
278+
279+
// Divide x directly by the denominator to get x * sigmoid(z)
280+
var result = x.div(denominator);
281+
282+
result.intoArray(this.data, this.offset + idx);
264283
}
265284
}
285+
286+
// Scalar fallback matching the exact Tanh/Sigmoid curve behavior
266287
for (; idx < len; idx++) {
267-
float x = data[offset + idx];
288+
float x = this.data[this.offset + idx];
268289
float x3 = x * x * x;
269-
float t = (float) Math.tanh(GELU_C1 * (x + GELU_C2 * x3));
270-
data[offset + idx] = 0.5f * x * (1.0f + t);
290+
float t = (float) Math.tanh(0.79788456f * (x + 0.044715f * x3));
291+
this.data[this.offset + idx] = 0.5f * x * (1.0f + t);
271292
}
272293
}
273294

274295
/**
275-
* GeGLU in-place: output = inputs[0] ⊙ GELU(inputs[1])
296+
* Faster than sigmoid, but less accurate GELU in-place using the fast
297+
* BERT/OpenAI 1.702 * x approximation. Eliminates cubic calculations
298+
* entirely for maximum throughput.
299+
*/
300+
public void geluInPlaceBert() {
301+
int len = rows * cols;
302+
int idx = 0;
303+
304+
if (HAS_VECTOR && isContiguous()) {
305+
final FloatVector V_BERT_C = FloatVector.broadcast(F_SPECIES, 1.702f);
306+
final FloatVector V_ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
307+
308+
for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
309+
var x = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
310+
311+
// z = 1.702 * x
312+
var z = x.mul(V_BERT_C);
313+
314+
// sigmoid(z) = 1 / (1 + fastExp(-z))
315+
var denominator = V_ONE.add(fastExp(z.neg()));
316+
317+
// x * sigmoid(z)
318+
var result = x.div(denominator);
319+
result.intoArray(this.data, this.offset + idx);
320+
}
321+
}
322+
323+
// Scalar fallback matching the exact 1.702 sigmoid calculation
324+
for (; idx < len; idx++) {
325+
float x = this.data[this.offset + idx];
326+
float sigmoid = 1.0f / (1.0f + (float) Math.exp(-1.702f * x));
327+
this.data[this.offset + idx] = x * sigmoid;
328+
}
329+
}
330+
331+
/**
332+
* Slower than Bert, but MORE accurate GeGLU in-place using the
333+
* Sigmoid-based exact identity math. Computes: output = value *
334+
* GELU_sigmoid(gate)
276335
*
277-
* @param gate
336+
* @param gate The gating matrix (inputs[1])
278337
*/
279-
public void gegluInPlaceOld(FlatMatrixF gate) { // gate is inputs[1]
338+
public void gegluInPlaceSigmoid(FlatMatrixF gate) {
280339
if (rows != gate.rows || cols != gate.cols) {
281340
throw new IllegalArgumentException("Shape mismatch for GeGLU");
282341
}
@@ -285,35 +344,88 @@ public void gegluInPlaceOld(FlatMatrixF gate) { // gate is inputs[1]
285344
int idx = 0;
286345

287346
if (HAS_VECTOR && isContiguous() && gate.isContiguous()) {
347+
// 2 * 0.79788456f = 1.59576912f
348+
final FloatVector V_C1 = FloatVector.broadcast(F_SPECIES, 1.59576912f);
349+
final FloatVector V_C2 = FloatVector.broadcast(F_SPECIES, 0.044715f);
350+
final FloatVector V_ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
351+
288352
for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
289353
var value = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
290354
var g = FloatVector.fromArray(F_SPECIES, gate.data, gate.offset + idx);
291355

292-
// GELU on gate
293-
var g2 = g.mul(g);
294-
var g3 = g2.mul(g);
295-
var inner = g.add(g3.mul(V_GELU_C2)).mul(V_GELU_C1);
296-
var t = inner.lanewise(VectorOperators.TANH);
297-
var geluGate = g.mul(V_HALF).mul(t.add(V_ONE));
356+
// z = 1.59576912 * (g + 0.044715 * g^3)
357+
var g3 = g.mul(g).mul(g);
358+
var z = g.add(g3.mul(V_C2)).mul(V_C1);
359+
360+
// GELU(gate) = g / (1 + fastExp(-z))
361+
var denominator = V_ONE.add(fastExp(z.neg()));
362+
var geluGate = g.div(denominator);
298363

299364
// value * GELU(gate)
300365
value.mul(geluGate).intoArray(this.data, this.offset + idx);
301366
}
302367
}
303368

304-
// Scalar fallback
369+
// Scalar fallback matching the exact Tanh/Sigmoid curve behavior
305370
for (; idx < len; idx++) {
306371
float value = this.data[this.offset + idx];
307372
float g = gate.data[gate.offset + idx];
308373

309374
float g3 = g * g * g;
310-
float t = (float) Math.tanh(GELU_C1 * (g + GELU_C2 * g3));
375+
float t = (float) Math.tanh(0.79788456f * (g + 0.044715f * g3));
311376
float geluGate = 0.5f * g * (1.0f + t);
312377

313378
this.data[this.offset + idx] = value * geluGate;
314379
}
315380
}
316381

382+
/**
383+
* Faster than sigmoid, but less accurate GeGLU in-place using the fast
384+
* BERT/OpenAI 1.702 * g approximation. Computes: output = value *
385+
* GELU_bert(gate)
386+
*
387+
* @param gate The gating matrix (inputs[1])
388+
*/
389+
public void gegluInPlaceBert(FlatMatrixF gate) {
390+
if (rows != gate.rows || cols != gate.cols) {
391+
throw new IllegalArgumentException("Shape mismatch for GeGLU");
392+
}
393+
394+
int len = rows * cols;
395+
int idx = 0;
396+
397+
if (HAS_VECTOR && isContiguous() && gate.isContiguous()) {
398+
final FloatVector V_BERT_C = FloatVector.broadcast(F_SPECIES, 1.702f);
399+
final FloatVector V_ONE = FloatVector.broadcast(F_SPECIES, 1.0f);
400+
401+
for (; idx < F_SPECIES.loopBound(len); idx += VF_LEN) {
402+
var value = FloatVector.fromArray(F_SPECIES, this.data, this.offset + idx);
403+
var g = FloatVector.fromArray(F_SPECIES, gate.data, gate.offset + idx);
404+
405+
// z = 1.702 * g
406+
var z = g.mul(V_BERT_C);
407+
408+
// GELU(gate) = g / (1 + fastExp(-z))
409+
var denominator = V_ONE.add(fastExp(z.neg()));
410+
var geluGate = g.div(denominator);
411+
412+
// value * GELU(gate)
413+
value.mul(geluGate).intoArray(this.data, this.offset + idx);
414+
}
415+
}
416+
417+
// Scalar fallback matching the exact 1.702 sigmoid calculation
418+
for (; idx < len; idx++) {
419+
float value = this.data[this.offset + idx];
420+
float g = gate.data[gate.offset + idx];
421+
422+
float sigmoid = 1.0f / (1.0f + (float) Math.exp(-1.702f * g));
423+
float geluGate = g * sigmoid;
424+
425+
this.data[this.offset + idx] = value * geluGate;
426+
}
427+
}
428+
317429
public void gegluInPlace(FlatMatrixF gate) {
318430
if (rows != gate.rows || cols != gate.cols) {
319431
throw new IllegalArgumentException("Shape mismatch for GeGLU");

‎parser-ng-simd/src/main/java/com/github/gbenroscience/simd/turbo/tools/VectorTurboEvaluator.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1785,8 +1785,7 @@ public void applyMatrixKernel(FlatMatrixF[] inputs, FlatMatrixF output, String o
17851785
if (inputs[0] != output) {
17861786
System.arraycopy(inputs[0].data, inputs[0].offset, output.data, output.offset, output.rows * output.cols);
17871787
}
1788-
output.geluInPlace();
1789-
//output.geluInPlaceHighSpeed();
1788+
output.geluInPlaceSigmoid();
17901789
}
17911790
case "geglu", "geglu_in_place" -> {
17921791
if (inputs.length < 2) {
@@ -1800,7 +1799,7 @@ public void applyMatrixKernel(FlatMatrixF[] inputs, FlatMatrixF output, String o
18001799
output.rows * output.cols);
18011800
}
18021801

1803-
output.gegluInPlace(inputs[1]); // <- gate
1802+
output.gegluInPlaceSigmoid(inputs[1]); // <- gate
18041803
}
18051804
// === New Q8 + Attention kernels ===
18061805
case "q8_quantize" -> {

‎parser-ng/src/main/java/com/github/gbenroscience/README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -901,7 +901,7 @@ is: true
901901
true
902902
```
903903

904-
See [jenkins-report-generic-chart-column](https://github.com/jenkinsci/report-generic-chart-column-plugin/ as real-world user of `ExpandingParser`
904+
See [jenkins-report-generic-chart-column](https://github.com/jenkinsci/report-generic-chart-column-plugin/) as real-world user of `ExpandingParser`
905905

906906

907907

0 commit comments

Comments
 (0)