Receiver-streaming convolutions for channel-wise UVU tensor products

[pl-fuchs]

I recently came across the Sobek paper and, with the help of Codex, explored whether its receiver-side streaming idea could accelerate channel-wise UVU tensor products in MACE. My implementation is available on the streaming-kernel branch.

The main difference from the standard OEQ convolution kernels is the ownership. In the forward pass, each thread owns one (receiver, channel) pair, scans the edges incident to that receiver, and reduces the convolution result locally. In the backward pass, each warp owns one edge, processes the channels in parallel, writes radial-weight gradients directly, reduces spherical-harmonic gradients across channels within the warp, and performs atomic additions for node feature gradients.

I benchmarked prepared, receiver-sorted tensor-product workloads on a synthetic 5,000-node, 280,000-edge topology with a uniform degree of 56 edges per receiver, chosen to approximate the coordination of bulk HCP titanium with a 6 Å cutoff. I obtained the following timings on an NVIDIA A100 GPU (script streaming_speedup.py):

Case	Forward speedup	Full TP VJP speedup	HVP speedup
MACE medium	2.34×	1.32×	3.93×
MACE large	2.86×	1.67×	4.75×
MH-1-like wide synthetic proxy	2.75×	3.42×	Not measured¹

¹ The standard HVP exceeded the A100's 80 GB memory, while the streaming HVP completed.

Graph construction, including sorting, was outside the timed region. In my LAMMPS workflow, neighbor lists group neighbors by their central atom. Preserving the central atom as the MACE receiver when flattening the list therefore produces receiver-major edges without a separate sorting pass.

These results look promising to me, and I intend to use the kernels in my own work.
However, I would be happy to use this fast path directly from OEQ.
I have checked the kernel code and integration and provided it here as proof of concept.

My questions are:

1. Are you already investigating related receiver-streaming approaches?
2. Would a specialized path for supported UVU convolutions be useful in OEQ?
3. If you find this direction useful, how could I best support its integration into OEQ?

[vbharadwaj-bk]

Huh, cool - didn't know their optimization would be this straightforward to apply- yeah agreed I think we should put this in the package and welcome the contribution. I'm a bit fuzzy on how the scheme compares to what we're doing (need to stare at the code), but I can't argue with the speedup, and the memory movement logic seems reasonable.

In terms of integration: if you can slot this in by modifying the existing kernel templates (in a minimal-ish way), go ahead. If you find you need a total rewrite of the templates or would have to modify them in a way that would slow down other tensor products, then the best thing to do would be adding a new template file and we can branch on which one is selected for compilation at runtime.

Feel free to open a PR and credit Sobek appropriately; we can version this so that their paper gets credit for speed up over the older version. @asglover was looking at Sobek as well, so he can review this diff.

[pl-fuchs]

Thanks for the quick reply. My implementation is inspired by SOBEK rather than a one-to-one realization of the paper’s proposal. The streaming kernel accelerates UVU convolutions, including UVW cases that can be reduced to UVU, but does not implement SOBEK’s general UVW scheme. It also still relies on radial weights being materialized beforehand in JAX.

I’ve implemented the streaming path in a separate template, since I found it difficult to reuse the standard convolution templates. Both paths share a common interface, with a mode argument for explicitly selecting either implementation where supported. The auto mode selects the streaming path for UVU and reducible UVW convolutions and falls back to the standard path otherwise.
Since my primary use case is MACE with UVU convolutions, this specialization has been simple and effective.

[asglover]

I just saw this paper, the ideas are very cool, and the performance makes sense. I'll take a look at the proposed diff. I would like improvements to be general, vs special cased. So supporting uvu, uvw all datatypes etc.

I think we need to make two implementations of the convolution op. One which "source stationary" One which is "destination stationary". The better performer will be the one that moves less data, so should be determinable by analyzing the input and output irreps. So we should be able to dispatch at init time.

I think we can aim to incorporate these techniques to make the speedups available from OEQ

Like you mentioned, fusing the realization of the radial weights implies a new op. I don't know how important this fusion is to overall performance. So I'm not going to address that part of the paper for now either.

But we should be able to start with the two different versions of convolution straightforwardly.

I'll take a closer look at your diff and the paper and come back with a plan. I'm busy at work right now but will try to push this forward. I might try to push a few bug fix prs before adding this functionality. But I will commit to integrating it and trying to achieve performance parity in the long run.

[vbharadwaj-bk]

@pl-fuchs I'd open a PR with this and I can take a look this week - just drained some other thing from the queue and @asglover can iterate when it gets merged