[0z5a]

Summary

Add a bufferless PyTorch backend for MoeDispatch and MoeCombine when the expert-parallel group has one rank. Dispatch orders local routes by expert and combine sums expert outputs back into source tokens. Routing metadata belongs to each invocation, so two outstanding forwards do not overwrite each other. The local backend rejects remote expert IDs and communication settings it cannot honor.

This is the EP=1 follow-up requested on NVIDIA/TransformerEngine#3503. This draft is stacked on vthumbe1503:nccl_ep_ops at 261a81c; its signed-off implementation commit and pre-commit.ci formatting commit change only the EP=1 path, tests, and API documentation. The main-targeting upstream draft is open.

Validation

Check	Result
CUDA_VISIBLE_DEVICES=3 PYTHONPATH=$PWD python -m pytest tests/pytorch/test_moe_ep_local.py -q	13 passed in 11.57 s on RTX 5090, PyTorch 2.13.0+cu130, at head 6ce761b
Five-op MoeDispatch → GroupedLinear → ScaledSwiGLU → GroupedLinear → MoeCombine	BF16 forward and backward match an independent dense PyTorch MoE within the test's rtol=atol=2e-2, including token, gate-weight, and expert-weight gradients
Twelve-step training E2E with changing inputs and top-k=2 routes	Full five-op forward, backward, SGD updates, and post-update inference matched the dense PyTorch reference; 8 parameter tensors changed, maximum output absolute difference 0.00178438 and maximum gradient absolute difference 0.00277808
Routing cases	top-k=1 and 2, repeated and dropped routes, empty experts, two outstanding forwards, and unsupported settings covered
Repository license checker	Passed

The test imported the PR source tree, ran without an initialized torch.distributed process group, and the loaded native extension exposed no EP entry points. The NCCL EP distributed suite was not run on this machine. The full L0 launcher was not run because it installs dependencies; the focused tests used the existing environment without package changes.