[CUDA][Improvement] Fast path for gather

[nastya236]

Fast path for the case when gathered axis = 0, input is contiguous and slice size is a full stride. It is useful for gather before gather_mm and scatter after (for MoEs).

Some numbers on B200:

D = 2048 
B = 2
L = 8192
top_k = 8
T = B*L*top_k

input shape is (B*L, D), output shape is (T, D)

Before:
1.5219 ms, 705.5 GB/s

After:
0.2272 ms 4725.9 GB/s

Torch:
0.1864 ms 5759.1 GB/s

I think the remaining ~0.05 ms is due to kernel launch overhead rather than kernel execution itself. That said, if you spot anything that could be optimized further to reduce it, I'm happy to address it.

[zcbenz]

Have you tried swapping grid.x and grid.y? It seems that currently the blocks would read in this sequence:

block0: src[0, 0:1024]
block1: src[1, 0:1024]
block2: src[2, 0:1024]
...
block8: src[0, 1024:2048]
block9: src[1, 1024:2048]

we might get better coalesce by swapping:

block0: src[0, 0:1024]
block1: src[0, 1024:2048]
block2: src[0, 2048:4096]
...
block8: src[1, 0:1024]
block9: src[1, 1024:2048]

For benchmarking you can use env var MLX_USE_CUDA_GRAPHS=0 to remove the overhead of cuda graph, there would still be overhead of lazy evaluation though.

[zcbenz]

Some nitpickings:

Suggested change

	uint32_t n_vec = (slice_size + vec_size - 1) / vec_size;
	uint32_t max_block_dim = 256;
	uint32_t block_dim = n_vec < max_block_dim ? n_vec : max_block_dim;
	dim3 block(block_dim, 1, 1);
	dim3 grid(
	uint32_t n_vec = cuda::ceil_div(slice_size, vec_size);
	uint32_t max_block_dim = 256;
	uint32_t block_dim = std::min(n_vec, max_block_dim);
	dim3 block_dims(block_dim, 1, 1);
	dim3 num_blocks(

We don't have standards on what names to call block and grid dims but we should use the same names in the same file. (We did not use the more common grid_dims as name because it means a different thing in Metal unfortunately.)

[zcbenz]

It seems that the kernel is capable of handling the large size?