rocSHMEM Symmetric Memory in PyTorch for AMD GPUs#

Prachi Gupta — AMD

• Why: Ports of GPU-initiated one-sided RMA to a second vendor are evidence this is becoming a cross-vendor fabric primitive rather than an NVIDIA trick — worth tracking for multi-vendor estate planning.
• What: rocSHMEM backend for PyTorch symmetric memory, now upstream: the device-side OpenSHMEM API (one-sided put/get, non-blocking get, put-with-signal, wait-until) exposed as Triton-callable primitives so a kernel initiates its own communication. MoE all-to-all dispatch and combine are constructed within the kernel and overlapped with expert compute; comparison against RCCL host-driven all-to-all on representative MoE workloads promised. (from abstract only)
• How: AMD-specific engineering — per-architecture device-bitcode linking and post-compile HIP-module initialization — organized behind a backend-agnostic layer shared with NVSHMEM so a single Triton kernel runs on both AMD and NVIDIA.
• Where applicable: AMD GPUs (ROCm); the transferable design is the backend-agnostic layering that keeps one kernel portable across vendors.