Reduce CUDA memory use in non-local SKALA layers

[Dynamics of Condensed Matter (DCM-Uni-Paderborn)]

Summary

• use the model precision for CUDA fine-to-coarse reductions and geometry tensors while retaining the existing double-precision CPU path
• avoid materializing the large fine-grid skip-connection concatenation on CUDA by applying the existing linear projection in two parts
• preserve the model parameters, state-dict layout, and CPU execution path

Motivation

CUDA profiles of native-grid SKALA in CP2K show that temporary fine-grid tensors contribute substantially to peak device memory. In particular, the skip-connection concatenation creates a large short-lived allocation and an associated copy kernel. The double-precision CUDA reduction and geometry intermediates also add work without improving the float32 model evaluation.

The split projection is mathematically equivalent to the existing concatenated linear layer:

W [features, update] + b = W_features features + b + W_update update

It therefore reuses the existing weights and does not change checkpoint compatibility.

Validation

• pytest -q --ignore=tests/test_pyscf_gradients.py: 179 passed, 1 skipped
• pytest -q tests/test_pyscf_gradients.py: 36 passed
• ruff check . and ruff format --check .
• CUDA tests verify reduction and geometry dtypes, forward equivalence, gradients, and the absence of the concatenation path

Small-system CP2K checks found maximum candidate-to-reference differences of approximately 1.2e-8 Ha in energy, 3e-9 Ha/bohr in forces, and 0.0014 bar in stress.

For a 96-atom CP2K SKALA workload, using three-run medians:

• NVIDIA GB10, one GPU: CP2K time 33.648 -> 30.997 s (-7.9%), SKALA time 22.418 -> 19.833 s (-11.5%), and peak GPU memory 11416 -> 10016 MiB (-12.3%) relative to the current Rev1 CUDA model
• two NVIDIA A40 GPUs, isolated split-projection comparison: peak GPU memory 10224/8896 -> 8864/7414 MiB (about 13-17% lower per GPU) and CP2K time 25.086 -> 24.918 s (-0.7%)

Model artifact

The source changes take effect after scripting/tracing a new CUDA .fun artifact. Existing published artifacts are not modified by this PR. A new CUDA model revision can be uploaded and added to the version mapping after this code is accepted; the CPU artifact can remain unchanged.

Related to cp2k/cp2k#5439.

[Dynamics of Condensed Matter (DCM-Uni-Paderborn)]

@microsoft-github-policy-service agree

[Jens (JensWehner)]

cool, thanks for the PR. We will have a look

[Chin-Wei Huang (CW-Huang)]

thanks Dynamics of Condensed Matter (@DCM-Uni-Paderborn) for the PR. There are two changes in this PR:

1. dtype casting: We explicitly hardcode some parts to double precision to handle summation of grid points. Casting this to model.dtype would lose precision. I wouldn't recommend merging this without extensively testing SCF convergence statistics.
2. splitting linear(merge(...)): splitting this in theory has a smaller footprint + slightly more cuda launch overhead. the report you shared didn't test this separately right? I suspect it's very marginal only with this. But if the goal is to reduce memory, may I suggest to look into gradient checkpointing or chunking by atoms?