Fuse compiled single-system periodic shifts

[mlazos]

Summary

For a single system, Inductor currently lowers the fixed-width cell-offset
matrix multiply as an external operation. That prevents it from fusing with
the following edge-vector construction.

While Dynamo is tracing, express the three-column transform as scaled cell
rows so Inductor can fuse the work. Eager execution retains the existing
matrix multiply, and the multi-system path is unchanged.

Performance

With torch.compile(dynamic=True), periodic fusion improved the full endpoint from 16.87322 to 16.83233 ms with fresh caches, a 0.04089 ms (0.24%) gain, and from 16.87650 to 16.82935 ms with warm caches, a 0.04715 ms (0.28%) gain. It removed four external calls.

The isolated dynamic shift and edge-vector benchmark improved from 0.07501 to 0.03298 ms, a 56.0% gain.

Configuration: H100, TF32, UMA-S-1p2, 1,000 atoms, energy/forces/stress, umas_fast_gpu, merge_mole=True, external_graph_gen=False, internal graph v3, internal_graph_skin=0, compile=True, compile_dynamic_shapes=True, and full CUDA graph replay. Steady-state execution had zero recompiles and zero dynamic graph fallbacks.

Testing

CUDA_VISIBLE_DEVICES=0 TORCHINDUCTOR_CACHE_DIR=$PWD/agent_space/cache_shift_compile_test PYTHONPATH=$PWD/src:/data/users/mlazos/pytorch /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m pytest -q tests/core/models/uma/test_compile.py::test_compile_single_system_external_graph
ruff check --ignore PT023,B905 src/fairchem/core/models/uma/escn_md.py tests/core/models/uma/test_compile.py
git diff --check origin/main...HEAD

Pytest reported 1 passed; the process subsequently hit the previously
observed CUDA interpreter-shutdown abort.

Authored with assistance from Codex.

[mlazos]

AI-Review

Automated multi-agent review (2 Claude + 2 Codex + synthesis). This is AI-generated — please verify findings manually.

---

No verified issues found.