openai swiglu

ml/backend.go

@@ -283,6 +283,7 @@ type Tensor interface {
 	SILU(ctx Context) Tensor
 	RELU(ctx Context) Tensor
 	Sigmoid(ctx Context) Tensor
+	SwiGLU(ctx Context, up Tensor, alpha, limit float32) Tensor
 
 	Reshape(ctx Context, shape ...int) Tensor
 	View(ctx Context, offset int, shape ...int) Tensor

ml/backend/ggml/ggml.go

@@ -1409,6 +1409,13 @@ func (t *Tensor) RELU(ctx ml.Context) ml.Tensor {
 	}
 }
 
+func (t *Tensor) SwiGLU(ctx ml.Context, up ml.Tensor, alpha, limit float32) ml.Tensor {
+	return &Tensor{
+		b: t.b,
+		t: C.ggml_swiglu_oai(ctx.(*Context).ctx, t.t, up.(*Tensor).t, C.float(alpha), C.float(limit)),
+	}
+}
+
 func (t *Tensor) Conv2D(ctx ml.Context, t2 ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
 	return &Tensor{
 		b: t.b,

model/models/gptoss/model.go

@@ -190,19 +190,15 @@ func (mlp *MLPBlock) Forward(ctx ml.Context, hiddenStates, one ml.Tensor, opts *
 
 	dimStride := []int{hiddenStates.Dim(0) / 2, hiddenStates.Stride(1), hiddenStates.Dim(1), hiddenStates.Stride(2), hiddenStates.Dim(2), hiddenStates.Stride(3), hiddenStates.Dim(3)}
 
-	glu := hiddenStates.View(ctx, 0, dimStride...)
-	glu = glu.Contiguous(ctx)
-	glu = glu.Reshape(ctx, glu.Dim(0)*glu.Dim(1), glu.Dim(2), glu.Dim(3))
+	gate := hiddenStates.View(ctx, 0, dimStride...)
+	gate = gate.Contiguous(ctx)
+	gate = gate.Reshape(ctx, gate.Dim(0)*gate.Dim(1), gate.Dim(2), gate.Dim(3))
 
-	glu = glu.Clamp(ctx, float32(math.Inf(-1)), 7.0)
-	glu = glu.QuickGELU(ctx)
+	up := hiddenStates.View(ctx, hiddenStates.Stride(0), dimStride...)
+	up = up.Contiguous(ctx)
+	up = up.Reshape(ctx, up.Dim(0)*up.Dim(1), up.Dim(2), up.Dim(3))
 
-	linear := hiddenStates.View(ctx, hiddenStates.Stride(0), dimStride...)
-	linear = linear.Contiguous(ctx)
-	linear = linear.Reshape(ctx, linear.Dim(0)*linear.Dim(1), linear.Dim(2), linear.Dim(3))
-	linear = linear.Clamp(ctx, -7.0, 7.0)
-
-	hiddenStates = glu.Mul(ctx, linear.Add(ctx, one))
+	hiddenStates = gate.SwiGLU(ctx, up, 1.702, 7)
 
 	experts := mlp.Down.Forward(ctx, hiddenStates, selectedExperts)
 	experts = experts.Mul(ctx, routingWeights)