Symmetric memory support

🚀 Feature

Goal: Add PyTorch Symmetric Memory Support

What is Symmetric Memory?

Symmetric memory enables GPUs to directly access each other's memory without CPU intervention, treating multiple GPUs as a unified memory space (PyTorch API doc).

Benefits

• In-kernel compute-communication fusion: Overlap computation with data movement within custom kernels
• Custom communication patterns: Implement application-specific strategies beyond standard collectives
• Low-latency remote access: Direct RDMA-based GPU-to-GPU communication without CPU's orchestration

High-Level TODO

• Support torch.distributed._symmetric_memory.empty() and rendezvous()
• Support torch.ops.symm_mem.* operations (all_reduce, all_to_all variants)
• Create DTensor for symmetric tensors
• Allow DTensors to use implementations of primitives backed with symmetric memory
• Add a transform to apply symmetric memory based communication for efficiency

cc @crcrpar @ali-alshaar7

Start with the linked PyTorch symmetric memory API and the requested entry points, torch.distributed._symmetric_memory.empty(), rendezvous(), and torch.ops.symm_mem.*. Map these against DTensor and the listed communication TODOs; the work is done when the requested symmetric tensor support, primitives, and communication transform are implemented and verified.