# 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`.

```python
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:

```python
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:

```python
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 {doc}`developer benchmark harness </user_guide/benchmarking>`.

## Examples and reference

{doc}`Batching tutorial </tutorial/01_tmol_how_to_guide>` · {doc}`Pose API </api/pose>` · {doc}`CPU threading </user_guide/cpu_threading>`