cuequivariance vs stock e3nn: a flat sub-unity region below break-even, and where break-even lands (MACE-MP-0, A5000)

[liulangdog1]

While profiling small-system MLIP workloads I ended up characterising where cuequivariance starts to pay for itself against stock e3nn. The shape turned out to be more consistent than the crossing point, so I'm posting both. Everything below is a single fused forward with forces on a pre-built batch — kernel-level, not an end-to-end workflow number.

The shape: a flat penalty, then a rise

Below break-even the ratio is not gradually degrading, it is constant. Small/float32 sits at 0.72–0.73 across every batch size from 1 through 48, then lifts. The same flat-then-rise appears in all six model × precision combinations tested.

At batch 1, cuequivariance is 0.62–0.73× of stock e3nn in all six — a narrow band across a 4× parameter range and both precisions.

That is consistent with a fixed per-launch kernel overhead: constant in relative terms until there is enough work per launch to hide it.

Ratio (cuequivariance ÷ stock e3nn), float32:

model	params	b1	b8	b16	b24	b32	b48	b64	b96	b128
small	3,847,696	0.72	0.73	0.72	0.72	0.72	0.73	0.79	0.85	1.07
medium	4,688,656	0.62	0.63	0.63	0.63	0.62	0.72	0.94	1.34	1.82
large	15,847,440	0.64	0.65	0.65	0.79	1.01	1.44	1.88	2.75	3.67

float64:

model	b1	b2	b3	b4	b8	b10	b16	b32	b64	b128
small	0.73	0.73	0.72	0.72	0.90	1.08	1.56	2.96	5.48	6.44
medium	0.65	0.70	0.88	1.11	2.01	2.48	3.81	7.50	14.21	17.75
large	0.69	1.15	1.58	2.05	3.88	4.84	7.61	11.42	12.53	13.09

Break-even, as brackets

Reported as the pair of measured batch sizes that straddle 1.0. Nothing interpolated — the sweep spacing does not support a single value.

model	float32	float64
small	b96 → b128	b8 → b10
medium	b64 → b96	b3 → b4
large	b24 → b32	b1 → b2

Monotonic in model size in both precisions.

What moves it, and by how much

The crossing is not portable. Three things move it:

• Model size — monotonic, as above.
• Precision — 12–20× for the same model on the same card. Small crosses near b96–128 in float32 and b8–10 in float64. On an A5000 the float64 rate is heavily limited, so compute dominates sooner and the kernels pay off much earlier. This should look different on a card with a stronger float64 path.
• Atoms per structure — see the caveat below.

Note that parameter count is a poor abscissa here. All three checkpoints report num_interactions: 2, so they differ in angular momentum rather than depth: small → medium is +22% parameters but moves the crossing by roughly 1.5 brackets, while medium → large moves it much further. I did not capture hidden_irreps on this run, so I can't yet plot against a proper work-per-structure axis.

One thing I could not reconcile

A separate run on 32-atom fcc Cu, using a different harness, put the small model's crossing near batch 17 rather than b96–128. Recasting the crossing as total atoms per forward pass rather than batch count narrows the discrepancy from roughly 7× to under 4×, which suggests atoms-per-batch is the better abscissa — but the two runs still do not fully reconcile, and I have not established why. Flagging it rather than smoothing it over.

Method

RTX A5000 (sm86, 23 GB), driver 575.57.08, CUDA 12.6. torch 2.12.0+cu126, mace-torch 0.3.16, cuequivariance 0.10.0. Conversion via mace.cli.convert_e3nn_cueq.run with "convert": true recorded. Ethanol (C₂H₆O, 9 atoms) from ASE's G2 database. MACE-MP-0 small/medium/large. Best of 5 after 3 warm-ups. Single node, single card, single driver throughout.

Instrument check. Before any timing was read, torch.ops was confirmed to expose ['cuequivariance', 'cuequivariance_ops'] with fused_tensor_product present. Worth stating because the namespace moved: on cuequivariance 0.6.0 the ops live under torch.ops.cuequivariance_ops and the forward is fused_tensor_product_fwd, so the obvious check false-negatives on older versions and you can silently benchmark a fallback path.

Cross-script reproduction. The float64 brackets above (small b8→b10, medium b3→b4) match an independent run with a different harness on the same system (small b8→b16, medium b2→b4).

One question

The contributing guide says direct code contributions are paused during the public beta. Is that still current? If a reproducible benchmark script for this would be useful to have in the repo I'm happy to open one; otherwise I'll leave the numbers here.