Optimize UMA SO2 and Wigner inference kernels

[mlazos]

What this changes

UMA spends much of its inference time moving edge data between GPU operations. This PR removes several large temporary tensors and combines adjacent operations so the GPU reads and writes less memory.

In plain terms, the six optimization commits do the following:

1. Keep edge features in the format already produced by the first layer instead of converting them to another format and immediately converting them back.
2. Combine the gate activation into one GPU operation.
3. Process more edges in parallel.
4. Write backward gradients directly to atoms instead of first building a large per-edge gradient tensor.
5. Rotate edge messages and add them to atoms in one GPU operation.
6. Store only the nonzero parts of each rotation matrix.

On a 1000-atom UMA-S-1p2 inference workload, the complete stack improves GPU throughput from 19.722 to 25.881 QPS (+31.22%) and reduces steady peak allocated memory by 11.18%. CUDA graphs were not enabled for this comparison.

How this relates to #2073

This work builds on #2073, which is already included in the base revision.

#2073 added GPU kernels that combine the initial node-to-edge rotation and the final edge rotation. However, the model still created large temporary edge tensors around those kernels. This PR removes more of those temporary tensors and reduces the amount of rotation data stored for every edge.

A simplified view is:

Before this PR:

atom features
  -> rotate and gather them onto edges
  -> build a large edge tensor
  -> apply several gate operations
  -> convert the result back into blocks
  -> run the second layer
  -> build another large rotated edge tensor
  -> add edge messages to atoms

After this PR:

atom features
  -> rotate and gather them onto edges
  -> apply the gate directly to the existing blocks
  -> run the second layer
  -> rotate and add messages directly to atoms

The six optimizations

1. Avoid an unnecessary format conversion

The first convolution already produces three feature blocks in the format expected by the second convolution. Previously, UMA joined these blocks into one [edges, 9, channels] tensor for the gate activation and then split the tensor back into blocks.

The model now keeps the blocks in their original format across the gate. This removes a full-tensor join, split, and memory round trip.

For the representative 78,848-edge graph with 128 channels, one FP32 [edges, 9, channels] tensor occupies about 346.5 MiB.

Forward and backward results were compared with the old path using:

relative tolerance: 1e-6
absolute tolerance: 1e-7

2. Run the gate as one GPU operation

The gate previously launched separate operations for sigmoid, SiLU, multiplication, splitting, and joining. The new Triton kernel performs those steps together and writes the block-form output directly.

The fused operation supports normal force inference and its first derivative. When a caller requests second derivatives through this gate, it uses a differentiable PyTorch implementation rather than the specialized first-derivative kernel. This fallback applies to the gate itself; the other existing umas_fast_gpu Triton operations retain the backend's existing higher-order derivative restrictions.

Measured agreement with the original gate:

Result	Relative tolerance	Absolute tolerance
Forward output	1e-5	1e-6
First derivative	1e-5	1e-6
Second derivative	2e-5	2e-6

Compiled and uncompiled executions of the new kernel matched exactly in the dynamic-shape tests.

3. Process more edges in parallel

The existing Wigner kernels always launched 2,048 GPU programs. A graph with 78,848 edges therefore made each program process many edges one after another.

The launch size now follows the number of edges, up to a limit of 131,072 programs. For the representative graph, this allows roughly one program per edge.

This does not change the formula or combine results from different edges; it only exposes more independent work to the GPU. The standard forward, backward, empty-graph, and dynamic-shape correctness tests pass. Across the four affected kernels at 78,848 edges, combined kernel time decreased from 2.8196 ms to 2.7111 ms.

4. Write input gradients directly to atoms

Backward previously created an [edges, 9, 2 * channels] tensor and then ran two additional reductions to add its values to source and target atoms.

The new kernel writes those contributions directly into the atom-gradient tensor. At 78,848 edges and 128 channels, this avoids an approximately 693.1 MiB temporary tensor. The isolated backward operation improved from 2.361 ms to 0.935 ms.

