Tutorial 02 — GPU batching#

Build repeated and mixed-size batches, score them in one call, and measure throughput and memory. Start with structure I/O. Small examples run on CPU; GPU measurements require CUDA. Each PoseStack uses one device.

Setup#

In Colab, select T4 GPU, then Run all. Locally, follow the installation guide. Setup downloads the fixtures and fixes random seeds.

Timing uses 1UBQ chain A residues 1–20 with no_optH=True. Slicing creates a new terminus; hydrogens remain at ideal kinematic positions. This fragment keeps the CPU example small.

The heterogeneous example uses full model-1 structures with no_optH=False, optimizing hydrogens/NHQ flips and rebuilding supported missing atoms. Its table records preparation and deposited-versus-built atom counts.

[1]:

try:
    import google.colab  # noqa: F401
except ImportError:
    IN_COLAB = False
else:
    IN_COLAB = True

if IN_COLAB:
    from urllib.request import urlopen

    exec(
        urlopen(
            "https://raw.githubusercontent.com/uw-ipd/tmol/"
            "master/docs/tutorial/colab_setup.py"
        ).read(),
        globals(),
    )
    setup_colab(
        [
            "tmol/tests/data/cif/1UBQ.cif",
            "tmol/tests/data/cif/1R21.cif",
            "tmol/tests/data/cif/1BL8.cif",
        ]
    )

[2]:

from contextlib import redirect_stderr, redirect_stdout
from io import StringIO
from pathlib import Path
import platform
from time import perf_counter
import warnings

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from IPython.display import display
from biotite.structure import AtomArrayStack
from biotite.structure.io import load_structure

import tmol
from tmol.io import biotite_from_pose_stack, pose_stack_from_biotite
from tmol.pose import PoseStackBuilder
from tmol.score import beta2016_score_function

SEED = 20260807
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

device = (
    torch.device("cuda", torch.cuda.current_device())
    if torch.cuda.is_available()
    else torch.device("cpu")
)
repo_root = Path.cwd()
if not (repo_root / "tmol/tests/data/cif/1UBQ.cif").exists():
    repo_root = Path(tmol.__file__).resolve().parents[1]
cif_path = repo_root / "tmol/tests/data/cif/1UBQ.cif"

atom_array = load_structure(str(cif_path), model=1, include_bonds=True)
protein_slice = atom_array[(atom_array.chain_id == "A") & (atom_array.res_id <= 20)]
pose_diagnostics = StringIO()
try:
    with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):
        single_pose = pose_stack_from_biotite(protein_slice, device, no_optH=True)
except Exception:
    print(pose_diagnostics.getvalue())
    raise
score_function = beta2016_score_function(device)


def show_table(frame):
    """Use sortable tables in rendered docs, with a pandas fallback."""
    try:
        from itables import show
    except ImportError:
        return display(frame)
    return show(frame)


print(f"benchmark device: {device}; input: {cif_path.name}")

