# rocSHMEM Symmetric Memory in PyTorch for AMD GPUs

PyTorch Conference North America 2026 · [All sessions](https://sessionize.com/api/v2/s2rq62nk/view/markdown)

- **Session ID:** 1256215
- **When:** Tuesday, October 20, 2026 4:35 pm - 4:45 pm (Pacific Standard Time, UTC-08)
- **Room:** 210B (Concourse Level)
- **Speakers:** Prachi Gupta

## Description

Mixture-of-Experts (MoE) models are increasingly bottlenecked by all-to-all token dispatch and combine rather than the expert computation. PyTorch symmetric memory addresses this by issuing communication from within device kernels, including Triton, allowing overlap with computation. To date it has been available only on NVIDIA, backed by NVSHMEM.

This work extends that capability to AMD GPUs through rocSHMEM, which is now available in upstream PyTorch. We describe how rocSHMEM's device-side OpenSHMEM API is exposed as Triton-callable primitives, including one-sided put/get, non-blocking get, put-with-signal, and wait-until, enabling a kernel to initiate its own communication. We discuss the AMD-specific engineering: per-architecture device-bitcode linking and post-compile HIP-module initialization, organized behind a backend-agnostic layer shared with NVSHMEM so that a single Triton kernel runs on both AMD and NVIDIA. Using these one-sided RMA and signal/wait primitives, all-to-all dispatch and combine are constructed within the kernel and overlapped with expert computation. We conclude with a comparison against RCCL host-driven all-to-all on representative MoE workloads.

## Speaker

- [Prachi Gupta](https://sessionize.com/api/v2/s2rq62nk/view/markdown/speaker/afc23c94-0247-4201-a2cc-b78853bd9ac3) — Engineering Manager, PyTorch Dev Team @ AMD