Skip to content

Beyond Collectives: Building Fused Compute-Communication Kernels with PyTorch Symmetric Memory#

PyTorch Conference 2026 digest Program Adjacent pick 20–21 October 2026

Abstract

Ke Wen (NVIDIA), Natalia Gimelshein (Meta), Kapil Sharma (Meta, PyTorch Distributed)

  • Why: One-sided RMA issued from inside GPU kernels is the biggest change in how accelerators use the fabric since collectives; this is the guided tour of the programming model, and it pairs directly with the new torch.distributed window APIs.
  • What: PyTorch Symmetric Memory gives every rank a direct window into its peers' GPU memory, erasing the compute/communication boundary. A practical tour across five themes: compute–communication fusion (pushing partial results to a peer from inside a kernel), one-sided puts/gets plus torch.distributed.TokenSwitch supporting expert-parallel MoE with autograd, binding NCCL device primitives into a CUDA DSL (cute-DSL), device-side profiling plus CUDA Graph capture, and community adoption. (from abstract only)
  • How: Communication becomes something a kernel does inline, hidden behind compute, instead of a black-box collective the compute waits on.
  • Where applicable: NVIDIA-backed (NVSHMEM) today, with the AMD rocSHMEM counterpart covered separately; relevant to anyone reasoning about when fused communication beats a collective — and about what next-gen RDMA APIs need to expose.