Environment variable CCD_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
  (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
  (2) set it in your current shell with 'export VAR_NAME=path/to/variable'
  (3) write it to a .env file in the root of the atomworks.io repository
Environment variable PDB_MIRROR_PATH not set. Will not be able to use function requiring this variable. To set it you may:
  (1) add the line 'export VAR_NAME=path/to/variable' to your .bashrc or .zshrc file
  (2) set it in your current shell with 'export VAR_NAME=path/to/variable'
  (3) write it to a .env file in the root of the atomworks.io repository

benchmark device: cpu; input: 1UBQ.cif

Record benchmark metadata first#

Latency numbers are not portable without the PyTorch/TMol versions, operating system, CUDA runtime, and GPU identity. Record these before timing. Peak memory below is CUDA allocator memory, not total process or driver memory.

[3]:

metadata = {
    "platform": platform.platform(),
    "python": platform.python_version(),
    "tmol": tmol.__version__,
    "torch": torch.__version__,
    "device": str(device),
    "cuda_runtime": torch.version.cuda,
}
if device.type == "cuda":
    props = torch.cuda.get_device_properties(device)
    metadata.update(
        {
            "gpu_name": props.name,
            "compute_capability": f"{props.major}.{props.minor}",
            "gpu_memory_GiB": props.total_memory / 2**30,
        }
    )
metadata_frame = pd.DataFrame.from_dict(metadata, orient="index", columns=["value"])
show_table(metadata_frame.reset_index(names="property"))

^ⓘproperty	value
platform	Linux-6.17.0-1022-azure-x86_64-with-glibc2.39
python	3.12.14
tmol	0.1.62
torch	2.14.1+cpu
device	cpu
cuda_runtime	NaN

Build batches explicitly#

PoseStackBuilder.from_poses concatenates compatible pose stacks and repacks their tensors on the requested device. Repeating the same pose is useful for a controlled throughput demonstration; real workloads normally batch distinct but chemistry-compatible structures.

[4]:

preview_batch = PoseStackBuilder.from_poses([single_pose] * 4, device)
print("single coords:", tuple(single_pose.coords.shape))
print("batch coords: ", tuple(preview_batch.coords.shape))
print("batch poses:  ", preview_batch.n_poses)

single coords: (1, 327, 3)
batch coords:  (4, 327, 3)
batch poses:   4

Score a heterogeneous CIF batch#

Prepare model 1 from three checked-in CIF entries with optimized hydrogens, combine them, and score the batch. The selector labels each input; 1R21 is an NMR entry, so its model count is reported separately.

Batches pad atom and block dimensions to the largest member. Group similarly sized structures to reduce unused slots.

The preparation table reports net atom-count changes, not per-atom provenance. In 1BL8, TMol excludes an unrecognized deposited K ion. Conversion may also select chemical variants and build missing atoms.

Assertions compare each batched total with an independently rendered scorer. These unrelated proteins differ in size and composition: their absolute score totals demonstrate batch indexing, not a meaningful ranking.

[5]:

structure_specs = [
    ("1UBQ — ubiquitin", "1UBQ.cif"),
    ("1R21 — model 1", "1R21.cif"),
    ("1BL8", "1BL8.cif"),
]
individual_poses = {}
preparation_records = {}
for label, filename in structure_specs:
    structure_file = load_structure(
        str(repo_root / "tmol" / "tests" / "data" / "cif" / filename),
        include_bonds=True,
    )
    if isinstance(structure_file, AtomArrayStack):
        deposited_model_count = structure_file.stack_depth()
        structure = structure_file[0]
    else:
        deposited_model_count = 1
        structure = structure_file

    pose_diagnostics = StringIO()
    try:
        with redirect_stdout(pose_diagnostics), redirect_stderr(pose_diagnostics):
            pose = pose_stack_from_biotite(structure, device, no_optH=False)
    except Exception:
        print(pose_diagnostics.getvalue())
        raise
    individual_poses[label] = pose
    tmol_built_structure = biotite_from_pose_stack(pose)
    preparation_records[label] = {
        "deposited_models": deposited_model_count,
        "selected_model": 1,
        "deposited_atoms": structure.array_length(),
        "TMol_built_atoms": tmol_built_structure.array_length(),
        "net_atom_count_change": (
            tmol_built_structure.array_length() - structure.array_length()
        ),
        "count_interpretation": (
            "net filtering plus selected/built chemistry; no per-atom provenance"
        ),
        "preparation": "no_optH=False (OptH optimized)",
        "known_excluded_material": "K ion" if filename == "1BL8.cif" else "",
    }

heterogeneous_batch = PoseStackBuilder.from_poses(
    list(individual_poses.values()), device
)
heterogeneous_scorer = score_function.render_whole_pose_scoring_module(
    heterogeneous_batch
)
with warnings.catch_warnings(), torch.no_grad():
    warnings.filterwarnings("ignore", message=r"Sparse index lookup.*")
    heterogeneous_scores = heterogeneous_scorer(
        heterogeneous_batch.coords
    ).detach().cpu().numpy()

independent_scores = []
with warnings.catch_warnings(), torch.no_grad():
    warnings.filterwarnings("ignore", message=r"Sparse index lookup.*")
    for pose in individual_poses.values():
        independent_scorer = score_function.render_whole_pose_scoring_module(pose)
        independent_score = independent_scorer(pose.coords).detach().cpu().numpy()
        independent_scores.append(float(independent_score.reshape(-1)[0]))
independent_scores = np.asarray(independent_scores)
np.testing.assert_allclose(
    heterogeneous_scores,
    independent_scores,
    rtol=1e-5,
    atol=1e-3,
)

batch_rows = []
for pose_index, (label, pose) in enumerate(individual_poses.items()):
    preparation = preparation_records[label]
    batch_rows.append(
        {
            "pose_index": pose_index,
            "structure": label,
            **preparation,
            "blocks": pose.max_n_blocks,
            "pose_atoms": pose.max_n_pose_atoms,
            "batched_weighted_score": float(heterogeneous_scores[pose_index]),
            "independent_weighted_score": float(independent_scores[pose_index]),
            "absolute_score_difference": float(
                abs(heterogeneous_scores[pose_index] - independent_scores[pose_index])
            ),
        }
    )
batch_frame = pd.DataFrame(batch_rows)
show_table(batch_frame)
print(
    "maximum batched-vs-independent absolute difference:",
    f"{batch_frame['absolute_score_difference'].max():.6g}",
)
display(
    tmol.switchable_view(
        individual_poses,
        notes={
            row["structure"]: (
                f"batch index {row['pose_index']}; model {row['selected_model']} of "
                f"{row['deposited_models']}; {row['blocks']} blocks; "
                f"score {row['batched_weighted_score']:.3f}; optimized H"
            )
            for row in batch_rows
        },
    )
)

^ⓘpose_index	structure	deposited_models	selected_model	deposited_atoms	TMol_built_atoms	net_atom_count_change	count_interpretation	preparation	known_excluded_material	blocks	pose_atoms	batched_weighted_score	independent_weighted_score
0	1UBQ — ubiquitin	1	1	660	1231	571	net filtering plus selected/built chemistry; no per-atom provenance	no_optH=False (OptH optimized)		76	1231	201.508240	201.508240
1	1R21 — model 1	23	1	1612	1615	3	net filtering plus selected/built chemistry; no per-atom provenance	no_optH=False (OptH optimized)		100	1615	1041.679932	1041.679932
2	1BL8	1	1	2824	5891	3067	net filtering plus selected/built chemistry; no per-atom provenance	no_optH=False (OptH optimized)	K ion	391	5891	3663.375244	3663.375244

maximum batched-vs-independent absolute difference: 0

Structure:

batch index 0; model 1 of 1; 76 blocks; score 201.508; optimized H

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

batch index 1; model 1 of 23; 100 blocks; score 1041.680; optimized H

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

batch index 2; model 1 of 1; 391 blocks; score 3663.375; optimized H

3Dmol.js failed to load for some reason. Please check your browser console for error messages.

Compare mixed and size-bucketed batches#

Repeat each prepared structure twice and score the same six poses as one mixed batch and three two-pose buckets. Render scorers before timing and verify scores in the same order.

The table reports padding, throughput, and incremental CUDA allocation above the post-warmup baseline. Bucketing reduces padding but adds kernel launches; either layout may be faster.

[6]:

workload_labels = [label for label in individual_poses for _ in range(2)]
workload_poses = [individual_poses[label] for label in workload_labels]
layout_groups = {
    "one heterogeneous batch": [workload_poses],
    "three size buckets": [
        [individual_poses[label], individual_poses[label]] for label in individual_poses
    ],
}


def build_layout(groups):
    """Build all batches and rendered scorers for one workload layout."""
    batches = [PoseStackBuilder.from_poses(group, device) for group in groups]
    scorers = [
        score_function.render_whole_pose_scoring_module(batch) for batch in batches
    ]
    return batches, scorers


def run_layout(batches, scorers):
    """Score each pre-rendered group and concatenate in workload order."""
    return torch.cat(
        [scorer(batch.coords) for batch, scorer in zip(batches, scorers, strict=True)]
    )


def time_layout(batches, scorers, repeats, warmup):
    """Measure aggregate latency for all calls in one workload layout."""
    with torch.no_grad():
        for _ in range(warmup):
            run_layout(batches, scorers)
        if device.type == "cuda":
            torch.cuda.synchronize(device)
            baseline_bytes = torch.cuda.memory_allocated(device)
            torch.cuda.reset_peak_memory_stats(device)
        else:
            baseline_bytes = np.nan
        elapsed = []
        for _ in range(repeats):
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            start_time = perf_counter()
            run_layout(batches, scorers)
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            elapsed.append(perf_counter() - start_time)
    if device.type == "cuda":
        incremental_peak_bytes = max(
            0, torch.cuda.max_memory_allocated(device) - baseline_bytes
        )
    else:
        incremental_peak_bytes = np.nan
    return 1e3 * np.asarray(elapsed), incremental_peak_bytes


built_layouts = {label: build_layout(groups) for label, groups in layout_groups.items()}
with torch.no_grad():
    layout_scores = {
        label: run_layout(*layout).detach().cpu()
        for label, layout in built_layouts.items()
    }
torch.testing.assert_close(
    layout_scores["one heterogeneous batch"],
    layout_scores["three size buckets"],
    rtol=1e-5,
    atol=1e-3,
)

useful_atoms = sum(int(pose.real_atoms.sum().item()) for pose in workload_poses)
useful_blocks = sum(
    int((pose.block_type_ind64 >= 0).sum().item()) for pose in workload_poses
)
layout_rows = []
layout_repeats = 7 if device.type == "cuda" else 2
layout_warmup = 3 if device.type == "cuda" else 1
for label, (batches, scorers) in built_layouts.items():
    elapsed_ms, incremental_peak_bytes = time_layout(
        batches,
        scorers,
        repeats=layout_repeats,
        warmup=layout_warmup,
    )
    atom_slots = sum(batch.n_poses * batch.max_n_pose_atoms for batch in batches)
    block_slots = sum(batch.n_poses * batch.max_n_blocks for batch in batches)
    layout_rows.append(
        {
            "layout": label,
            "score_calls_per_workload": len(batches),
            "poses": len(workload_poses),
            "useful_atom_slots": useful_atoms,
            "rectangular_atom_slots": atom_slots,
            "atom_padding_fraction": 1.0 - useful_atoms / atom_slots,
            "useful_block_slots": useful_blocks,
            "rectangular_block_slots": block_slots,
            "block_padding_fraction": 1.0 - useful_blocks / block_slots,
            "median_total_latency_ms": float(np.median(elapsed_ms)),
            "throughput_poses_s": float(
                1e3 * len(workload_poses) / np.median(elapsed_ms)
            ),
            "latency_IQR_ms": float(
                np.percentile(elapsed_ms, 75) - np.percentile(elapsed_ms, 25)
            ),
            "incremental_peak_cuda_MiB": incremental_peak_bytes / 2**20,
            "measurement": (
                "CUDA workload timing"
                if device.type == "cuda"
                else "CPU smoke timing; not GPU guidance"
            ),
        }
    )
layout_frame = pd.DataFrame(layout_rows)
show_table(layout_frame)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
padding_plot = layout_frame.set_index("layout")[
    ["atom_padding_fraction", "block_padding_fraction"]
]
padding_plot.plot.bar(ax=axes[0], rot=10)
axes[0].set(
    ylabel="fraction of rectangular slots that are padding",
    title="Bucketing removes avoidable padding",
)
axes[1].bar(
    layout_frame["layout"],
    layout_frame["median_total_latency_ms"],
    color=["#3b82f6", "#f59e0b"],
)
axes[1].set(
    ylabel="median total workload latency (ms)",
    title=(
        "Measured CUDA workload"
        if device.type == "cuda"
        else "CPU smoke timing — not GPU guidance"
    ),
)
axes[1].tick_params(axis="x", rotation=10)
for axis in axes:
    axis.grid(axis="y", alpha=0.3)
plt.tight_layout()
plt.show()

^ⓘlayout	score_calls_per_workload	poses	useful_atom_slots	rectangular_atom_slots	atom_padding_fraction	useful_block_slots	rectangular_block_slots	block_padding_fraction	median_total_latency_ms	throughput_poses_s	latency_IQR_ms	incremental_peak_cuda_MiB	measurement
one heterogeneous batch	1	6	17474	35346	0.50563	1134	2346	0.516624	99.923915	60.045686	1.693125	NaN	CPU smoke timing; not GPU guidance
three size buckets	3	6	17474	17474	0.00000	1134	1134	0.000000	87.498579	68.572542	2.400695	NaN	CPU smoke timing; not GPU guidance

../_images/tutorial_02_gpu_batching_11_1.png

Choose bucket boundaries from measured throughput and memory, not padding alone. CUDA peaks cover PyTorch allocations, excluding driver/process memory. CPU rows check execution and do not measure GPU scaling.

Benchmark methodology#

render_whole_pose_scoring_module(batch) creates a scorer for that batch’s fixed block, atom, and connectivity layout. Reuse it while only coordinates change within the same layout. Re-render when batch size, membership, padding dimensions, or chemical layout changes; benchmark_scoring() therefore renders once for each batch size and excludes that construction from the timed region.

CUDA launches are asynchronous. A valid wall-clock measurement therefore:

renders the scorer before timing and records the post-construction allocator baseline;
resets peak-memory statistics relative to that live baseline;
runs untimed warmup calls so lazy compilation and caches are not charged to steady-state latency;
synchronizes before starting and after finishing each timed call; and
reports multiple repeats rather than one launch.

The table’s latency_ms is median total latency for one batch call and is accompanied by Q1, Q3, IQR, minimum, and maximum. latency_per_pose_ms divides the median total by batch size and is an amortized throughput metric, not the time at which one pose’s result becomes independently available. Plot error bars span Q1–Q3. CUDA memory reports the post-construction live allocation, total observed peak, and the incremental peak above that baseline. The CPU path uses the same timing function without CUDA synchronization and deliberately tiny sizes.

[7]:

def benchmark_scoring(batch_size, repeats=5, warmup=2):
    batch = PoseStackBuilder.from_poses([single_pose] * batch_size, device)
    scorer = score_function.render_whole_pose_scoring_module(batch)

    if device.type == "cuda":
        torch.cuda.synchronize(device)
        post_construction_bytes = torch.cuda.memory_allocated(device)
        torch.cuda.reset_peak_memory_stats(device)
    else:
        post_construction_bytes = np.nan

    with torch.no_grad():
        for _ in range(warmup):
            scorer(batch.coords)
        if device.type == "cuda":
            torch.cuda.synchronize(device)

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

    elapsed_ms = 1e3 * np.asarray(elapsed)
    latency_q1_ms, latency_q3_ms = np.percentile(elapsed_ms, [25, 75])
    latency_ms = float(np.median(elapsed_ms))
    if device.type == "cuda":
        peak_bytes = torch.cuda.max_memory_allocated(device)
        incremental_peak_bytes = max(0, peak_bytes - post_construction_bytes)
    else:
        peak_bytes = np.nan
        incremental_peak_bytes = np.nan
    return {
        "batch_size": batch_size,
        "latency_min_ms": float(elapsed_ms.min()),
        "latency_q1_ms": float(latency_q1_ms),
        "latency_ms": latency_ms,
        "latency_q3_ms": float(latency_q3_ms),
        "latency_max_ms": float(elapsed_ms.max()),
        "latency_iqr_ms": float(latency_q3_ms - latency_q1_ms),
        "throughput_poses_s": 1e3 * batch_size / latency_ms,
        "post_construction_cuda_MiB": post_construction_bytes / 2**20,
        "peak_cuda_MiB": peak_bytes / 2**20,
        "incremental_peak_cuda_MiB": incremental_peak_bytes / 2**20,
        "repeats": repeats,
    }


batch_sizes = [1, 4, 16, 64] if device.type == "cuda" else [1, 2, 4]
repeats = 5 if device.type == "cuda" else 2
benchmark_frame = pd.DataFrame(
    [benchmark_scoring(size, repeats=repeats) for size in batch_sizes]
)
benchmark_frame["latency_per_pose_ms"] = (
    benchmark_frame["latency_ms"] / benchmark_frame["batch_size"]
)
benchmark_frame["latency_per_pose_q1_ms"] = (
    benchmark_frame["latency_q1_ms"] / benchmark_frame["batch_size"]
)
benchmark_frame["latency_per_pose_q3_ms"] = (
    benchmark_frame["latency_q3_ms"] / benchmark_frame["batch_size"]
)
benchmark_frame["throughput_q1_poses_s"] = (
    1e3 * benchmark_frame["batch_size"] / benchmark_frame["latency_q3_ms"]
)
benchmark_frame["throughput_q3_poses_s"] = (
    1e3 * benchmark_frame["batch_size"] / benchmark_frame["latency_q1_ms"]
)
benchmark_frame["throughput_vs_batch_1"] = (
    benchmark_frame["throughput_poses_s"]
    / benchmark_frame.loc[0, "throughput_poses_s"]
)
benchmark_frame.insert(
    0, "measurement", "CUDA throughput" if device.type == "cuda" else "CPU smoke check"
)
show_table(benchmark_frame)

^ⓘmeasurement	batch_size	latency_min_ms	latency_q1_ms	latency_ms	latency_q3_ms	latency_max_ms	latency_iqr_ms	throughput_poses_s	post_construction_cuda_MiB	peak_cuda_MiB	incremental_peak_cuda_MiB	repeats	latency_per_pose_ms	latency_per_pose_q1_ms	latency_per_pose_q3_ms	throughput_q1_poses_s	throughput_q3_poses_s	throughput_vs_batch_1
CPU smoke check	1	2.353236	2.393311	2.433386	2.473461	2.513536	0.080150	410.950010	NaN	NaN	NaN	2	2.433386	2.393311	2.473461	404.291800	417.831197	1.000000
CPU smoke check	2	2.963070	2.996011	3.028952	3.061893	3.094834	0.065882	660.294386	NaN	NaN	NaN	2	1.514476	1.498006	1.530947	653.190690	667.554291	1.606751
CPU smoke check	4	4.250931	4.256515	4.262099	4.267683	4.273267	0.011168	938.504713	NaN	NaN	NaN	2	1.065525	1.064129	1.066921	937.276738	939.735911	2.283744

[8]:

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
latency_yerr = np.vstack(
    [
        benchmark_frame["latency_per_pose_ms"]
        - benchmark_frame["latency_per_pose_q1_ms"],
        benchmark_frame["latency_per_pose_q3_ms"]
        - benchmark_frame["latency_per_pose_ms"],
    ]
)
axes[0].errorbar(
    benchmark_frame["batch_size"],
    benchmark_frame["latency_per_pose_ms"],
    yerr=latency_yerr,
    marker="o",
    capsize=4,
)
axes[0].set(
    xlabel="batch size",
    ylabel="median milliseconds / pose",
    title="Amortized latency (Q1–Q3)",
)
throughput_yerr = np.vstack(
    [
        benchmark_frame["throughput_poses_s"]
        - benchmark_frame["throughput_q1_poses_s"],
        benchmark_frame["throughput_q3_poses_s"]
        - benchmark_frame["throughput_poses_s"],
    ]
)
axes[1].errorbar(
    benchmark_frame["batch_size"],
    benchmark_frame["throughput_poses_s"],
    yerr=throughput_yerr,
    marker="o",
    capsize=4,
)
axes[1].set(
    xlabel="batch size",
    ylabel="poses / second",
    title="Scoring throughput (Q1–Q3)",
)
for axis in axes:
    axis.grid(alpha=0.3)
fig.suptitle(
    "Measured CUDA batching" if device.type == "cuda" else "CPU smoke check — not GPU scaling",
    fontweight="bold",
)
plt.tight_layout()
plt.show()

../_images/tutorial_02_gpu_batching_15_0.png

Batching can improve GPU utilization: call latency grows more slowly than batch size until compute or memory bandwidth saturates. Per-pose latency is call latency divided by batch size; the underlying scoring work still grows with the number of poses.

CUDA measurements depend on hardware, software, structure sizes, and competing work. Rerun for your workload. Warmup excludes first-call compilation.

CUDA capacity probe#

The gpu-only cell measures a larger batch. GPU CI executes it; CPU documentation builds retain its checked-in output.

[ ]:

#| tags: [gpu-only]
# gpu-only: this capacity point is omitted on CPU.
if device.type != "cuda":
    print("Skipped: this cell requires CUDA.")
    gpu_capacity_point = None
else:
    gpu_capacity_point = benchmark_scoring(128, repeats=5, warmup=2)
    show_table(pd.DataFrame([gpu_capacity_point]))
gpu_capacity_point

Memory and chunking#

Padding scales with the largest atom and block counts; pairwise intermediates can add further memory costs.

post_construction_cuda_MiB: live PyTorch allocation after batch and scorer construction.
peak_cuda_MiB: peak allocation during warmup and timed calls.
incremental_peak_cuda_MiB: peak minus the post-construction allocation.

These exclude driver/process memory and are not per-pose quantities. Shared tensors and padding make division by batch size misleading.

Leave headroom for compilation and downstream work. Score conservative chunks, detach or transfer the needed results, then release each chunk. Group similar sizes to limit padding.

For multiple GPUs, use an external scheduler such as Slurm or Dask, normally with one process per GPU. TMol has no multi-GPU scheduler.

Next#

Use scoring and analysis to inspect terms and block-pair contributions.

Exercises#

Increase the repeat count and compare IQR with the full min–max timing range.
Increase CUDA batch sizes until throughput plateaus, stopping well before memory exhaustion.
Repeat the executed bucketing experiment with at least two members per size range, then choose bucket boundaries from measured memory and throughput rather than residue count alone.
Implement an outer loop that scores a large list in conservative chunks and concatenates detached totals.
Design a Slurm array where each task selects one GPU and processes many TMol batches; keep random seeds and metadata per task.

References#

Download this notebook