Fast Custom Ops

mlx_array mlx_fast_layer_norm(mlx_array x, mlx_array weight, mlx_array bias, float eps, mlx_stream s)

mlx_array mlx_fast_rms_norm(mlx_array x, mlx_array weight, float eps, mlx_stream s)

mlx_array mlx_fast_rope(mlx_array x, int dims, bool traditional, float base, float scale, int offset, mlx_stream s)

mlx_array mlx_fast_scaled_dot_product_attention(mlx_array queries, mlx_array keys, mlx_array values, float scale, mlx_array mask, mlx_stream s)