Function ffn_cpu

src/infer.cpp:1166–1197

}

void ffn_cpu(
  float* xout, float* x, 
  float* w1, float* w2, float* w3, 
  int hidden_dim, int dim,
  ActivationType act
) {
  float* hb = new float[hidden_dim];
  float* hb2 = new float[hidden_dim];
  // mix self.w2(F.silu(self.w1(x)) * self.w3(x))
  // Note this is a feedforward with a GLU, not a simple MLP.
  matmul_unscaled(hb, x, QTensor(Quant::F32, {dim, hidden_dim}, w1, dim*hidden_dim*sizeof(float)));
  matmul_unscaled(hb2, x, QTensor(Quant::F32, {dim, hidden_dim}, w3, dim*hidden_dim*sizeof(float)));
  switch (act) {
    case ActivationType::GELU: {
      for (int i = 0; i < hidden_dim; ++i) {
        hb[i] = gelu(hb[i]) * hb2[i];
      }
      break;
    }
    case ActivationType::SILU: {
      for (int i = 0; i < hidden_dim; ++i) {
        hb[i] = silu(hb[i]) * hb2[i];
      }
      break;
    }
  }

  matmul_unscaled(xout, hb, QTensor(Quant::F32, {hidden_dim, dim}, w2, hidden_dim*dim*sizeof(float)));
  
  delete[] hb;
  delete[] hb2;
}

template void Block::_block_cpu<float>(InferenceState&, int, int, int, int) const;
template void Block::_block_cpu<f16_t>(InferenceState&, int, int, int, int) const;