[a2av] Add autograd support for token combine op

[kwen2501]

Added class TokenCombiner which combines tokens from different experts, with backward support.

Usage:

combiner = TokenCombiner(group_name, align, max_inp_len, max_out_len, inp.shape[1:], world_size, ne, dtype)
# inp, out, in_splits_offsets, out_splits_offsets must be symmetric tensors
output = combiner(inp, out, in_splits_offsets, out_splits_offsets)

Supports:

torch.compile(combiner)

[tianyu-l]

similar request: please allow device as an input to constructor

[vwxyzjn]

Hi @kwen2501 is this kernels for all to all communications? Any ways to test it out with deepseek v3?

[kwen2501]

@vwxyzjn Yes, the inner part of the PR is an op named torch.ops.symm_mem.all_to_all_vdev_2d_offset, which accepts token splits + offsets on GPU. The PR just wraps that op with Autograd support.

Similarly, for token dispatch, we have an op named all_to_all_vdev_2d. There is more documentation about these two ops here:
https://github.com/pytorch/pytorch/blob/main/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu#L557
https://github.com/pytorch/pytorch/blob/main/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.cu#L704