Add some flash attn test (#253)

* Add some flash-attn test.

* Add the cpu test.

* Fail when the head is not a multiple of 8.

* Polish the flash attention test.
This commit is contained in:
Laurent Mazare
2023-07-26 20:56:00 +01:00
committed by GitHub
parent ded197497c
commit 4f92420132
5 changed files with 125 additions and 14 deletions

View File

@ -121,8 +121,8 @@ UNARY_OP(float, usin_f32, sing(x))
UNARY_OP(double, usin_f64, sing(x))
UNARY_OP(float, ucos_f32, cosg(x))
UNARY_OP(double, ucos_f64, cosg(x))
UNARY_OP(float, uabsg_f32, absg(x))
UNARY_OP(double, uabsg_f64, absg(x))
UNARY_OP(float, uabs_f32, absg(x))
UNARY_OP(double, uabs_f64, absg(x))
UNARY_OP(float, usqr_f32, x*x)
UNARY_OP(double, usqr_f64, x*x)
UNARY_OP(float, usqrt_f32, sqrtg(x))