update

Signed-off-by: daquexian <daquexian566@gmail.com>
daquexian · Oct 20, 2023 · 6f65056 · 6f65056
1 parent be19206
commit 6f65056
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/kernels/cuda/ffn.cu b/kernels/cuda/ffn.cu
@@ -81,6 +81,8 @@ Tensor _FFN(const Tensor &x, const Tensor &sx, const Tensor &ln_w,
   element_wise(InplaceReLUAndSquare{vx}, kw.size(1));
   gemm_cublas(vx, vw.data_ptr<half>(), x_plus_out.data_ptr<half>(), 1, 1,
               vw.size(1), vw.size(0));
+  Tensor r_t = Tensor::FromPtr(r, x.sizes(), DType::kFloat16, x.device());
+  // hfma loses precision
   x_plus_out = x_plus_out * r_t + x;
   return xx;
 }