[PyTorch] Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy#3035

Open

phu0ngng wants to merge 39 commits into

NVIDIA:mainfrom

phu0ngng:phuong/ep-3-pytorch-on-commwindow

Commits on Jun 23, 2026

EP PyTorch: NCCL EP backend + autograd ops + tests, route zero_copy via cfg.zero_copy
phu0ngng
committed
EP PyTorch: wire maybe_make_window into per-step ops for zero_copy
phu0ngng
committed
EP PyTorch: merge EpHandle into EpBuffer; ep_dispatch/ep_combine take a single buffer
phu0ngng
committed
EP PyTorch example: drop stale ep_group kwarg from EpBuffer call
phu0ngng
committed
EP PyTorch: drop ep_scope; ep_finalize is optional with atexit fallback
phu0ngng
committed
EP PyTorch: restrict payload dtype to bf16; refresh stale ep.cpp comments; drop unused ep_group kwargs
phu0ngng
committed
EP PyTorch: clear pylint warnings in ep.py (broad-except suppression, _-prefixed stub args, autograd docstrings)
phu0ngng
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
phu0ngng
committed
EP PyTorch: skip combine_in staging copy in non-zero-copy mode
phu0ngng
committed
EP PyTorch: gate symm-mem slot allocation on zero-copy; require expert_out to alias combine_in in zero-copy
phu0ngng
committed
EP PyTorch: rename buffer slots to dispatch_/combine_symm_buf; drop grad_tokens/grad_topk_weights; alias-check bwd grads in zero-copy
phu0ngng
committed
EP PyTorch: warn that ep_bootstrap(zero_copy=True) is experimental
phu0ngng
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
phu0ngng
committed
EP PyTorch: wire test_ep.py into L1 distributed QA suite
phu0ngng
committed
EP PyTorch: validate contiguity of dispatch/combine inputs and topk_weights fp32 dtype
phu0ngng
committed
EP PyTorch: check topk_idx/topk_weights token count matches tokens in ep_dispatch
phu0ngng
committed
EP PyTorch: move symm-mem allocation out of EpBuffer; make ep_dispatch/ep_combine accept caller-supplied output buffers with C++ symm-mem checks under zero-copy
phu0ngng
committed
EP PyTorch: enforce contiguous caller-supplied EP buffers in C++ and normalize bwd grad layout in Python
phu0ngng
committed
EP PyTorch: store autograd ctx tensors via save_for_backward, mark EP workspace buffers non-offloadable, and let backward stage autograd-allocated upstream grads instead of requiring symm-mem
phu0ngng
committed
EP PyTorch: reword EP buffer comments and drop redundant recv_topk_weights loss term in dispatch autograd test
phu0ngng
committed
EP PyTorch: rename build env var NVTE_BUILD_WITH_NCCL_EP to NVTE_WITH_NCCL_EP
phu0ngng
committed

Commits on Jun 24, 2026