NVSHMEM Deep Dive: One-Sided GPU Communication, and How It Differs from NCCL

This is a companion to the NCCL deep dive, which covered the physical interconnects (NVLink/NVSwitch/PCIe/IB) and how collectives are bootstrapped. I’ll reference that here rather than repeat it, and focus on what NVSHMEM does differently. The motivating use case — fused MoE dispatch/combine — is in the Mixture-of-Experts post.

Why another comms library?

If you’ve run any PyTorch distributed training, you’ve used NCCL: dist.all_reduce, dist.all_gather, the FSDP/TP/DDP machinery. NCCL is a collective library — every rank calls the same collective, the GPUs synchronize, bulk data moves. It is the right tool for “sum these gradients across 8 GPUs.”

But a whole class of modern kernels — fused tensor-parallel matmuls, and especially MoE dispatch/combine — wants something NCCL fundamentally isn’t built for: a single GPU thread, inside a running kernel, reaching out and writing a few bytes into another GPU’s memory, while the rest of the kernel keeps computing. That’s NVSHMEM, and it’s a different programming model entirely.

1. NVSHMEM vs NCCL — the core distinction

The one sentence: NCCL is host-launched, collective, two-sided; NVSHMEM is device-initiated, one-sided, PGAS.

	NCCL	NVSHMEM
Model	collective (all ranks participate)	one-sided RMA (any PE → any PE)
Who initiates	the host enqueues a collective on a stream	a GPU thread/warp, from inside a kernel
Primitive	all_reduce, all_gather, all_to_all	put, get, atomic, signal
Granularity	bulk (whole tensors)	fine (a row, a scalar, a tile)
Matching	two-sided (everyone calls the matching op)	one-sided (target is passive — no recv)
Overlap w/ compute	separate kernel on another stream	same kernel — comm warps + math warps
Addressing	rank + buffer you pass in	global: a symmetric pointer is valid on every PE

PGAS = Partitioned Global Address Space. NVSHMEM gives every rank (“PE”, Processing Element) a symmetric heap: when all PEs call nvshmem_malloc(n), they each get a buffer at the same offset, and a pointer to it is meaningful on every PE. PE 3 can put into PE 7’s copy of that buffer directly — no send/recv handshake, PE 7 doesn’t even have to call anything.

How each shows up in the PyTorch world

This is the part people are usually fuzzy on:

• NCCL is the default torch.distributed backend. init_process_group(backend="nccl") → every all_reduce/all_gather/reduce_scatter in DDP, FSDP, TP, PP is a NCCL collective. ~99% of PyTorch distributed is this. You never write a comm kernel; you call a collective and it runs on a stream.

• NVSHMEM is the “I’m writing a custom fused communication kernel” path. You reach for it when the comm has to be interleaved with compute at sub-kernel granularity. In PyTorch it surfaces through torch.distributed._symmetric_memory (the symm_mem API — see §4), which is the engine behind Async Tensor Parallel (overlapping the all-gather with the matmul that consumes it) and behind custom MoE kernels like DeepEP. You allocate symmetric tensors and write a Triton/CUDA kernel that does the remote loads/stores yourself.

So the rule of thumb: NCCL when you want a collective; NVSHMEM when you want to fuse the communication into a compute kernel. They’re not competitors so much as different layers — in fact NCCL ≥ 2.30 even added a device-side API (“Gin”) that moves it toward NVSHMEM’s territory.

2. The NVSHMEM programming model

The symmetric heap

// Every PE runs this. All get a buffer at the SAME symmetric address.
float* buf = (float*) nvshmem_malloc(N * sizeof(float));
int    mype  = nvshmem_my_pe();
int    npes  = nvshmem_n_pes();

buf is a symmetric object. nvshmem_malloc is collective (all PEs call it), but after that, any PE can address any other PE’s buf with the same pointer + a target PE id. (Heap-local malloc memory is not remotely accessible — only the symmetric heap is.)

One-sided primitives

// PUT: write my data INTO PE `peer`'s symmetric buffer. Peer does nothing.
nvshmem_float_put(dst /*remote symm addr*/, src /*my local*/, nelems, peer);

// GET: pull data FROM PE `peer`'s buffer into my local memory.
nvshmem_float_get(dst /*my local*/, src /*remote symm addr*/, nelems, peer);

// ATOMIC + SIGNAL: increment a remote counter, or set a remote flag.
nvshmem_uint64_atomic_add(remote_counter, 1, peer);
nvshmemx_signal_op(remote_flag, value, NVSHMEM_SIGNAL_SET, peer);

