[PyTorch] NCCL EP zero copy with symmem pool and user provided recv_topk_weight tensor

[YangFei1990]

Description

Addition support for NCCL EP zero copy feature, including

• Expose user provided recv_topk_weights buffer
• Use symm memory pool to allocate symm buffer

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR replaces the static, EpBuffer-owned symm-mem buffer strategy with on-demand allocation from a process-global symm-mem MemPool, and exposes caller-supplied recv_tokens, recv_topk_weights, and grad_out kwargs on ep_dispatch/ep_combine so mcore can manage its own IO buffers.

• Pool allocation (_get_symm_mem_pool + _alloc_io): A single process-wide MemPool backed by the NCCL symm-mem allocator is created once at first dispatch; tensors are allocated per-call, ref-counted by PyTorch, and auto-registered with the NCCL comm — removing the need for per-buffer rendezvous and static lifetime management in EpBuffer.
• API simplification: EpBuffer.__init__ drops dispatch_recv_tokens / combine_grad_expert_out kwargs and all three *_symm_buf slots; callers pass buffers directly at ep_dispatch/ep_combine call time, and the new is_symm_backed helper lets callers verify pool-backed tensors.
• Tests updated: Tests migrate from buffer-construction kwargs to call-site kwargs; new test_zero_copy_pool_auto_alloc covers the mcore-primary path where neither recv nor grad buffer is supplied.

Confidence Score: 5/5

Safe to merge. The refactor correctly moves symm-mem buffer ownership out of EpBuffer and into on-demand pool allocation, with no changes to the C++ dispatch/combine kernels.

The allocation logic is consistent across forward and backward, mark_not_offload is applied everywhere pool tensors are created, and the backend-mismatch guard added in the previous review round is in place. No correctness regression was found in the dispatch or combine paths.

No files require special attention. The two observations are minor: ep_group is unused in the pool path of symm_mem_alloc, and the pool auto-alloc test could add a symmetry check for recv_topk_weights.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Adds _get_symm_mem_pool (process-global, created once with backend assertion) and extends symm_mem_alloc with use_pool/backend kwargs. Pool creation is defensive (tries two API shapes). ep_group is silently dropped in the pool path — a minor API design issue but no runtime impact on the current call sites.
transformer_engine/pytorch/ep.py	Replaces static EpBuffer-owned symm-mem buffers with on-demand pool allocation via _alloc_io. ep_dispatch/ep_combine now accept optional caller-supplied recv_tokens, recv_topk_weights, and grad_out kwargs. is_symm_backed added to __all__. Logic is sound; pool allocations call mark_not_offload correctly.
tests/pytorch/distributed/run_ep.py	Tests migrated from buffer-construction kwargs to dispatch/combine call-site kwargs, matching the new API. New test_zero_copy_pool_auto_alloc covers the primary mcore path. recv_w symmetry not asserted in the pool auto-alloc test.

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into nccl_ep_zero_cp" | Re-trigger Greptile}

[phu0ngng]

/te-ci pytorch L1

[phu0ngng]

/te-ci pytorch L1

[KshitijLakhani]

@phu0ngng and @YangFei1990 is this supposed to be cherry-picked for TE 2.18
If yes, please tag it. Thank you :)

cc: @fheinecke (Just FYI - no action needed)