diff --git a/content/distributed/2026-07-21-rdma4py.md b/content/distributed/2026-07-21-rdma4py.md new file mode 100644 index 0000000..b3a6130 --- /dev/null +++ b/content/distributed/2026-07-21-rdma4py.md @@ -0,0 +1,425 @@ +--- +title: "rdma4py: Do We Need Transfer Engines for PyTorch and RDMA?" +date: 2026-07-21 +author: "Tristan Rice (@d4l3k)" +tags: [distributed, rdma, networking, gpudirect, storage] +--- + +> **TL;DR** - Python can drive GPUDirect RDMA at line rate without a large +> transfer-engine abstraction. `rdma4py` provides 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](https://github.com/kvcache-ai/Mooncake), +[NIXL](https://github.com/ai-dynamo/nixl), and +[Uniflow](https://github.com/meta-pytorch/torchcomms/tree/main/comms/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: + +1. 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. +2. 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. +3. Most projects use only one hardware platform and network. Abstractions across + many platforms can limit performance by giving up control in favor of + generalization. +4. 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. +5. RDMA APIs are much simpler than many people assume. + +## rdma4py + +This brings us to [rdma4py](https://github.com/d4l3k/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: + +```shell +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. + +![Measured GPUDirect bandwidth by message size](/devlogs/images/distributed/rdma4py-bandwidth.png) + +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. + +![Measured GPUDirect completion latency by message size](/devlogs/images/distributed/rdma4py-latency.png) + +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: + +- [EFA benchmarks](https://github.com/d4l3k/rdma4py/blob/main/efa/BENCHMARKS.md) +- [ibverbs benchmarks](https://github.com/d4l3k/rdma4py/blob/main/ibverbs/BENCHMARKS.md) +- [NVMe-oF benchmarks](https://github.com/d4l3k/rdma4py/blob/main/nvmeof/BENCHMARKS.md) + +## 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: + +```python +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: + +```python +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: + +```python +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: + +```python +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](https://pypi.org/project/ucxx-cu12/), 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. diff --git a/static/images/distributed/rdma4py-bandwidth.png b/static/images/distributed/rdma4py-bandwidth.png new file mode 100644 index 0000000..d61d6c4 Binary files /dev/null and b/static/images/distributed/rdma4py-bandwidth.png differ diff --git a/static/images/distributed/rdma4py-latency.png b/static/images/distributed/rdma4py-latency.png new file mode 100644 index 0000000..736c2f1 Binary files /dev/null and b/static/images/distributed/rdma4py-latency.png differ