Skip to content

Get rid of old multi-kernel DG and ETK energy approach - #378

Merged
scal444 merged 2 commits into
NVIDIA-BioNeMo:mainfrom
scal444:perf-campaign-dg_etk-full_energy
Oct 26, 2026
Merged

scal444 merged 2 commits into
NVIDIA-BioNeMo:mainfrom
scal444:perf-campaign-dg_etk-full_energy

Conversation

@scal444

@scal444 scal444 commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator

The block-per-mol approach that the e2e per_mol backend has is better, even for large mols that don't go through the full per_mol backend.

The double-precision DG energy issued a scratch memset, one term-parallel
launch per term family and a reduction over the per-term energy buffer;
the ETK energy did the same for its six term families. Evaluate both with
the existing block-per-molecule energy kernels instead (one launch, one
block reduction per molecule), which SINGLE precision and the per-molecule
path already use. On sm_89 this takes the DG energy from about 1.18 to
0.90 ms per call and the ETK energy from about 1.05 to 0.71 ms per call on
ChEMBL 60-80 batches of 1000 conformers.

The block-per-molecule kernel writes each active system's energy and zero
for inactive systems; every caller zeroes the outputs first, so results
match up to summation order. ETK term subsets other than PLAIN without
improper torsions keep the term-parallel kernels. The DG forcefield no
longer allocates the per-term energy scratch buffer.

The FULL gradients keep the term-parallel kernels: the existing block-per-
molecule DG gradient measured 25% slower per call (2.19 vs 1.75 ms).
@greptile-apps

greptile-apps Bot commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The reviewed changes appear safe to merge; no blocking issues were found.

Summary

The PR routes DG energy evaluation and eligible ETK evaluations through one block per molecule.

  • DG batches calculate each molecule’s energy in one CUDA block.
  • Full ETK energy calls use one block per molecule when their terms fit.

No actionable new issues were found in the changes since the previous review.

Reviews (2) · Last reviewed commit: "Remove unused per-term DG energy kernels..." · Reviewed by Greptile

@scal444
scal444 requested a review from evasnow1992 October 23, 2026 16:13
Comment thread src/forcefields/dist_geom.h Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The docstring still says both energyOuts and energyBuffer must be zeroed. With the new ALL/PLAIN fast path, energyOuts is overwritten directly and energyBuffer is no longer used, while the other terms still require zero initialization. May consider updating the documentation to clarify these two behaviors.

@evasnow1992

Copy link
Copy Markdown
Collaborator

It looks like the old per-term DG energy kernels and related energy-buffer indexing logic are no longer used in production, but are still referenced by test_distgeom_kernels.cu. Is the intention to keep them as reference implementations for testing? If so, might be worth documenting; otherwise, we could consider removing the unused code.

@evasnow1992 evasnow1992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two comments, none blocker. Changes look good to me.

The DG energy path now runs only the block-per-molecule kernel, so drop the
per-term DG energy launchers, the DG energy buffer and its index arrays, the
unused DG allocateIntermediateBuffers and computeEnergyBlockPerMol, and the
host argument setupDeviceBuffers no longer reads. DG per-term energy tests now
go through computeEnergy, isolating terms with the chiral and fourth-dimension
weights. Update the computeEnergyETK docstring for the overwrite and
accumulate paths.
@scal444
scal444 merged commit 1876a73 into NVIDIA-BioNeMo:main Oct 26, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants