3 Commits

Author	SHA1	Message	Date
Angelos Katharopoulos	2225374060	Adds mx.fast.layer_norm (#870)	2024-03-21 13:55:51 -07:00
Awni Hannun	a54f06b16f	Fast RMS Norm (#862) ... * fast rmsnorm * no rms gpu * kernel * fix shared mem * looped rms and donation in softmax * Make the squaring in float32 to avoid underflow * Fix the default StreamOrDevice for rope and rms_norm in fast * nits --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2024-03-21 07:20:54 -07:00
Awni Hannun	ccf1645995	Custom primitive + RoPE fat op (#676) ... * extensions start * rope custom op * fix build * docs + rope benchmark * fix test * Add a Metal kernel for RoPE * Fix position of traditional * transform tests * Move rope computation to float and fix tests * Fix the test and a typo * change to fast * fix no metal build --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2024-02-14 14:04:25 -08:00