Originally published on the PyTorch Developer Log.
TL;DR - Python can drive GPUDirect RDMA at line rate without a large transfer-engine abstraction.
rdma4pyprovides lightweight, backend-specific bindings for ibverbs, AWS EFA, and NVMe-oF, reaching about 400 gbps in our benchmarks while preserving direct access to the underlying APIs.
Introduction
Recently, quite a few different “transfer engines” have been developed for fast weight synchronization between machines using RDMA-capable transports and libraries such as NVLink, ibverbs, EFA, NVMe-oF, and SPDK. These include projects such as Mooncake, NIXL, and Uniflow. As with any large project and abstraction, these projects make trade-offs around specific use cases that might not be optimal for yours.
CPU-initiated RDMA transfers have started appearing in quite a few places in large-scale training and inference. They can be especially helpful for complex operations that do not fit well into existing collectives, such as checkpointing (GPU to disk and GPU to remote VRAM), reinforcement-learning weight synchronization (GPU to GPU), and fault tolerance (GPU to GPU).
There are a few key observations:
- Underlying libraries such as libibverbs have historically required C or C++, which creates a significant barrier to entry for ML researchers and projects that largely live in Python.
- Unlike protocols such as TCP, RDMA transfers tend to have very little CPU overhead for large payloads. It is therefore straightforward to reach line-rate speeds while driving them from Python.
- Most projects use only one hardware platform and network. Abstractions across many platforms can limit performance by giving up control in favor of generalization.
- Coding agents make it much easier to get RDMA working quickly, even with advanced features such as multi-QP or multi-lane transfers, which are required for maximum performance across machines in real-world clusters.
- RDMA APIs are much simpler than many people assume.
rdma4py
This brings us to rdma4py, a set of
lightweight bindings that exposes raw control over the underlying libraries
rather than trying to be an abstraction. rdma4py does not aim to be a
full-fledged transfer engine, though one could be built on top of it as needed.
It is currently a prototype and may have breaking API changes.
It provides bindings for the most common patterns we see:
- ibverbs: InfiniBand and RoCE
- EFA: Elastic Fabric Adapter on AWS
- NVMe-oF: NVMe over RDMA fabrics, for remote NVMe drives using
nvmet_rdma
Because each backend is largely standalone, it is easy to add new hardware bindings with full control without worrying about compatibility with the other backends. UCX and libfabric are both good candidates for future bindings.
Each backend is built as a separate package, so downstream users can install only the bindings they need:
uv pip install ibverbs efa nvmeof
The packages use Cython bindings with dlopen and dlsym, so the prebuilt
wheels should work across most Python, ibverbs, EFA, and kernel versions without
a custom build. They use Python’s stable ABI3 and support Python 3.9 and newer.
Benchmarks
To validate the design and show that it is more than a Python toy, we benchmarked
rdma4py on real hardware. The results demonstrate that Python is not a
bottleneck, even at the 400 Gbps line rate of the NICs.
We used both an ibverbs machine and an AWS EFA machine with four H100 GPUs. Each benchmark performed a point-to-point transfer between multiple NICs and two hardware devices:
- ibverbs: H100 -> NIC -> remote NIC -> H100
- EFA: H100 -> 4 EFA NICs -> 4 remote EFA NICs -> H100
- NVMe-oF: H100 -> NIC -> remote NIC -> NVMe, and the reverse read path
Although these benchmarks ran on single machines, that does not invalidate the result: they demonstrate that Python application code can achieve line-rate transfers. We used multiple queue pairs or lanes where possible and paired GPUs and NICs within a NUMA/PCIe node for maximum bandwidth.
EFA peaked at 48.88 GB/s, ibverbs at 48.55 GB/s, and NVMe-oF at 27.85 GB/s. NVMe-oF did not hit the RDMA NICs’ line rate, likely because of NVMe read/write I/O limits rather than overhead in our implementation.
For small transfers, EFA latency was 20-30 microseconds, ibverbs was about 5 microseconds, and NVMe-oF was 50-100 microseconds. NVMe-oF requires additional synchronization, which accounts for some of its extra latency.
The full setup and results are available in the backend-specific reports:
API
The high-level APIs for these backends are relatively simple, resulting in fairly minimal Python examples. A complete example is about 150 lines without abstractions. It is not a one-line solution, but it provides raw control and is tractable to implement without the overhead of a full transfer engine. LLMs make these algorithms even more accessible.
ibverbs
The following example performs a one-sided write from rank 0 to rank 1:
import ibverbs as ib
import ibverbs.cuda as ibcuda
...
WRITE_COMPLETE = 0xC0DE
def wait(cq):
while True:
completions = cq.poll(1)
if completions:
completions[0].raise_for_status()
return completions[0]
hca = ib.get_device_list()[0]
with hca.open() as ctx, ctx.alloc_pd() as pd:
port_attr = ctx.query_port(RDMA_PORT)
gid = ctx.query_gid(RDMA_PORT, GID_INDEX)
with ctx.create_cq(16) as cq:
qp_init = ib.QPInitAttr(
send_cq=cq,
recv_cq=cq,
qp_type=ib.QPType.RC,
max_send_wr=16,
max_recv_wr=16,
)
with pd.create_qp(qp_init) as qp:
# Exchange the information needed to connect the RC QPs.
local_qp = ib.local_qp_info(
qp, port_attr, gid, port=RDMA_PORT
)
local_wire = torch.tensor(
list(local_qp.to_bytes()), dtype=torch.uint8
)
qp_wires = [
torch.empty_like(local_wire)
for _ in range(world_size)
]
dist.all_gather(qp_wires, local_wire)
remote_qp = ib.QPInfo.from_bytes(
bytes(qp_wires[peer].tolist())
)
ib.connect_rc(
qp,
remote_qp,
port=RDMA_PORT,
sgid_index=GID_INDEX,
access=ib.AccessFlags.REMOTE_WRITE,
)
tensor = (
torch.arange(1024, dtype=torch.int32, device=gpu)
if rank == 0
else torch.zeros(1024, dtype=torch.int32, device=gpu)
)
mr_access = (
ib.AccessFlags.LOCAL_WRITE
| ib.AccessFlags.REMOTE_WRITE
)
with ibcuda.register_tensor(pd, tensor, mr_access) as mr:
target = torch.zeros(2, dtype=torch.int64)
if rank == 1:
target[0] = mr.addr
target[1] = mr.rkey
# WRITE_WITH_IMM consumes an RQ entry but does not
# place its payload in this receive request.
qp.post_recv(
ib.RecvWR(wr_id=1, sg_list=[])
)
# Publish rank 1's destination address and rkey. Because
# rank 1 posts its receive first, this also establishes
# notification readiness.
dist.broadcast(target, src=1)
if rank == 0:
# Wait only for work previously submitted to the
# stream that produced the source tensor.
torch.cuda.current_stream(gpu).synchronize()
qp.post_send(
ib.SendWR(
wr_id=2,
sg_list=[mr.sge()],
opcode=ib.WROpcode.RDMA_WRITE_WITH_IMM,
send_flags=ib.SendFlags.SIGNALED,
remote_addr=int(target[0]),
rkey=int(target[1]),
imm_data=WRITE_COMPLETE,
)
)
completion = wait(cq)
assert completion.opcode == ib.WCOpcode.RDMA_WRITE
else:
notification = wait(cq)
assert (
notification.opcode
== ib.WCOpcode.RECV_RDMA_WITH_IMM
)
assert notification.wc_flags & ib.WCFlags.WITH_IMM
assert notification.imm_data == WRITE_COMPLETE
with torch.cuda.device(gpu):
ibcuda.flush_gpudirect_writes()
expected = torch.arange(
1024,
dtype=torch.int32,
device=gpu,
)
assert torch.equal(tensor, expected)
Multiple queue pairs
For more performance, create multiple queue pairs and shard transfers across them. In practice, a thin wrapper could hide the multiple QPs behind register, send, receive, and signaling APIs:
with hca.open() as ctx, ctx.alloc_pd() as pd, ExitStack() as stack:
port_attr = ctx.query_port(RDMA_PORT)
gid = ctx.query_gid(RDMA_PORT, GID_INDEX)
cq = stack.enter_context(ctx.create_cq(NUM_QPS * 4))
qp_init = ib.QPInitAttr(
send_cq=cq,
recv_cq=cq,
qp_type=ib.QPType.RC,
max_send_wr=4,
max_recv_wr=4,
)
qps = [
stack.enter_context(pd.create_qp(qp_init))
for _ in range(NUM_QPS)
]
...
qp_access = ib.AccessFlags.REMOTE_WRITE if rank == 1 else 0
for index, qp in enumerate(qps):
...
ib.connect_rc(
qp,
remote_qp,
port=RDMA_PORT,
sgid_index=GID_INDEX,
access=qp_access,
)
...
with ibcuda.register_tensor(pd, tensor, mr_access) as mr:
...
for index, qp in enumerate(qps):
offset = index * chunk_size
qp.post_send(
ib.SendWR(
wr_id=index,
sg_list=[mr.sge(chunk_size, offset=offset)],
opcode=ib.WROpcode.RDMA_WRITE_WITH_IMM,
send_flags=ib.SendFlags.SIGNALED,
remote_addr=remote_addr + offset,
rkey=remote_rkey,
imm_data=index,
)
)
EFA
EFA is very similar to ibverbs, but it uses connectionless SRD queue pairs instead of establishing a persistent RC connection:
import efa
import efa.cuda as efa_cuda
...
device = efa.get_efa_device_list()[0]
with device.open() as ctx, ctx.alloc_pd() as pd, ctx.create_cq(16) as cq:
qp = pd.create_qp(
efa.QPInitAttr(send_cq=cq, recv_cq=cq)
).prepare(qkey=0x1234)
local_info = efa.local_endpoint_info(qp, qkey=0x1234)
# Exchange local_info.to_bytes() between ranks.
...
remote_info = efa.EndpointInfo.from_bytes(remote_bytes)
# EFA SRD is connectionless; each send identifies its destination.
with remote_info.peer(pd) as peer:
...
with efa_cuda.register_tensor(pd, tensor, access) as mr:
if rank == 1:
target = (mr.addr, mr.rkey)
qp.post_recv(efa.RecvWR(wr_id=1, sg_list=[]))
# Exchange rank 1's target address and rkey.
...
if rank == 0:
torch.cuda.current_stream().synchronize()
qp.post_send(
efa.SendWR(
wr_id=2,
sg_list=[mr.sge()],
opcode=efa.WROpcode.RDMA_WRITE_WITH_IMM,
send_flags=efa.SendFlags.SIGNALED,
remote_addr=target_addr,
rkey=target_rkey,
imm_data=0xC0DE,
dest=peer,
)
)
wait(cq)
else:
notification = wait(cq)
assert notification.imm_data == 0xC0DE
efa_cuda.flush_gpudirect_writes()
...
NVMe over Fabrics
The NVMe-oF API is the simplest of the three. It uses ibverbs under the hood with a protocol for reading from and writing to the storage device:
import nvmeof
...
with nvmeof.Controller.connect(
TARGET,
SUBSYSTEM_NQN,
source=SOURCE,
) as controller:
namespace = controller.namespace(NSID)
tensor = torch.empty(
namespace.lba_size, dtype=torch.uint8, device=gpu
)
with controller.register_gpu(tensor) as mr:
with torch.cuda.device(gpu):
namespace.read(mr, slba=SLBA, blocks=1)
print(
f"read LBA {SLBA} ({namespace.lba_size} bytes):",
tensor[:16].cpu(),
)
Next Steps
A PyTorch transfer engine
Now that we have flexible bindings for these transports, we could build a pure Python, hackable, topology-aware transfer engine on top of them. If there is interest, we would like to explore this further.
We would also like to experiment with integrating these APIs directly into
torchstore as optional transports, leveraging its existing weight-management
and coordination layers.
Other backends
It should be straightforward to add backends such as libfabric, which do not currently have Python bindings.
UCX has ucxx on PyPI, but it is unclear
how widely it has been adopted, and we have not yet investigated all the backends
it supports out of the box. It is also an abstraction, so it does not provide
full access to each backend API.
GPU-initiated operations
rdma4py currently supports only CPU-initiated operations. We should be able to
add helper code that triggers transfers from custom kernels. This could make it
possible to write extremely lightweight custom communication kernels entirely in
Python, without depending on existing backends such as NCCL GIN.