Reduce CUDA memory use in non-local SKALA layers

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

Summary

• 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, numerical precision, and CPU execution path
• retain the explicit double-precision grid reductions and geometry calculations

Motivation

CUDA profiles of native-grid SKALA in CP2K show that the skip-connection concatenation creates a large short-lived fine-grid allocation and an associated copy kernel.

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

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

It reuses the existing weights and does not change checkpoint compatibility. This is primarily a memory optimization; the runtime effect is workload dependent.

Validation

• pytest -q tests/test_model.py: 14 passed on an NVIDIA GB10
• ruff check src/skala/functional/model.py tests/test_model.py
• ruff format --check src/skala/functional/model.py tests/test_model.py
• CUDA tests compare the forward result, all five relevant input gradients, and the projection weight and bias gradients with the original concatenation path

A six-step, 96-atom CP2K SCF trajectory on the GB10 differed from the unchanged reference by at most 5.13e-8 Ha in energy and 1e-8 in the convergence measure.

For the isolated split-projection comparison on two NVIDIA A40 GPUs, using three-run medians with identical weights and the unchanged double-precision reduction path:

• peak GPU memory decreased from 10253/8945 MiB to 9093/7883 MiB, or about 11-12% per GPU
• CP2K time changed from 26.943 s to 25.792 s and wall time from 30.174 s to 29.001 s
• the maximum energy spread across all six runs was 3.34e-8 Ha

An earlier isolated comparison showed only a 0.7% runtime difference, so no general speedup is claimed.

Model artifact

The source change takes 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?

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

Thanks, we agree with both points.

We have removed the two dtype changes and retained the explicit double-precision grid reductions and geometry calculations. The updated PR now contains only the split linear projection.

We repeated the split-only comparison using identical weights and the unchanged FP64 reduction path. On two A40 GPUs, three-run medians show peak memory decreasing from 10253/8945 MiB to 9093/7883 MiB, or about 11-12% per GPU. CP2K time changed from 26.943 s to 25.792 s in this run, although an earlier isolated comparison showed only a 0.7% difference, so we regard this primarily as a memory optimization and do not claim a general speedup.

We also strengthened the CUDA test to compare both the forward result and all five relevant input gradients against the original concatenation path. A six-step, 96-atom CP2K SCF trajectory differed by at most 5.13e-8 Ha in energy and 1e-8 in the convergence measure.

Atom chunking is already used independently by the CP2K native SKALA path, so the split projection reduces the remaining allocation within each atom subchunk. Gradient checkpointing is worth investigating separately because it trades activation memory for recomputation during every VXC backward evaluation.

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

Rebased onto the current main after #102 and extended the CUDA equivalence test to compare the projection weight and bias gradients in addition to the forward result and all five input gradients. The previous seven PySCF jobs did not reach the tests: collection failed because the environment combined h5py 3.16.0 built against HDF5 2.1.0 with HDF5 2.2.0. Current main now pins hdf5 >=2.1,<2.2, addressing that incompatibility. The new Tests and Examples workflow runs are awaiting maintainer approval.