Repository: ollama/ollama
File: model/models/llama4/model_text.go (Go)

Code before:
	return mlp.Down.Forward(ctx, hiddenStates)
}

type TextExperts struct {
	Gate ml.Tensor `gguf:"ffn_gate_exps.weight"`
	Up   ml.Tensor `gguf:"ffn_up_exps.weight"`
	Down ml.Tensor `gguf:"ffn_down_exps.weight"`
}

func (e *TextExperts) Forward(ctx ml.Context, hiddenStates, routerLogits ml.Tensor, opts *TextOptions) ml.Tensor {
	experts := routerLogits.TopK(ctx, opts.numExpertsUsed)
	scores := routerLogits.Sigmoid(ctx).Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, experts)

	hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
	hiddenStates = hiddenStates.Repeat(ctx, 1, opts.numExpertsUsed)
	hiddenStates = hiddenStates.Mul(ctx, scores)

	upStates := e.Up.MulmatID(ctx, hiddenStates, experts)
	gateStates := e.Gate.MulmatID(ctx, hiddenStates, experts)
	downStates := e.Down.MulmatID(ctx, upStates.Mul(ctx, gateStates.SILU(ctx)), experts)

	nextStates := downStates.View(ctx, 0, hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2))
	for i := 1; i < opts.numExpertsUsed; i++ {
		nextStates = nextStates.Add(ctx, downStates.View(ctx, i*downStates.Stride(1), hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2)))
	}

	return nextStates
}

// TextSharedExpert is TextMLP with different tensor names
type TextSharedExpert struct {
	Gate *nn.Linear `gguf:"ffn_gate_shexp"`
	Up   *nn.Linear `gguf:"ffn_up_shexp"`
	Down *nn.Linear `gguf:"ffn_down_shexp"`
}

func (mlp *TextSharedExpert) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenStates))
	return mlp.Down.Forward(ctx, hiddenStates)
}

type TextMOE struct {
	Router       *nn.Linear `gguf:"ffn_gate_inp"`
	Experts      *TextExperts
	SharedExpert *TextSharedExpert
}

func (moe *TextMOE) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
	hiddenDim, sequenceLength, batchSize := hiddenStates.Dim(0), hiddenStates.Dim(1), hiddenStates.Dim(2)
	hiddenStates = hiddenStates.Reshape(ctx, hiddenDim, sequenceLength*batchSize)
	routerLogits := moe.Router.Forward(ctx, hiddenStates)

Review comment (bug):
This fixes a bug where additional experts were being skipped. There's no impact to Llama 4 which has a single expert

Code after:
func (mlp *TextMLP) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenStates))
	return mlp.Down.Forward(ctx, hiddenStates)
}

type TextExperts struct {
	Gate ml.Tensor `gguf:"ffn_gate_exps.weight"`
	Up   ml.Tensor `gguf:"ffn_up_exps.weight"`
	Down ml.Tensor `gguf:"ffn_down_exps.weight"`
}

func (e *TextExperts) Forward(ctx ml.Context, hiddenStates, routerLogits ml.Tensor, opts *TextOptions) ml.Tensor {
	experts := routerLogits.TopK(ctx, opts.numExpertsUsed)
	scores := routerLogits.Sigmoid(ctx).Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, experts)

	hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
	hiddenStates = hiddenStates.Repeat(ctx, 1, opts.numExpertsUsed)
	hiddenStates = hiddenStates.Mul(ctx, scores)

	upStates := e.Up.MulmatID(ctx, hiddenStates, experts)
	gateStates := e.Gate.MulmatID(ctx, hiddenStates, experts)
	downStates := e.Down.MulmatID(ctx, upStates.Mul(ctx, gateStates.SILU(ctx)), experts)

	nextStates := downStates.View(ctx, 0, hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2))
	for i := 1; i < opts.numExpertsUsed; i++ {
		nextStates = nextStates.Add(ctx, downStates.View(ctx, i*downStates.Stride(1), hiddenStates.Dim(0), downStates.Stride(2), hiddenStates.Dim(2)))
	}

	return nextStates
}

// TextSharedExpert is TextMLP with different tensor names
type TextSharedExpert struct {
	Gate *nn.Linear `gguf:"ffn_gate_shexp"`
	Up   *nn.Linear `gguf:"ffn_up_shexp"`
	Down *nn.Linear `gguf:"ffn_down_shexp"`
}

func (mlp *TextSharedExpert) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
	hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx).Mul(ctx, mlp.Up.Forward(ctx, hiddenStates))
	return mlp.Down.Forward(ctx, hiddenStates)
}

type TextMOE struct {
	Router       *nn.Linear `gguf:"ffn_gate_inp"`
	Experts      *TextExperts
	SharedExpert *TextSharedExpert
}

func (moe *TextMOE) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *TextOptions) ml.Tensor {
	hiddenDim, sequenceLength, batchSize := hiddenStates.Dim(0), hiddenStates.Dim(1), hiddenStates.Dim(2)