[DeepSeek][Kernels] MoE sorting - Scatter Gather kernels

[lessw2020]

encapsulating scatter/gather in Cuda kernels. Effectively does the following PyTorch actions but via CUDA:

def pytorch_sort_tokens(topk_ids, x, n_experts):
    """PyTorch implementation for comparison"""
    with torch.no_grad():
        # [seq_len, n_experts]
        cnts = topk_ids.new_zeros((topk_ids.shape[0], n_experts))
        # Fill 1 to the selected experts
        cnts.scatter_(1, topk_ids, 1)
        tokens_per_expert = cnts.sum(dim=0)
        # Token indices for each expert
        idxs = topk_ids.view(-1).argsort()
        sorted_tokens_shape = idxs.shape + x.shape[1:]
    sorted_tokens = x[idxs // topk_ids.shape[1]]

    return sorted_tokens, idxs, tokens_per_expert