The initiator does everything; the target is passive. There’s no matching recv. This is what makes irregular, data-dependent communication (a token going to whichever expert the router picked) natural — you don’t need every PE to agree on a collective shape up front.

The killer feature: the device API

NVSHMEM ops come in host and device flavors. The device ones are callable from inside a CUDA kernel:

__global__ void moe_dispatch(...) {
    // GEMM warps compute the current tile...
    // ...while THIS warp ships the next PE's tokens, concurrently:
    nvshmemx_putmem_warp(remote_tokens, local_tokens, nbytes, owner_pe);
    // signal the remote consumer that its tile has landed
    nvshmemx_signal_op(remote_ready, 1, NVSHMEM_SIGNAL_SET, owner_pe);
}

This is the whole reason NVSHMEM exists for MoE. A NCCL all_to_all is a separate kernel — you can’t start the expert GEMM until it finishes. With NVSHMEM, comm warps and GEMM warps run in the same kernel, so the dispatch hides behind the math (the warp-specialized fused MoE in the MoE post does exactly this). On NVLink, NVSHMEM can even hand you a raw peer pointer via nvshmem_ptr(addr, peer) and you just do ordinary ld/st to the other GPU’s memory.

Completion and ordering

One-sided comms need explicit ordering — there’s no recv to tell you data arrived:

• nvshmem_fence() — order puts to the same PE.
• nvshmem_quiet() — wait until all my outstanding puts have completed remotely.
• nvshmem_signal_wait_until(flag, CMP, val) / nvshmem_uint64_wait_until(...) — spin on a remote flag until a producer signals (the classic producer/consumer handshake).
• nvshmem_barrier_all() — collective barrier.

The idiom is put-then-signal / wait-then-read: producer writes data, fences, sets a flag; consumer waits on the flag, then reads. That’s how the fused MoE kernel coordinates comm warps with GEMM warps tile by tile.

Teams and bootstrap

• PEs are ranks; teams are subgroups (NCCL’s “communicator”): nvshmem_team_my_pe(team), nvshmem_team_n_pes(team), with NVSHMEM_TEAM_WORLD as the global one. You map a torch process group to a team to do expert-parallel comms over just the EP ranks.
• Bootstrap is how PEs discover each other — and it’s a real gotcha. Two common paths:
  • PMI (nvshmem_init under srun --mpi=pmi2 / MPI) — needs a PMI process manager (SLURM/mpirun).
  • UID (nvshmemx_get_uniqueid → broadcast the id → nvshmemx_set_attr_uniqueid_args → nvshmemx_init_attr) — no launcher required; you broadcast the unique id out-of-band, e.g. over a torch process group. This is the “init-from-PG” pattern, and it’s what lets NVSHMEM bootstrap under plain torchrun (vs. needing SLURM). I hit this exact wall recently: a tuner hard-wired to PMI wouldn’t run on a non-SLURM pod until I swapped it to the UID path and launched with torchrun --no-python.

3. PyTorch’s symmetric-memory API

You rarely call libnvshmem directly from Python. PyTorch wraps the symmetric-memory model in torch.distributed._symmetric_memory (still a private _-API, but it’s what Async-TP and the fused-comms ecosystem are built on). The shape of it:

import torch.distributed._symmetric_memory as symm_mem

symm_mem.set_backend("NVSHMEM")          # or "CUDA" (intranode P2P), "NCCL"
assert symm_mem.is_nvshmem_available()

# Allocate a SYMMETRIC tensor (same address on every rank)
t = symm_mem.empty(1024, 1024, dtype=torch.bfloat16, device="cuda")

# Establish the symmetric handle across the group
hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)

# Now reach a peer's buffer directly, and synchronize with signal pads
peer_buf = hdl.get_buffer(peer_rank, shape, dtype)   # P2P-accessible
hdl.barrier(channel=0)
symm_mem.put_signal(src, hdl, peer); symm_mem.wait_signal(hdl, peer)

The important pieces, by exact name:

• empty(*size, dtype, device) — allocate on the symmetric heap.
• rendezvous(tensor, group) -> _SymmetricMemory — the collective handshake that makes the tensor remotely addressable; returns a handle.
• handle .get_buffer(rank, shape, dtype, storage_offset=0) — a pointer to peer rank’s copy, usable directly (NVLink) or via NVSHMEM (multi-node).
• .barrier(), signal pads, put_signal/wait_signal — the ordering/handshake primitives.
• set_backend("NVSHMEM"|"CUDA"|"NCCL") — the same symmetric-memory front-end, three transports: CUDA = intranode NVLink P2P, NVSHMEM = multi-node, NCCL = via NCCL’s window API.

The payoff is Async Tensor Parallel: instead of all_gather(x) then matmul(x_full, W) as two stages, you write one kernel that streams the remote shards in (symmetric get_buffer / multimem) while the matmul consumes the shards that have already landed — overlapping the TP communication with the TP compute. Same trick as MoE, different collective.

4. The transport layer, and the advanced bit: IBGDA vs IBRC

The same put/get runs over whatever transport connects two PEs — and the NCCL post already covered the physical layer (NVLink/NVSwitch/PCIe/IB). NVSHMEM picks:

• Intranode → NVLink P2P: direct GPU-to-GPU loads/stores or copy-engine transfers. No NIC involved. (Our single-node 8-GPU MoE benchmarks live entirely here.)
• Internode → InfiniBand RDMA, and here’s where the advanced distinction lives: IBRC vs IBGDA.

Both are NVSHMEM IB transports; they differ in who rings the NIC doorbell to start a transfer.

IBRC — InfiniBand Reliable Connection (host-proxied)

The GPU can’t post to the NIC itself, so:

1. the GPU kernel drops a request into a host-visible queue,
2. a CPU proxy thread picks it up and rings the NIC doorbell on the GPU’s behalf,
3. the NIC does the RDMA.

Critical path: GPU → CPU proxy → NIC. Works on any IB system. (This is the same proxy-thread idea NCCL uses for its NET transport — see the NCCL post’s “Proxy Threads” section.)

IBGDA — InfiniBand GPUDirect Async (GPU-initiated)

The NIC’s doorbell + work queue are mapped into GPU memory, so a GPU warp rings the doorbell directly — “GPUDirect Async, Kernel-Initiated.” No CPU in the loop.

Critical path: GPU → NIC.

	IBRC	IBGDA
Rings the NIC	CPU proxy thread	GPU directly
Latency	higher (extra GPU→CPU→NIC hop)	much lower
Many small messages	proxy thread serializes → bottleneck	scales (no proxy)
Bandwidth (large)	~same	~same
Requirements	any IB	NIC + driver + NVSHMEM build must support GPUDirect Async

The headline: IBGDA wins on latency and small-message concurrency, not bandwidth. That is exactly the MoE decode regime — tiny per-step token counts, many small dispatch/combine transfers, latency-bound. It’s why DeepEP’s low-latency (decode) kernels are built on IBGDA, and why its internode kernels won’t even link against an NVSHMEM that lacks it.

And the practical catch I ran into: IBGDA needs the NIC/driver/firmware to support it, and NVSHMEM has to be built with it. Plenty of clusters aren’t there yet — ours wasn’t, so the NVSHMEM build was IBRC-only, and the team’s note was blunt: “our pods do not seem to support IBGDA (the good stuff) and need to use the older IBRC backend (higher latency, more or less same bandwidth).” That table, observed in the wild.

One clarification that trips people up: IBGDA vs IBRC only matters inter-node. Single-node multi-GPU is NVLink P2P — no IB, no doorbell-ringing question at all. The distinction only appears once your PEs span hosts.

5. When to reach for which

• Use NCCL for everything collective and bulk: gradient all-reduce (DDP), FSDP shard all-gather/reduce-scatter, plain TP/PP. It’s mature, autotuned, and you never write a kernel.
• Use NVSHMEM (via symm_mem) when the win is overlap at sub-kernel granularity: fused async-TP (hide all-gather behind the matmul), and fused MoE dispatch/combine (hide the all-to-all behind the expert GEMMs). You’re writing a comm kernel because no collective can express “send this row to that expert while I keep computing.”
• The line is blurring: PyTorch’s symm_mem can sit on a CUDA, NVSHMEM, or NCCL backend, and NCCL 2.30’s device-side API encroaches on NVSHMEM’s turf. But the mental model holds: collective + host-launched vs one-sided + device-initiated.

TL;DR

NCCL moves whole tensors between ranks with host-launched collectives; NVSHMEM lets a GPU thread reach into a peer’s symmetric memory mid-kernel with one-sided put/get/signal, so communication can be fused into compute. In PyTorch, NCCL is the dist backend you already use; NVSHMEM shows up through torch.distributed._symmetric_memory for async-TP and fused MoE. And once you go multi-node, the single most important knob is IBGDA vs IBRC — GPU-initiated vs CPU-proxied RDMA — which decides whether your decode-time all-to-all is latency-bound or not.

---

Companion posts: NCCL Deep Dive (collectives, transports, torchrun bootstrap) and Mixture-of-Experts Across the Stack (where the fused NVSHMEM dispatch/combine actually pays off).