nbshmem: Enabling GPU-Initiated Multi-GPU Communication in Python

In this paper we present nbshmem, a Python library for GPU-initiated GPU-to-GPU communication. The library can be used within Numba CUDA kernels that are compiled into GPU device code at runtime. nbshmem is designed with NVSHMEM in mind, but its implementation poses several challenges due to Python’s language constraints compared to C. One major challenge is that Python does not support pointer arithmetic, making it difficult to realise a symmetric address space. To overcome this, our approach uses Python tuples to manage shared memory objects. Another challenge is the lack of volatile operations in Python, which are required for certain synchronisation mechanisms. This paper presents and evaluates different solutions to this problem. As a proof of concept, we have implemented several collective operations. However, in most cases their performance remains below that of NCCL (in Python) and NVSHMEM, highlighting areas for potential improvement. To evaluate performance, we implemented a stencil computation kernel that requires frequent data exchange between neighbouring GPUs. Our results show that