Because multiple edges may update the same atom concurrently, the order of FP32 additions can change. The measured numerical effect was:

atom-feature gradient relative L2 difference: 2.97e-7
rotation gradient: bitwise equal
radial gradient: bitwise equal

Here, relative L2 means the total error magnitude divided by the reference result magnitude.

When deterministic algorithms are enabled, UMA keeps the previous reduction path. Repeated forward and backward executions of that path are bitwise identical.

5. Rotate and add edge messages in one operation

Previously, UMA wrote an inverse-rotated [edges, 9, channels] tensor and then read it again to add edge messages to their target atoms. Backward similarly gathered atom gradients into another edge tensor.

The new kernel performs the rotation and atom update together. The isolated forward operation improved from 0.565 ms to 0.201 ms.

This operation also changes the order of some FP32 additions:

output relative L2 difference: 2.15e-7
input gradients: bitwise equal
rotation gradients: bitwise equal

When deterministic algorithms are enabled, UMA uses the previous materialized reduction. Its forward and backward results are bitwise repeatable.

6. Store only the nonzero rotation values

UMA-S uses a 9 by 9 block-diagonal rotation matrix. Only three diagonal blocks can contain nonzero values:

1 by 1 block:  1 value
3 by 3 block:  9 values
5 by 5 block: 25 values
               ---------
               35 values

The old representation stored all 81 values, including 46 values that are always zero. The new representation stores only the 35 useful values. For 78,848 edges, this saves approximately 13.8 MiB per rotation tensor, or 27.7 MiB for the forward and inverse tensors before counting gradients.

The compact rotation generator is bitwise equal to the corresponding blocks of the dense generator, including its position gradients. The edge-degree output and input gradient are also bitwise equal. Double-precision radial and rotation gradient comparisons were within 2e-15 and 1e-14 relative tolerance respectively.

Legacy entry points still accept compact rotations by reconstructing the dense matrix when needed.

Independent review follow-up

An independent review found two compatibility gaps, fixed in follow-up commit 802979b9:

• The backend can be selected for non-power-of-two channel widths or under autocast, while the fused gate kernel requires power-of-two FP32 inputs. The gate now checks the actual convolution outputs and uses the existing block-form PyTorch activation when the Triton kernel does not apply. This preserves the packed convolution results rather than rerunning the convolutions.
• The exported fused Wigner operations had stopped accepting their previous dense [edges, 9, 9] rotation input. They now pack dense block-diagonal inputs before custom autograd, preserving the public input shape, dense gradient shape, and zero off-block gradients.

The added tests cover 96-channel FP32 fallback, 128-channel BF16 fallback, dynamic full-graph compilation, dense forward and backward compatibility, and changing edge counts. An independent re-review found no remaining blocking issues.

The fixed branch measured 25.946 QPS with a 38.541 ms GPU median and 5.956 GiB steady peak allocation under the same configuration below. This is effectively unchanged from the published 25.881 QPS result and produced no timed recompilations.

When these optimizations run

These changes are used by the umas_fast_gpu backend. That backend is selected explicitly with:

InferenceSettings(
    activation_checkpointing=False,
    merge_mole=True,
    execution_mode="umas_fast_gpu",
)

It can also be selected automatically when execution_mode=None and all of the following are true:

• CUDA is available
• the model uses lmax=2 and mmax=2
• activation checkpointing is disabled
• merge_mole=True
• the requested Hessian configuration is supported

The optimized kernels are not selected by execution_mode="general".

Additional fallbacks remain available:

• non-FP32 and non-power-of-two-channel gate inputs use the block-form PyTorch implementation
• S2 activation uses the existing materialized path
• deterministic algorithms use the previous deterministic reductions
• second-order gate differentiation uses differentiable PyTorch operations
• legacy backends can reconstruct dense rotation matrices
• exported fused Wigner operations accept either compact or dense rotation inputs

