GPU batching#

Use this compact recipe when an application needs to score many structures on one device. The linked tutorial provides the deeper walkthrough, measurements, and interpretation.

TMol gets application throughput by evaluating several structures in one PoseStack on one device. Build each input pose from the same chemical database, combine them, and render the scorer for the resulting batch layout:

from time import perf_counter

import torch

from tmol.pose import PoseStackBuilder

batch = PoseStackBuilder.from_poses(poses, device)
scorer = sfxn.render_whole_pose_scoring_module(batch)

with torch.no_grad():
    scores = scorer(batch.coords)

from_poses() accepts heterogeneous sizes but pads the batch to its largest member. Group similarly sized structures when practical, and use a score function built from the same parameter database as poses with custom residue or ligand types.

Time CUDA work correctly#

CUDA launches are asynchronous. Warm up the rendered module, synchronize on both sides of each timed call, and report multiple repeats:

for _ in range(warmup):
    scorer(batch.coords)
torch.cuda.synchronize(device)

elapsed = []
for _ in range(repeats):
    torch.cuda.synchronize(device)
    start = perf_counter()
    scorer(batch.coords)
    torch.cuda.synchronize(device)
    elapsed.append(perf_counter() - start)

Measure total batch latency and throughput separately. First-call compilation time should not be mixed into steady-state kernel timing.

For repeated default weighted scoring with a fixed coordinate shape, dtype, and device, capture the launches once and replay them:

graphed_scorer = sfxn.render_whole_pose_scoring_module(
    batch, cuda_graph="forward"
)
with torch.no_grad():
    scores = graphed_scorer(batch.coords)

Forward-only replay reuses its output buffer. Clone scores when it must remain unchanged across the next call. CUDA graphs mainly reduce launch overhead; they do not replace batching or improve kernels that already dominate runtime.

Chunk larger workloads#

Batch size is limited by padding and operation-specific intermediates. Choose a conservative chunk size, render and score one chunk, move only the needed results off device, and then release that chunk before continuing. Monitor both the live CUDA allocation after construction and the incremental peak during the operation.

TMol does not provide a multi-GPU scheduler. Shard independent chunks outside TMol, normally with one process owning one GPU. This application workflow is separate from the developer benchmark harness.