RoPE for CUDA (#2293)

* First working CUDA rope * Fix random
2025-12-16 01:49:05 +08:00 · 2025-06-15 06:08:07 -07:00
parent a14aaa7c9d
commit 580776559b
8 changed files with 443 additions and 29 deletions
--- a/mlx/backend/cuda/kernel_utils.cuh
+++ b/mlx/backend/cuda/kernel_utils.cuh
@@ -121,6 +121,7 @@ dim3 get_2d_grid_dims(
    const Shape& shape,
    const Strides& strides,
    size_t divisor);
+std::pair<dim3, dim3> get_grid_and_block(int dim0, int dim1, int dim2);

 // Return a block size that achieves maximum potential occupancy for kernel.
 template <typename T>