CUDA graphs are not required to use these kernels.

End-to-end performance

Standalone H100 measurements compare clean detached checkouts of base a3fa357c3 and optimization head 522ce9df. Follow-up head 802979b9 contains the compatibility fixes described above and has unchanged performance.

Configuration:

model: uma-s-1p2
system: 1000-atom Cu
GPU: NVIDIA H100
PyTorch: 2.14.0a0+gitc0c824e
CUDA: 12.9
TF32: enabled
activation checkpointing: disabled
merge_mole: enabled
torch.compile: enabled, dynamic shapes
execution mode: umas_fast_gpu
external graph generation: disabled
internal graph generation: NVIDIA v3
CUDA graphs: disabled

Measurement	Base	This PR	Change
GPU median latency / QPS	50.704 ms / 19.722	38.639 ms / 25.881	+31.22% QPS
GPU minimum latency	49.121 ms	36.905 ms	-24.87%
Wall median latency / QPS	50.523 ms / 19.793	39.256 ms / 25.474	+28.70% QPS
Wall minimum latency	49.272 ms	37.053 ms	-24.80%
Overall peak allocated memory	6.705 GiB	5.970 GiB	-10.95%
Steady peak allocated memory	6.705 GiB	5.955 GiB	-11.18%
Reserved memory after warmup	8.219 GiB	7.561 GiB	-8.01%

Each process performed 24 warmup predictions followed by eight blocks of 100 timed predictions. Base and head processes were interleaved in base/head/head/base order, and each process used a fresh TorchInductor cache. No recompilation occurred during timed blocks. QPS is 1000 / median_latency_ms.

End-to-end numerical comparison

Five frames used deterministic position displacements of -0.02, -0.003, 0, 0.003, and 0.02. The table compares the complete optimization stack with the base. The final column shows the normal difference observed when the same revision is run twice.

Output	Maximum absolute difference	Relative L2 difference	Same-revision maximum difference
Energy	3.813e-3	6.270e-7	4.073e-3
Forces	2.199e-4	5.778e-4	2.407e-4
Stress	1.518e-5	1.279e-4	1.467e-5

The complete-stack differences are comparable to the run-to-run variation already observed from the same revision. The two operations that intentionally change FP32 addition order were also tested directly, producing relative L2 differences of 2.97e-7 and 2.15e-7.

Suggested review order

The first six commits follow the model's execution order:

1. preserve the existing block format
2. combine the gate operations
3. increase edge parallelism
4. remove the large backward temporary tensor
5. combine rotation and atom accumulation
6. compact the rotation representation

The seventh commit contains the independent-review compatibility fixes and their regression tests.

Test plan

PYTHONPATH="$PWD/src" python -m pytest -q   tests/core/models/uma/nn/test_so2_layers.py   tests/core/models/uma/test_quaternion_wigner.py   tests/core/models/uma/uma_fast/test_execution_backends.py   tests/core/models/uma/uma_fast/test_fused_edgewise.py   -c packages/fairchem-core/pyproject.toml

/home/mlazos/.conda/envs/pytorch-3.12/bin/pre-commit run   --files $(git diff --name-only a3fa357c3..HEAD)

All selected correctness tests and all 21 fused-edgewise tests reached pass status. The final review-focused run passed four gate/backend cases and three dense/compact dynamic Wigner cases. In this local custom PyTorch/Triton environment, one Triton-heavy test process subsequently hit a known interpreter teardown abort after pytest reported 100% passed.

Authored with assistance from an AI coding tool.

[bkmi]

this is all I got!

[bkmi]

is there any backwards compatibility we should think about here?

[mlazos]

fairchem/src/fairchem/core/models/uma/escn_md_block.py

Line 276 in e4425dc

g0, g1, g2 = self.so2_conv_2.gemms_to_buffers(x_message)

is the only usage, so I think we are good!

[bkmi]

seems good to me... maybe it's worth having someone else look too.