Skip to content

Symmetric memory pytorch backends#6023

Merged
saivishal1999 merged 23 commits into
mainfrom
symmetric-memory-pytorch-backends
Apr 13, 2026
Merged

Symmetric memory pytorch backends#6023
saivishal1999 merged 23 commits into
mainfrom
symmetric-memory-pytorch-backends

Conversation

@saivishal1999
Copy link
Copy Markdown
Collaborator

@saivishal1999 saivishal1999 commented Mar 2, 2026

Adds support for PyTorch-backed symmetric memory in nvFuser alongside the existing native implementation. This enables using c10d::symmetric_memory (NCCL/NVSHMEM/CUDA) as an alternative backend while keeping the current SymmetricTensor API unchanged.

  • Introduced SymmetricMemoryBackend (Native, PyTorchNccl, PyTorchNvshmem, PyTorchCuda) with runtime selection.
    --Default behavior unchanged (native backend).
    --PyTorch backend is opt-in via NVFUSER_ENABLE=symmetric_memory_backend(...).
  • Integrated PyTorch symmetric memory into SymmetricTensor:
    allocate (empty_strided_p2p)
    setupRemoteHandles (c10d::symmetric_memory::rendezvous)
    remoteTensor(get_remote_tensor)
  • Added process group tracking + registration in Communicator to support rendezvous and ensure proper cleanup.
  • Extended c10d mocks to allow builds without distributed support.

Loading
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants