Host-side

Create a context once and reuse it for your collectives.

First, set up workspace with symmetric staging and signal buffers.

You can do that via makePurlinWorkspace or for a more generic but involved way, see purlin_initialize.

#include <purlin/host.cuh>

auto ctx = purlin::initialize(rank, world, workspace, stream);
purlin::allReduce<ARCH, __nv_bfloat16>(src, dst, bytes, ctx, stream);
purlin::finalize(ctx, stream);

src and dst are regular pointers to device memory.

Set ARCH for your GPU, for example 900 for SM90. The collective runs on stream.

Call finalize after your last collective to free internal buffers.

Device-side

Use <purlin/core.cuh> to call an Atom or a collective from inside your kernel.

See examples as a starting point.

For tuning, see host code, tuning policies.

Samples

In the example below, PurlinAtom::THREADS copy one aligned slice.

src or dst pointers can refer to local GPU memory or remote memory on a peer GPU:

#include <purlin/core.cuh>

template<typename PurlinAtom>
__global__ void copyKernel(cuda::std::byte* dst, const cuda::std::byte* src,
                          size_t bytesPerBlock) {
  extern __shared__ __align__(128) cuda::std::byte workspace[];
  const auto offset = blockIdx.x * bytesPerBlock;
  PurlinAtom::copy(dst + offset, src + offset, bytesPerBlock, workspace);
}

Launch it with:

• PurlinAtom::THREADS threads per block.
• PurlinAtom::COPY_PIPELINE_SMEM_BYTES bytes of dynamic shared memory.

See the point-to-point example for a full example.

To call a collective, pass the context you created on the host like below

#include <purlin/core.cuh>
template<typename PurlinAtom, typename CollConfig>
__global__ void allGatherKernel(const cuda::std::byte* src, cuda::std::byte* dst,
                            size_t bytesPerRank, const __grid_constant__ purlin::Context ctx) {
  extern __shared__ __align__(128) cuda::std::byte workspace[];
  // do some work
  const int blocks = static_cast<int>(gridDim.x);
  const purlin::SnacArgs args{
    .dst = dst,
    .src = src,
    .bytes = bytesPerRank,
    .workspace = workspace,
    .blocks = blocks,
    .collBlocks = blocks,
    .bIdx = static_cast<int>(blockIdx.x)
  };
  purlin::allGather<PurlinAtom, CollConfig>(args, ctx);
  // do some other work
}

Use PurlinAtom::THREADS threads per block. See here for a more complete example

Exactly args.blocks CTAs must enter the collective, where each CTA has PurlinAtom::THREADS threads.

First, initialize torch.distributed and select the local CUDA device.

Here, x is a contiguous CUDA tensor of float32 values. Create an output tensor, then call AllReduce:

import torch
import torch.distributed as dist
import purlin

stream = torch.cuda.current_stream(device)
major, minor = torch.cuda.get_device_capability(device)
handle = purlin.initialize(
    dist.group.WORLD, device, major * 10 + minor, stream.cuda_stream
)
out = torch.empty_like(x)
purlin.all_reduce(x, out, handle, stream.cuda_stream)
purlin.finalize(handle, stream.cuda_stream)

The first initialize call compiles the CUDA code and caches it for reuse. See quickstart.py for a complete script.

Dependencies:

• CUDA Toolkit
• NVSHMEM (for C++ symmetric memory)
• MPI (to launch processes)
• CMake 3.27+
• Ninja
• CPM (for CMake dependency management)

Set NVSHMEM_LIB_HOME to your NVSHMEM library directory. From the repository root:

export NVSHMEM_LIB_HOME=/path/to/nvshmem/lib
cmake -S csrc -B csrc/build -G Ninja -DCMAKE_BUILD_TYPE=Release
cmake --build csrc/build --target testAR testA2A testA2AV testAG testAGV testRS testRSV

📏 Understanding totalBytes

The reported totalBytes is what we use to compute bandwidth. It is the output size for AllGather. For the other collectives, it is the input size per rank (also the output size for AllReduce and AlltoAll).

Run with one MPI process per GPU; for example, on 8 GPUs:

# each rank gets as input 1K...1G with output also the same size for the below
NVSHMEM_BOOTSTRAP=MPI NVSHMEM_REMOTE_TRANSPORT=none mpirun -n 8 ./csrc/build/testAR 1K 1G
# each rank gets as input: 128...128M but output is multiplied by world for allGather so 1K...1G  
NVSHMEM_BOOTSTRAP=MPI NVSHMEM_REMOTE_TRANSPORT=none mpirun -n 8 ./csrc/build/testAG 128 128M

Arguments: [minBytes] [maxBytes] [graphLaunches] [runs] [warmup] [seed].

Enable deterministic mode for bit-wise repeatable results with the same inputs and operator.

👉 This configuration only affects Hopper and above because non-deterministic mode allows NVLS, while deterministic mode disables it.

For Ampere and below, both modes use the deterministic path always which gives peak performance there.

All ranks must use the same mode.

C++:

// this call uses deterministic mode
purlin::allReduce<..., purlin::ReductionMode::deterministic>(src, dst, bytes, ctx, stream);
// this call uses the default mode
purlin::allReduce<...>(src, dst, bytes, ctx, stream);

Python:

purlin.all_reduce(..., reduction_mode=purlin.ReductionMode.DETERMINISTIC)

Default is non-deterministic.