[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]

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