pjhkorea/branchless-nccl-router

branchless-nccl-router is an experimental Python project that explores a different way of coordinating communication between many GPUs working together on the same computation. In large distributed training or inference setups, multiple accelerators need to constantly check in with each other over the network to stay in sync, and normally this involves conditional checks: if this packet arrived correctly, do this, otherwise wait. Those checks can create delays because some devices end up waiting on others. This project's core idea is to remove those conditional branches entirely and replace them with pure math instead. Rather than asking a yes or no question about whether a packet is corrupted, it turns that answer into a number, zero or one, and multiplies data by that number directly. Because every device runs the exact same arithmetic operation regardless of what actually happened on the network, there is no branching logic for the hardware to wait around on, which the author argues removes a source of synchronization delay across a 32 GPU cluster. The project is built on JAX and its XLA compiler, along with NCCL, the library GPUs use to communicate with each other during distributed computation. It also describes techniques for keeping data in fast on chip memory instead of spilling it to slower memory, and for making sure this synchronization logic does not interfere with how a model learns during training, by explicitly cutting it off from the gradient calculations used in backpropagation. Much of the README is a mathematical argument for why this approach should scale well as clusters grow larger, comparing the cost of the arithmetic masking approach against the cost of network communication as the number of devices increases. This is a research style proof of concept rather than a ready to use product, and it assumes the reader already has background in distributed GPU training, JAX, and low level performance engineering. There is no indication in the README of an established user base or a lightweight setup path, it looks like an experimental blueprint for engineers or researchers working on large multi GPU systems.