Message passing on TPU

Animation of E3J’s fused message-passing convolution forward kernel, e3j/pallas_ops/convolution/mosaic_tpu/fwd.py, which computes:

\[{\bf m}_b = \frac 1 N \sum_a ({\bf x}_a \otimes {\bf y}_{ab}) \odot {\bf w}_{ab}\]

The kernel fuses the whole convolution into one pass over receiver-sorted edges. Per block of edges, the sender features \({\bf x}_a\) are gathered into VMEM, where the Clebsch-Gordan paths, unrolled at trace time, build the whole block’s messages. The block is then reduced eight edges at a time, each edge accumulating on its own sublane, and a change of receiver sums the sublanes into one row and writes it out as \({\bf m}_b\) — one write to HBM per receiver.

The vector registers of a TPU hold 8x128 tiles, and VMEM lays arrays out to match. Channels are therefore padded to a multiple of 128 lanes, and the harmonics and output axes to 8 sublanes. The edge block used is 128 and the reduction is made in chunks of 8 edges at a time to fill the TPU’s vector registers.

Note

The kernel actually expects sender-sorted edges, so that the backward pass, which operates on the transposed graph, can be more efficient. Graphs are assumed symmetric, so senders and receivers are morally interchangeable.