Collectives

An all-reduce, an all-gather, a reduce-scatter, an all-to-all, a collective-permute: each of these shows up in an HLO dump as a single instruction, and that appearance is misleading in a specific way worth correcting early. No single device is in charge of a collective. Every participant in what the runtime calls a clique has to independently issue the matching call, and the operation only completes once all of them have.

That is the sense in which a collective is a contract rather than a command. Nobody sends an all-reduce to the fleet the way you might send a message to one recipient; every device's own compiled program contains the identical instruction, at the point the partitioner decided a sharding boundary demanded it, and the correctness of the whole thing depends on every copy of that program agreeing to run it.

A collective is a contract between every participant.

Break that agreement and the failure mode is deadlock: one device waiting on a collective that another device never issues, or issues in a different order relative to some other collective it also needed. The kernel path met this exact failure at the lowest possible level, raw semaphoresThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici → in a Pallas kernel at /s/distributed, a send waiting on a recv that is itself waiting on that same send.

This chapter meets the identical failure two layers higher up, where the ordering constraint is enforced by the compiled schedule rather than by a semaphoreThe counter a DMA signals on completion and a kernel waits on; the synchronization primitive under every transfer.taught in /l/ici → you wrote by hand. The mechanism looks different at each layer. The failure is the same failure, because the contract being violated, every participant issues the same collective in the same order, does not change shape depending on which layer you are reading it at.

A synchronous all-reduce forces compute to stop and wait for the network. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → has an escape from that: an async pair, all-reduce-start followed later by all-reduce-done, splits one logical collective into an instruction that launches the communication and a second instruction that waits for it to finish.

Everything scheduled between the start and the done runs concurrently with the network transfer, which is the entire point. A scheduling pass decides how much independent compute fits into that gap, stretching it as wide as the dataflow allows, and the wider that gap gets filled, the less the collective costs the step that contains it.

Host offloading, moving a tensor out to host memory across some region of the program and back, sounds unrelated to a collective on its face, but it rests on the identical piece of machinery: an async value that resolves once an operation has actually finished, not once it was merely launched. A collective's start instruction and a host transfer both hand back a token immediately; what differs is only what that token is waiting on underneath.

That shared machinery is what makes both overlap patterns real rather than theoretical. Compute can run ahead of a collective's completion, or ahead of a host transfer's completion, because the program has an honest way to ask later whether the thing it is waiting on is actually done, instead of guessing based on how much time has passed.