|
|
@ -66,9 +66,9 @@ type MLP interface { |
|
|
|
|
|
|
|
|
type sparse struct { |
|
|
type sparse struct { |
|
|
Router *nn.Linear `gguf:"ffn_gate_inp"` |
|
|
Router *nn.Linear `gguf:"ffn_gate_inp"` |
|
|
Gate ml.Tensor `gguf:"ffn_gate_exps.weight"` |
|
|
Gate *nn.Linear `gguf:"ffn_gate_exps"` |
|
|
Up ml.Tensor `gguf:"ffn_up_exps.weight"` |
|
|
Up *nn.Linear `gguf:"ffn_up_exps"` |
|
|
Down ml.Tensor `gguf:"ffn_down_exps.weight"` |
|
|
Down *nn.Linear `gguf:"ffn_down_exps"` |
|
|
} |
|
|
} |
|
|
|
|
|
|
|
|
func (mlp *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor { |
|
|
func (mlp *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor { |
|
|
@ -87,13 +87,13 @@ func (mlp *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options |
|
|
|
|
|
|
|
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1)) |
|
|
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1)) |
|
|
|
|
|
|
|
|
upStates := mlp.Up.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
upStates := mlp.Up.Weight.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
|
|
|
|
|
|
hiddenStates = mlp.Gate.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
hiddenStates = mlp.Gate.Weight.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
hiddenStates = hiddenStates.SILU(ctx) |
|
|
hiddenStates = hiddenStates.SILU(ctx) |
|
|
hiddenStates = hiddenStates.Mul(ctx, upStates) |
|
|
hiddenStates = hiddenStates.Mul(ctx, upStates) |
|
|
|
|
|
|
|
|
experts := mlp.Down.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
experts := mlp.Down.Weight.MulmatID(ctx, hiddenStates, selectedExperts) |
|
|
experts = experts.Mul(ctx, routingWeights) |
|
|
experts = experts.Mul(ctx, routingWeights) |
|
|
|
|
|
|
|
|
nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2)) |
|
|
nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2)) |
|
|
|