[CUDA] Grouped mm

[nastya236]

Currently for grouped_mm operation (that is used in MoEs for example) we use the gather_mm with sorted indices. On cuda it is cutlass unaligned kernel for sm80 (Ampere). Also, since gather_mm expects indices as input (not token counts), in case of MoEs swiglu we need to recompute tokens offset in all projections.

This PR adds:

• grouped_mm op that inputs sorted activations, token offsets and experts weights
• port of moe_grouped_matmul
• GroupedMM primitive [for now only forward]

D, H, E, S, k = 2048, 1024, 32, 8192, 4
dtype = mx.bfloat16

def gather_sort(x, indices):
    *_, top_k = indices.shape
    indices = indices.flatten()
    order = mx.argsort(indices)
    return x[order // top_k], indices[order], mx.argsort(order)

def make_offsets(sorted_indices, num_experts):
    experts = mx.arange(num_experts, dtype=sorted_indices.dtype)
    offsets = mx.searchsorted(sorted_indices, experts)
    return offsets.astype(mx.int32).reshape(num_experts, 1, 1)

def mlp_gather(x, wg, wu, wd, indices):
    up = mx.gather_mm(x, wu, rhs_indices=indices, sorted_indices=True)
    gate = mx.gather_mm(x, wg, rhs_indices=indices, sorted_indices=True)
    return  mx.gather_mm(nn.silu(gate) * up, wd, indices)

def mlp_grouped(x, wg, wu, wd, indices):
    offsets = make_offsets(indices, wg.shape[0])
    up = mx.grouped_mm(x, wu, token_offsets=offsets)
    gate = mx.grouped_mm(x, wg, token_offsets=offsets)
    return mx.grouped_mm(nn.silu(gate) * up, wd, token_offsets=offsets)

Before: gather_mm=47.922ms
After: grouped_mm=0.552ms

I hope that the numbers are correct, I recomputed it multiple times. 80x difference looks weird tho..
TODO:

Currently tests for grouped_mm are skipped because sm 7.5 is not supported by cudnn.
I decided to split the change in two pull requests: this one implements cudnn port, primitive and operation.
The second one will implement a fallback for sm < 80 using cutlass_grouped_mm.

• fallback for cuda by using cutlass_grouped_mm without offset calculation
• forward for metal
• backward for cuda
• backward for metal
• grouped_qmm for metal
• grouped_qqmm for cuda

After this we can train MoEs with mlx 🎉

[zcbenz]

We don't need to check CUDNN_VERSION ourselves, the cudnn-frontend C++ APIs we use are capable of detecting cudnn version and throw errors. And we can ensure minimum cudnn version in setup.py by setting the version of nvidia-cudnn-cu12/13 dependencies.

Also since cudnn_grouped_mm requires sm80 and later, this function should check it here.