Metadata-Version: 2.4
Name: purlin
Version: 0.6.4
Summary: Adaptable, template-based substrate for high-performance inter-GPU communication
Author-email: Osayamen Jonathan Aimuyo <osayamen@stanford.edu>
License-Expression: BSD-3-Clause
Project-URL: Homepage, https://github.com/osayamenja/purlin
Project-URL: Repository, https://github.com/osayamenja/purlin
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.8.0
Dynamic: license-file

# Purlin
**Purlin is a high-performance GPU communication framework for moving and reducing data across GPUs within a scale-up domain, such as an NVLink-connected server.** 

It provides both ready-to-use, single-kernel collectives and reusable device-side primitives for building custom communication and fusing it with computation.

Purlin includes AllReduce, AllGather, ReduceScatter, and AllToAll, plus variable-length variants of the latter three. 

Underneath these collectives are two hardware-aware primitives: **copy** and **N-to-1 reduce**, which combines multiple input buffers into one output.

## Key Idea
The key idea is **decoupling orchestration from the datapath**: separating *where and when* data moves from *how* the GPU moves it. 

- **Layouts describe the collective:** how inputs and outputs are distributed, and whether to copy or reduce.
- **SNAC coordinates execution:** the shared *Stage, Notify, And Consume* (SNAC) protocol derives orchestration from those layouts.
- **Atoms move the data:** hardware-specific implementations perform the copies and reductions as coordinated by SNAC.

This separation makes Purlin **evolvable**: 
- New hardware mechanisms can be added through Atoms without rewriting orchestration 
- Communication is much easier to customize either at the collective level through our layouts or via composing our Atom building blocks. 

Also, Purlin allows for extensive tuning (see [codesign](csrc/include/purlin/host/codesign.cuh)) to achieve peak performance.

## 🧨 QuickStart
```bash
uv pip install purlin # best to use a venv here
torchrun --nproc-per-node <num-of-gpus> quickstart.py 
```

## C++ benchmarks
<details>
<summary>Click here to see steps</summary>

Dependencies:

- CUDA Toolkit
- [NVSHMEM](https://developer.nvidia.com/nvshmem-downloads?target_os=Linux) (for C++ symmetric memory)
- MPI (to launch processes)
- CMake 3.27+
- Ninja
- [CPM](https://github.com/cpm-cmake/cpm.cmake#adding-cpm) (for CMake dependency management)

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

```bash
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
```

| Collective | Binary |
| --- | --- |
| AllReduce | `testAR` |
| AllToAll / AllToAllV | `testA2A` / `testA2AV` |
| AllGather / AllGatherV | `testAG` / `testAGV` |
| ReduceScatter / ReduceScatterV | `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:

```bash
# 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]`.

</details>

## Atom implementations

| Atom | File |
| --- | --- |
| `Atom<700>` | [fascia.cuh](csrc/include/purlin/fascia.cuh) |
| `Atom<800>` | [tendon.cuh](csrc/include/purlin/tendon.cuh) |
| `Atom<900>` | [ligament.cuh](csrc/include/purlin/ligament.cuh) |
| `Atom<1000>` | [cortex.cuh](csrc/include/purlin/cortex.cuh) |

## Reduction determinism
Enable deterministic mode for repeatable results with the same inputs and operator. All ranks must use the same mode.

C++:

```cpp
purlin::allReduce<..., purlin::ReductionMode::deterministic>(src, dst, bytes, ctx, stream);
```

Python:

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

The mode only changes the datapath on SM90 and newer: non-deterministic mode
allows multimem, while deterministic mode disables it. On older GPUs (Ampere and below), both
modes use the deterministic path.

Default is non-deterministic.
