pytorch - 💡(How to fix) Fix [AMD] MoRI support for MoE EP Combine/Dispatch torch.distributed.TokenSwitch Interface & AMD Pollara Support [1 pull requests]

🚀 The feature, motivation and pitch

Hi @powderluv @chunfangamd @andyluo7 @pragupta @simondanielsson

Following up on #178712 by @kwen2501, there should be MoRI support to the new TokenSwitch abstraction

TokenSwitch is an abstract class defining MoE-style token exchange via three interfaces: • bind_routing • dispatch • combine

The initial implementation (TokenSwitchNCCL) is backed by nccl-ep. To make this first class on AMD/ROCm hardware, we should add a MoRI-backed implementation.

In the age of MoE, for inference & training, MoE combine & MoE dispatch is one of the most important comms op and takes any where from 60-70% of the communication volume. Having the ROCm optimized MoRI is essential to optimized first class pytorch

Details

• Add a TokenSwitchMoRI implementation conforming to the TokenSwitch interface (bind_routing / dispatch / combine). • Wire it up so users can select the backend (NCCL vs MoRI) at construction time, e.g. TokenSwitch(group, num_experts, num_tokens, backend="mori"). • MoRI PyTorch should work out of the box — the Dockerfile and prebuilt ROCm PyTorch images at https://hub.docker.com/r/rocm/pytorch should ship with MoRI included and TokenSwitchMoRI usable with no extra install steps. Please make sure the published Dockerfile and image tags on https://hub.docker.com/r/rocm/pytorch bundle MoRI + the dependencies needed for TokenSwitch. • AMD first-party Pollara / AINIC NIC support out of the box — just like CUDA images bundle ConnectX NIC userspace bits, the ROCm PyTorch image should ship with the AMD Pollara (AINIC ) userspace RoCE libraries pre-installed, mirroring the approach in https://github.com/vllm-project/vllm/pull/40453

Alternatives

No response

Additional context

No response

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @aditvenk @weifengpy