GPU batching#
Combine prepared structures in one PoseStack and score the batch in one
call. The inputs in poses and the score function sfxn must use the same
parameter database and target device.
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.