How does OEq compare to FlashTP for higher l operations?

[CompRhys]

FlashTP seems to outperform CuEq by a lot for higher angular momentum ops. They sadly do not benchmark against OEq this is just a general inquiry as to whether this is on the mantainers radar.

https://openreview.net/pdf?id=wiQe95BPaB
https://github.com/SNU-ARC/flashTP?tab=readme-ov-file

[vbharadwaj-bk]

Hmm- we did some benchmarks on Nequip / MACE-style configs and for inference we outperform, training we are competitive against FlashTP.

Worth noting: their paper compares against an older version of CuEq. CuEq is now much more competitive for higher irreps (see our roofline plot in the paper, where we go up to L=7.

Can you define "higher angular momentum"? Some sample input representations would be helpful.

[CompRhys]

Can you define "higher angular momentum"? Some sample input representations would be helpful.

by higher angular momentum I meant l>2 from the table 2 in the paper.

[vbharadwaj-bk]

Sure - we can add to our radar to extract / benchmark some of those configs.

[vbharadwaj-bk]

Below: some benchmarks for flashTP, forward and backward, on what I believe are the irreps used in their experiments. Tested fused convolutions on 3 different graphs.

Takeaways:

1. The latest version of cuE (which was released in March) is much more competitive than the one benchmarked in the FlashTP paper (v0.2.0), and generally has the fastest backward pass.
2. We are very competitive on the backward pass and often 2-3x faster than flashTP.
3. Not shown are double-backward kernels for training. Haven't benchmarked, but given recent improvements to our kernel fusion strategy, I'd wager we'd beat flashTP and hold our own (although probably worse) against cuE.

def flashtp_higher():
    return [
        CTPP(
            "32x0e+32x0o+32x1e+32x2e+32x3e+32x4e+32x1o+32x2o+32x3o+32x4o",
            "1x0e+1x1o+1x2e+1x3o+1x4e",
            "32x0o+32x0o+32x0o+32x0o+32x0o+32x0e+32x0e+32x0e+32x0e+32x0e+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e",
            label="flashtp-l4",
        ),
        CTPP(
            "32x0e+32x0o+32x1e+32x2e+32x3e+32x4e+32x5e+32x1o+32x2o+32x3o+32x4o+32x5o", 
            "1x0e+1x1o+1x2e+1x3o+1x4e+1x5o", 
            "32x0o+32x0o+32x0o+32x0o+32x0o+32x0o+32x0e+32x0e+32x0e+32x0e+32x0e+32x0e+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1o+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x1e+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2o+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x2e+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3o+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x3e+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4o+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x4e+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5o+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e+32x5e", 
            label="flashtp-l5",
        ),
    ]

[vbharadwaj-bk]

Let's convert this to a discussion now that the main issue has been answered. Can follow up there if needed.