geforce-p2p

geforce-p2p moves data between GPUs whose driver denies peer-to-peer, and is loadable through kernels. NVIDIA's driver refuses P2P between GeForce cards, so a device-to-device copy detours through a host buffer one direction at a time. Two cards sit on two PCIe links, and the hardware will run both at once; this package does, with a pipelined cross-device copy, a ring all_reduce and all_gather over pinned host memory, and a probe that names the exact reason P2P is off for a pair and whether the driver patch in PATCH.md would turn it on. Where P2P is on, copy uses it directly.

Usage

import torch
from kernels import get_kernel

gp = get_kernel("phanerozoic/geforce-p2p", version=1, trust_remote_code=True)

gp.probe()                                   # devices, pairs, verdicts, link rates
a = torch.randn(1 << 28, device="cuda:0")
b = gp.copy(a, device="cuda:1")              # both links busy, not one
grads = [torch.randn(1 << 26, device=f"cuda:{i}") for i in range(2)]
gp.all_reduce(grads, op="avg")               # in place, every device gets the mean

API

Compiled ops, under gp.ops: geforce_p2p_staged_copy, geforce_p2p_reduce_from_host, geforce_p2p_copy_to_host, geforce_p2p_copy_from_host.

Method

staged_copy splits the transfer into chunks and drives a ring of pinned slots: each chunk's D2H runs on the source device's current stream, its H2D on the destination's, with an event from landing to upload and another from upload to slot reuse. The two DMAs are on two devices, so they overlap, and the destination stream is fenced at entry so a reused staging buffer is never overwritten while a previous call is still uploading from it. Both devices' torch stream orderings are preserved.

all_reduce is a ring reduce-scatter followed by a ring all-gather. A step stages one chunk with a D2H DMA on the sender; the receiver pulls it into a VRAM scratch with an H2D DMA and applies the op in place. Pinned memory that torch allocates is mapped into the device address space, so reduce_from_host can instead read the staged bytes directly from host memory and fold them in one pass with no scratch; zero_copy=True selects it. Each byte crosses each link once per ring step either way.

probe reads can_device_access_peer for each pair, BAR1 and VRAM sizes, the Windows driver model and the PCIe link through NVML when pynvml is present, and /proc/driver/nvidia/version on Linux for the module flavour. The verdict follows from those: WDDM never exposes P2P; the proprietary module cannot take the patch; a BAR1 smaller than VRAM needs Resizable BAR first; open modules with a full-size BAR1 leave only the driver's refusal.

Measured

RTX 6000 Ada, PCIe 4.0 x16, one card. Rates in GiB/s, 1 GiB transfers.

The same-card staged_copy puts every byte across one link twice, 23.2 GiB/s of link traffic against a 24.4 ceiling, so the pipeline runs the link flat out; on two cards each half has its own link. In-kernel access to mapped host memory runs at under half the DMA rate, which is why the DMA scratch path is the default for all_reduce.

Correctness

tests/test_geforce_p2p.py checks reduce_from_host against torch for every op and dtype at lengths including non-multiples of the vector width, the host copies for round-trip equality, staged_copy for chunking, slot reuse and back-to-back staging reuse on one card, and that a refused unpinned pointer leaves no sticky error for the next launch. Two-device cases for copy, all_reduce in both transports, and all_gather run when a second GPU is present and skip otherwise.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+. Host tensors passed to the compiled ops must be pinned (pin_memory=True); anything else is refused.
• copy and the collectives need every tensor on a distinct CUDA device with matching shape and dtype. all_reduce and all_gather require contiguous tensors.
• probe reports BAR1, driver model and link details only when pynvml is installed.
• The two-device paths are exercised by the test suite only on a host with two or more GPUs.

References

tinygrad, open-gpu-kernel-modules branch 565.57.01-p2p (the driver-side enablement, described in PATCH.md); NVIDIA CUDA Runtime API, cudaHostAlloc and peer access; Thakur, Rabenseifner, Gropp, "Optimization of Collective Communication Operations in MPICH" (IJHPCA 2005) for the ring algorithms.

License

Apache-2.0.