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

Merged

phu0ngng merged 44 commits into

NVIDIA:mainfrom

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

Jun 29, 2026

Commits on Jun 23, 2026

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

pre-commit-ci[bot]
authored and
phu0ngng
committed
d5a6c09
EP PyTorch: skip combine_in staging copy in non-zero-copy mode
phu0ngng
committed
2e9a89e
EP PyTorch: gate symm-mem slot allocation on zero-copy; require expert_out to alias combine_in in zero-copy
phu0ngng
committed
e9cdaed
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
5f29d90
EP PyTorch: warn that ep_bootstrap(zero_copy=True) is experimental
phu0ngng
committed
91edff0
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
phu0ngng
committed
0415e3a
EP PyTorch: wire test_ep.py into L1 distributed QA suite
phu0ngng
committed
bab2bdc
EP PyTorch: validate contiguity of dispatch/combine inputs and topk_weights fp32 dtype
phu0ngng
committed
6696030
EP PyTorch: check topk_idx/topk_weights token count matches tokens in ep_dispatch
phu0ngng
committed
090d05b
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
926bef2
EP PyTorch: enforce contiguous caller-supplied EP buffers in C++ and normalize bwd grad layout in Python
phu0ngng
committed
2ca14a7
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
8cb815f
EP PyTorch: reword EP buffer comments and drop redundant recv_topk_weights loss term in dispatch autograd test
phu0ngng
committed
53a3834
EP PyTorch: rename build env var NVTE_BUILD_WITH_NCCL_EP to NVTE_WITH_NCCL_EP
phu0ngng
committed
a4cb0cb

Commits on Jun 24, 2026

Commits on Jun 25, 2026

Commits on Jun 28, 2026

Commits on Jun 29, 2026

added test skip when no NVlink
phu0ngng
committed
a8b7702
Merge branch 'main' into phuong/ep-3-pytorch-on-commwindow
phu0ngng
authored
4e63776
Merge branch 'main' into phuong/ep-3-pytorch-on-commwindow
phu0ngng
authored
7ff285c