From 457265c1ffe6ff726a413a7a248d46dcfc447c33 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 29 Jul 2026 16:27:12 -0400 Subject: [PATCH 1/4] Widen substructure autotune search range, sanitize SMILES by default The batchSize search range was (128, 1024, step=128), and tuning runs were pinning to 1024 -- the ceiling of the range, not an interior optimum. Raise the ceiling to 8192, matching what benchmarks/substruct_config.csv already uses, and switch to log spacing since the useful range spans two orders of magnitude. The step=128 bought nothing: the runtime derives targetsPerBatch = batchSize // numQueries, and that truncation destroys any 128 alignment before it reaches a kernel. Default calibration_max_size to None (uncapped) for tune_substructure. The 2000-target cap left too few mini-batches to fill every worker's executor ring, so pipeline fill and drain dominated the measurement and large batchSize values were systematically under-rewarded. auto_subsample and normalize_calibration_set now accept max_size=None; other tuners keep the 2000 default. Also default substruct_bench.py to --sanitize. It was the only benchmark defaulting to sanitize=False, while etkdg_bench and ff_optimize_bench both default to True. Move the flag pair into a shared bench_utils.cli helper so the mutually-exclusive group and the safe default stay consistent. --- benchmarks/bench_utils/__init__.py | 2 ++ benchmarks/bench_utils/cli.py | 35 ++++++++++++++++++++++++ benchmarks/substruct_bench.py | 6 +--- benchmarks/tests/test_substruct_bench.py | 25 +++++++++++++++++ nvmolkit/autotune/_calibration.py | 14 ++++++---- nvmolkit/autotune/tune_substructure.py | 12 ++++---- 6 files changed, 79 insertions(+), 15 deletions(-) create mode 100644 benchmarks/bench_utils/cli.py create mode 100644 benchmarks/tests/test_substruct_bench.py diff --git a/benchmarks/bench_utils/__init__.py b/benchmarks/bench_utils/__init__.py index 4cabcc21..74366b35 100644 --- a/benchmarks/bench_utils/__init__.py +++ b/benchmarks/bench_utils/__init__.py @@ -19,6 +19,7 @@ so individual bench scripts can ``from bench_utils import time_it, load_smiles``. """ +from bench_utils.cli import add_smiles_sanitization_args from bench_utils.loaders import load_pickle, load_sdf, load_smarts, load_smiles from bench_utils.molprep import clone_mols_with_conformers, embed_and_jitter, perturb_conformer, prep_mols from bench_utils.timing import ( @@ -34,6 +35,7 @@ "Deadline", "TimingResult", "add_rdkit_max_seconds_arg", + "add_smiles_sanitization_args", "clone_mols_with_conformers", "embed_and_jitter", "load_pickle", diff --git a/benchmarks/bench_utils/cli.py b/benchmarks/bench_utils/cli.py new file mode 100644 index 00000000..106a9d04 --- /dev/null +++ b/benchmarks/bench_utils/cli.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared command-line argument helpers for nvMolKit benchmarks.""" + +import argparse + +def add_smiles_sanitization_args(parser: argparse.ArgumentParser) -> None: + """Add safe-by-default SMILES sanitization flags to ``parser``.""" + sanitize_group = parser.add_mutually_exclusive_group() + sanitize_group.add_argument( + "--sanitize", + action="store_true", + dest="sanitize", + default=True, + help="Sanitize SMILES during parsing (default)", + ) + sanitize_group.add_argument( + "--no_sanitize", + action="store_false", + dest="sanitize", + help="Skip sanitization (preprocessed SMILES only)", + ) diff --git a/benchmarks/substruct_bench.py b/benchmarks/substruct_bench.py index a3424318..78b9c2e4 100644 --- a/benchmarks/substruct_bench.py +++ b/benchmarks/substruct_bench.py @@ -335,11 +335,7 @@ def main(): default=42, help="Random seed for sampling SMILES when --num_mols > 0 (default: 42)", ) - parser.add_argument("--sanitize", action="store_true", dest="sanitize", help="Sanitize SMILES during parsing") - parser.add_argument( - "--no_sanitize", action="store_false", dest="sanitize", help="Skip sanitization (preprocessed SMILES)" - ) - parser.set_defaults(sanitize=False) + add_smiles_sanitization_args(parser) parser.add_argument("--runs", "-r", type=int, default=1, help="Number of timing runs (default: 1)") parser.add_argument( "--mode", diff --git a/benchmarks/tests/test_substruct_bench.py b/benchmarks/tests/test_substruct_bench.py new file mode 100644 index 00000000..9f055ae7 --- /dev/null +++ b/benchmarks/tests/test_substruct_bench.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the substructure benchmark CLI.""" + +import argparse + +import pytest +from bench_utils.cli import add_smiles_sanitization_args + +def test_smiles_sanitization_is_enabled_by_default(): + parser = argparse.ArgumentParser() + add_smiles_sanitization_args(parser) + + assert parser.parse_args([]).sanitize is True + assert parser.parse_args(["--sanitize"]).sanitize is True + assert parser.parse_args(["--no_sanitize"]).sanitize is False + + +def test_smiles_sanitization_flags_are_mutually_exclusive(): + parser = argparse.ArgumentParser() + add_smiles_sanitization_args(parser) + + with pytest.raises(SystemExit): + parser.parse_args(["--sanitize", "--no_sanitize"]) diff --git a/nvmolkit/autotune/_calibration.py b/nvmolkit/autotune/_calibration.py index e4a1b5af..a7782dea 100644 --- a/nvmolkit/autotune/_calibration.py +++ b/nvmolkit/autotune/_calibration.py @@ -25,7 +25,7 @@ def auto_subsample( workload_size: int, *, fraction: float = 0.1, - max_size: int = 2000, + max_size: Optional[int] = 2000, min_size: int = 1, seed: int = 0, ) -> list[int]: @@ -33,11 +33,14 @@ def auto_subsample( The slice size is ``min(max_size, max(min_size, round(fraction * workload_size)))`` and is shuffled deterministically with ``seed`` so trials sample roughly - uniformly across the workload. + uniformly across the workload. Passing ``max_size=None`` disables the cap, + leaving ``fraction`` as the only limit. """ if workload_size <= 0: raise ValueError("workload_size must be positive") - target = min(max_size, max(min_size, int(round(fraction * workload_size)))) + target = max(min_size, int(round(fraction * workload_size))) + if max_size is not None: + target = min(max_size, target) target = min(target, workload_size) rng = random.Random(seed) indices = list(range(workload_size)) @@ -50,7 +53,7 @@ def normalize_calibration_set( workload_size: int, *, fraction: float = 0.1, - max_size: int = 2000, + max_size: Optional[int] = 2000, seed: int = 0, ) -> list[int]: """Return calibration indices from an explicit set or auto-sample. @@ -63,7 +66,8 @@ def normalize_calibration_set( auto-subsample. workload_size: Total size of the user workload. fraction: Auto-subsample fraction when ``calibration_set`` is ``None``. - max_size: Auto-subsample cap when ``calibration_set`` is ``None``. + max_size: Auto-subsample cap when ``calibration_set`` is ``None``, or + ``None`` for no cap. seed: Seed for auto-subsampling. """ if calibration_set is None: diff --git a/nvmolkit/autotune/tune_substructure.py b/nvmolkit/autotune/tune_substructure.py index 49418688..e864ae65 100644 --- a/nvmolkit/autotune/tune_substructure.py +++ b/nvmolkit/autotune/tune_substructure.py @@ -64,8 +64,8 @@ def _default_substruct_search_space(num_gpus: int, cpus: int) -> dict: at trial sampling time by clamping ``preprocessingThreads`` to whatever cores are left after ``workerThreads`` is chosen. - * ``batchSize``: stepped int range in multiples of 128 (kernels are - tile-tuned for these sizes); stepping preserves numeric ordering for TPE. + * ``batchSize``: target-query pairs per mini-batch, log-uniform since the + useful range spans two orders of magnitude. * ``workerThreads`` (per-GPU): max = ``min(8, cpus // num_gpus)``. 8 is the empirical point of diminishing returns; the physical-core floor prevents oversubscribing across GPUs. @@ -74,7 +74,7 @@ def _default_substruct_search_space(num_gpus: int, cpus: int) -> dict: """ per_gpu_worker_max = max(1, min(8, cpus // max(1, num_gpus))) return { - "batchSize": (128, 1024, 128), + "batchSize": (128, 8192, "log"), "workerThreads": (1, per_gpu_worker_max), "preprocessingThreads": (1, cpus), } @@ -121,7 +121,7 @@ def tune_substructure( gpuIds: Optional[Iterable[int]] = None, calibration_set: Optional[Iterable[int]] = None, calibration_fraction: float = 0.1, - calibration_max_size: int = 2000, + calibration_max_size: Optional[int] = None, target_seconds_per_trial: float = 10.0, n_trials: int = 30, search_space_overrides: Optional[dict[str, Any]] = None, @@ -145,7 +145,9 @@ def tune_substructure( gpuIds: GPU device IDs to use. Fixed across the study. calibration_set: Optional explicit indices into ``targets``. calibration_fraction: Fraction of the workload to auto-sample. - calibration_max_size: Cap on the auto-sampled calibration size. + calibration_max_size: Cap on the auto-sampled calibration size, or + ``None`` (the default) for no cap. Small calibration slices + under-reward large ``batchSize`` values. target_seconds_per_trial: Target wall-clock budget for one trial. n_trials: Number of Optuna trials to run after warm-up. search_space_overrides: Optional overrides for ``batchSize``, From 30c9a15a577fce564306562618aa96001bad0a6f Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 29 Jul 2026 17:28:57 -0400 Subject: [PATCH 2/4] Upload the recursive pattern table to the device once preprocessMiniBatch re-staged and re-uploaded the recursive SMARTS pattern entries on every pattern sub-batch, of every mini-batch, on every worker. The data is identical every time: MiniBatchPlanner::prepareRecursiveMiniBatch points plan.patternsAtDepth at leafSubpatterns_.allQueriesPatternsAtDepth, a table built once by buildAllPatterns and shared across all mini-batches. On a 200k-target x 158-query hasSubstructMatch run that was 133,333 uploads costing 4,979 ms of aggregate worker time -- 44% of all time the six worker threads spent launching work, and the single largest item in the profile. Most of it was not the 428-byte copy itself but the handshake around it: the pinned staging area is only double-buffered, so with four sub-batches per mini-batch acquireBufferIndex kept returning a buffer whose copy was still queued behind a full depth level of paint kernels, and waitForBuffer blocked the worker on cudaEventSynchronize. Those accounted for roughly 135k of the 168k blocking event syncs on worker threads; the actual result drain was only 221 ms. Upload the table once in LeafSubpatterns::syncToDevice, which already runs on the setup stream that callers synchronize before starting any worker. Each sub-batch is a contiguous slice of its depth's entries, so the kernels now take a pointer into the device table rather than a per-sub-batch staging buffer, and the whole acquire/wait/fill/copy/record sequence is gone. The non-pipeline preprocessRecursiveSmarts path builds its own local pattern vectors and is left unchanged. --- benchmarks/bench_utils/cli.py | 1 + benchmarks/tests/test_substruct_bench.py | 1 + src/substruct/recursive_preprocessor.cu | 36 +++++++++++++----------- src/substruct/recursive_preprocessor.h | 6 ++++ 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/benchmarks/bench_utils/cli.py b/benchmarks/bench_utils/cli.py index 106a9d04..b4c0aaf3 100644 --- a/benchmarks/bench_utils/cli.py +++ b/benchmarks/bench_utils/cli.py @@ -17,6 +17,7 @@ import argparse + def add_smiles_sanitization_args(parser: argparse.ArgumentParser) -> None: """Add safe-by-default SMILES sanitization flags to ``parser``.""" sanitize_group = parser.add_mutually_exclusive_group() diff --git a/benchmarks/tests/test_substruct_bench.py b/benchmarks/tests/test_substruct_bench.py index 9f055ae7..85bcb6f6 100644 --- a/benchmarks/tests/test_substruct_bench.py +++ b/benchmarks/tests/test_substruct_bench.py @@ -8,6 +8,7 @@ import pytest from bench_utils.cli import add_smiles_sanitization_args + def test_smiles_sanitization_is_enabled_by_default(): parser = argparse.ArgumentParser() add_smiles_sanitization_args(parser) diff --git a/src/substruct/recursive_preprocessor.cu b/src/substruct/recursive_preprocessor.cu index a1bc2f52..d344ad55 100644 --- a/src/substruct/recursive_preprocessor.cu +++ b/src/substruct/recursive_preprocessor.cu @@ -167,6 +167,19 @@ void LeafSubpatterns::syncToDevice(cudaStream_t stream) { return; } patternsDevice.copyFromHost(patternsHost, stream); + + // The all-queries pattern table is the same for every mini-batch, so upload + // it once here. Callers synchronize this stream before any worker starts. + for (int depth = 0; depth <= kMaxSmartsNestingDepth; ++depth) { + const auto& entries = allQueriesPatternsAtDepth[depth]; + auto& entriesOnGpu = allQueriesPatternsAtDepthDevice[depth]; + entriesOnGpu.setStream(stream); + if (entries.empty()) { + continue; + } + entriesOnGpu.resize(entries.size()); + entriesOnGpu.copyFromHost(entries.data(), entries.size()); + } } void RecursivePatternPreprocessor::buildPatterns(const MoleculesHost& queriesHost) { @@ -230,14 +243,10 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( const size_t numPatternsInSubBatch = patternEnd - patternStart; const size_t numBlocksInSubBatch = numTargetsInMiniBatch * numPatternsInSubBatch; - ScopedNvtxRange prepareRange("GPU: Upload pattern entries"); - const int bufferIdx = scratch.acquireBufferIndex(); - scratch.waitForBuffer(bufferIdx); - scratch.ensureCapacity(bufferIdx, static_cast(numPatternsInSubBatch)); - for (size_t i = 0; i < numPatternsInSubBatch; ++i) { - scratch.patternsAtDepthHost[bufferIdx][i] = patternsForDepth[patternStart + i]; - } - prepareRange.pop(); + // The pattern table already lives on the device; this sub-batch is a + // contiguous slice of it, so no staging or H2D copy is needed here. + const BatchedPatternEntry* patternEntriesDevice = + leafSubpatterns_.allQueriesPatternsAtDepthDevice[currentDepth].data() + patternStart; const int buffersPerBlock = gsiBuffersPerBlock; const size_t overflowNeeded = numBlocksInSubBatch * buffersPerBlock * kOverflowEntriesPerBuffer; @@ -251,13 +260,6 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( scratch.labelMatrixBuffer.resize(static_cast(labelMatrixNeeded * 1.5)); } - if (scratch.patternEntries.size() < numPatternsInSubBatch) { - scratch.patternEntries.resize(static_cast(numPatternsInSubBatch * 1.5)); - } - - scratch.patternEntries.copyFromHost(scratch.patternsAtDepthHost[bufferIdx].data(), numPatternsInSubBatch); - scratch.recordCopy(bufferIdx, scratch.patternEntries.stream()); - const uint32_t* recursiveBitsForLabel = (currentDepth > 0) ? miniBatchResults.recursiveMatchBits() : nullptr; std::optional zeroBuffers; @@ -272,7 +274,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( launchLabelMatrixPaintKernel(paintConfig, targetsDevice.view(), leafSubpatterns_.view(), - scratch.patternEntries.data(), + patternEntriesDevice, static_cast(numPatternsInSubBatch), numBlocksInSubBatch, numQueries, @@ -289,7 +291,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( algorithm, targetsDevice.view(), leafSubpatterns_.view(), - scratch.patternEntries.data(), + patternEntriesDevice, static_cast(numPatternsInSubBatch), numBlocksInSubBatch, miniBatchResults.recursiveMatchBits(), diff --git a/src/substruct/recursive_preprocessor.h b/src/substruct/recursive_preprocessor.h index 2201c00f..beaa2669 100644 --- a/src/substruct/recursive_preprocessor.h +++ b/src/substruct/recursive_preprocessor.h @@ -87,6 +87,12 @@ struct LeafSubpatterns { /// Max recursion depth across all queries int allQueriesMaxDepth = 0; + /// Device-resident copy of allQueriesPatternsAtDepth, uploaded once by + /// syncToDevice. The table is invariant across mini-batches, so workers + /// index into it directly rather than re-staging it through pinned memory + /// on every pattern sub-batch. + std::array, kMaxSmartsNestingDepth + 1> allQueriesPatternsAtDepthDevice; + int maxPatternAtoms_ = 0; LeafSubpatterns() = default; From 0134dd082cadd1daffd188e82c92ff9201628e72 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Thu, 30 Jul 2026 11:10:38 -0400 Subject: [PATCH 3/4] Further improvemetns --- benchmarks/bench_utils/__init__.py | 2 -- benchmarks/bench_utils/cli.py | 36 ------------------------ benchmarks/substruct_bench.py | 4 ++- benchmarks/tests/test_substruct_bench.py | 26 ----------------- nvmolkit/autotune/_calibration.py | 14 ++++----- nvmolkit/autotune/tune_substructure.py | 6 ++-- src/substruct/recursive_preprocessor.cu | 6 ++-- src/substruct/recursive_preprocessor.h | 5 +--- 8 files changed, 13 insertions(+), 86 deletions(-) delete mode 100644 benchmarks/bench_utils/cli.py delete mode 100644 benchmarks/tests/test_substruct_bench.py diff --git a/benchmarks/bench_utils/__init__.py b/benchmarks/bench_utils/__init__.py index 74366b35..4cabcc21 100644 --- a/benchmarks/bench_utils/__init__.py +++ b/benchmarks/bench_utils/__init__.py @@ -19,7 +19,6 @@ so individual bench scripts can ``from bench_utils import time_it, load_smiles``. """ -from bench_utils.cli import add_smiles_sanitization_args from bench_utils.loaders import load_pickle, load_sdf, load_smarts, load_smiles from bench_utils.molprep import clone_mols_with_conformers, embed_and_jitter, perturb_conformer, prep_mols from bench_utils.timing import ( @@ -35,7 +34,6 @@ "Deadline", "TimingResult", "add_rdkit_max_seconds_arg", - "add_smiles_sanitization_args", "clone_mols_with_conformers", "embed_and_jitter", "load_pickle", diff --git a/benchmarks/bench_utils/cli.py b/benchmarks/bench_utils/cli.py deleted file mode 100644 index b4c0aaf3..00000000 --- a/benchmarks/bench_utils/cli.py +++ /dev/null @@ -1,36 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Shared command-line argument helpers for nvMolKit benchmarks.""" - -import argparse - - -def add_smiles_sanitization_args(parser: argparse.ArgumentParser) -> None: - """Add safe-by-default SMILES sanitization flags to ``parser``.""" - sanitize_group = parser.add_mutually_exclusive_group() - sanitize_group.add_argument( - "--sanitize", - action="store_true", - dest="sanitize", - default=True, - help="Sanitize SMILES during parsing (default)", - ) - sanitize_group.add_argument( - "--no_sanitize", - action="store_false", - dest="sanitize", - help="Skip sanitization (preprocessed SMILES only)", - ) diff --git a/benchmarks/substruct_bench.py b/benchmarks/substruct_bench.py index 78b9c2e4..4a88602e 100644 --- a/benchmarks/substruct_bench.py +++ b/benchmarks/substruct_bench.py @@ -335,7 +335,9 @@ def main(): default=42, help="Random seed for sampling SMILES when --num_mols > 0 (default: 42)", ) - add_smiles_sanitization_args(parser) + sanitize_group = parser.add_mutually_exclusive_group() + sanitize_group.add_argument("--sanitize", action="store_true", dest="sanitize", default=True) + sanitize_group.add_argument("--no_sanitize", action="store_false", dest="sanitize") parser.add_argument("--runs", "-r", type=int, default=1, help="Number of timing runs (default: 1)") parser.add_argument( "--mode", diff --git a/benchmarks/tests/test_substruct_bench.py b/benchmarks/tests/test_substruct_bench.py deleted file mode 100644 index 85bcb6f6..00000000 --- a/benchmarks/tests/test_substruct_bench.py +++ /dev/null @@ -1,26 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for the substructure benchmark CLI.""" - -import argparse - -import pytest -from bench_utils.cli import add_smiles_sanitization_args - - -def test_smiles_sanitization_is_enabled_by_default(): - parser = argparse.ArgumentParser() - add_smiles_sanitization_args(parser) - - assert parser.parse_args([]).sanitize is True - assert parser.parse_args(["--sanitize"]).sanitize is True - assert parser.parse_args(["--no_sanitize"]).sanitize is False - - -def test_smiles_sanitization_flags_are_mutually_exclusive(): - parser = argparse.ArgumentParser() - add_smiles_sanitization_args(parser) - - with pytest.raises(SystemExit): - parser.parse_args(["--sanitize", "--no_sanitize"]) diff --git a/nvmolkit/autotune/_calibration.py b/nvmolkit/autotune/_calibration.py index a7782dea..e4a1b5af 100644 --- a/nvmolkit/autotune/_calibration.py +++ b/nvmolkit/autotune/_calibration.py @@ -25,7 +25,7 @@ def auto_subsample( workload_size: int, *, fraction: float = 0.1, - max_size: Optional[int] = 2000, + max_size: int = 2000, min_size: int = 1, seed: int = 0, ) -> list[int]: @@ -33,14 +33,11 @@ def auto_subsample( The slice size is ``min(max_size, max(min_size, round(fraction * workload_size)))`` and is shuffled deterministically with ``seed`` so trials sample roughly - uniformly across the workload. Passing ``max_size=None`` disables the cap, - leaving ``fraction`` as the only limit. + uniformly across the workload. """ if workload_size <= 0: raise ValueError("workload_size must be positive") - target = max(min_size, int(round(fraction * workload_size))) - if max_size is not None: - target = min(max_size, target) + target = min(max_size, max(min_size, int(round(fraction * workload_size)))) target = min(target, workload_size) rng = random.Random(seed) indices = list(range(workload_size)) @@ -53,7 +50,7 @@ def normalize_calibration_set( workload_size: int, *, fraction: float = 0.1, - max_size: Optional[int] = 2000, + max_size: int = 2000, seed: int = 0, ) -> list[int]: """Return calibration indices from an explicit set or auto-sample. @@ -66,8 +63,7 @@ def normalize_calibration_set( auto-subsample. workload_size: Total size of the user workload. fraction: Auto-subsample fraction when ``calibration_set`` is ``None``. - max_size: Auto-subsample cap when ``calibration_set`` is ``None``, or - ``None`` for no cap. + max_size: Auto-subsample cap when ``calibration_set`` is ``None``. seed: Seed for auto-subsampling. """ if calibration_set is None: diff --git a/nvmolkit/autotune/tune_substructure.py b/nvmolkit/autotune/tune_substructure.py index e864ae65..dfbeffd7 100644 --- a/nvmolkit/autotune/tune_substructure.py +++ b/nvmolkit/autotune/tune_substructure.py @@ -121,7 +121,7 @@ def tune_substructure( gpuIds: Optional[Iterable[int]] = None, calibration_set: Optional[Iterable[int]] = None, calibration_fraction: float = 0.1, - calibration_max_size: Optional[int] = None, + calibration_max_size: int = 50_000, target_seconds_per_trial: float = 10.0, n_trials: int = 30, search_space_overrides: Optional[dict[str, Any]] = None, @@ -145,9 +145,7 @@ def tune_substructure( gpuIds: GPU device IDs to use. Fixed across the study. calibration_set: Optional explicit indices into ``targets``. calibration_fraction: Fraction of the workload to auto-sample. - calibration_max_size: Cap on the auto-sampled calibration size, or - ``None`` (the default) for no cap. Small calibration slices - under-reward large ``batchSize`` values. + calibration_max_size: Cap on the auto-sampled calibration size. target_seconds_per_trial: Target wall-clock budget for one trial. n_trials: Number of Optuna trials to run after warm-up. search_space_overrides: Optional overrides for ``batchSize``, diff --git a/src/substruct/recursive_preprocessor.cu b/src/substruct/recursive_preprocessor.cu index d344ad55..5fcff86c 100644 --- a/src/substruct/recursive_preprocessor.cu +++ b/src/substruct/recursive_preprocessor.cu @@ -168,8 +168,7 @@ void LeafSubpatterns::syncToDevice(cudaStream_t stream) { } patternsDevice.copyFromHost(patternsHost, stream); - // The all-queries pattern table is the same for every mini-batch, so upload - // it once here. Callers synchronize this stream before any worker starts. + // Upload the immutable pattern table once. for (int depth = 0; depth <= kMaxSmartsNestingDepth; ++depth) { const auto& entries = allQueriesPatternsAtDepth[depth]; auto& entriesOnGpu = allQueriesPatternsAtDepthDevice[depth]; @@ -243,8 +242,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( const size_t numPatternsInSubBatch = patternEnd - patternStart; const size_t numBlocksInSubBatch = numTargetsInMiniBatch * numPatternsInSubBatch; - // The pattern table already lives on the device; this sub-batch is a - // contiguous slice of it, so no staging or H2D copy is needed here. + // Use the device-resident pattern table directly. const BatchedPatternEntry* patternEntriesDevice = leafSubpatterns_.allQueriesPatternsAtDepthDevice[currentDepth].data() + patternStart; diff --git a/src/substruct/recursive_preprocessor.h b/src/substruct/recursive_preprocessor.h index beaa2669..5a5e37c1 100644 --- a/src/substruct/recursive_preprocessor.h +++ b/src/substruct/recursive_preprocessor.h @@ -87,10 +87,7 @@ struct LeafSubpatterns { /// Max recursion depth across all queries int allQueriesMaxDepth = 0; - /// Device-resident copy of allQueriesPatternsAtDepth, uploaded once by - /// syncToDevice. The table is invariant across mini-batches, so workers - /// index into it directly rather than re-staging it through pinned memory - /// on every pattern sub-batch. + /// Device-resident pattern table shared across mini-batches. std::array, kMaxSmartsNestingDepth + 1> allQueriesPatternsAtDepthDevice; int maxPatternAtoms_ = 0; From f8d5710cb186e89d95dd0c443470c90f00d86ba7 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Thu, 30 Jul 2026 11:46:12 -0400 Subject: [PATCH 4/4] Fix substructure setup ordering and autotune test --- nvmolkit/tests/test_autotune.py | 9 +++++---- src/substruct/substruct_search.cu | 5 +++++ 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/nvmolkit/tests/test_autotune.py b/nvmolkit/tests/test_autotune.py index d4a52803..e8c3833b 100644 --- a/nvmolkit/tests/test_autotune.py +++ b/nvmolkit/tests/test_autotune.py @@ -289,10 +289,11 @@ def test_default_substruct_search_space_caps_per_pool(): assert space_4gpu["preprocessingThreads"] == (1, 16) assert space_64gpu["preprocessingThreads"] == (1, 16) - low, high, step = space_1gpu["batchSize"] - assert step % 64 == 0 - assert low % step == 0 and high % step == 0 - assert low <= high + # Unlike the FF and embed spaces, substruct batchSize is log-uniform rather + # than a stepped range: the useful span is two orders of magnitude. + low, high, distribution = space_1gpu["batchSize"] + assert distribution == "log" + assert 0 < low <= high class _RecordingTrial: diff --git a/src/substruct/substruct_search.cu b/src/substruct/substruct_search.cu index 4dee07da..99a88f23 100644 --- a/src/substruct/substruct_search.cu +++ b/src/substruct/substruct_search.cu @@ -439,6 +439,11 @@ void runGpuCoordinator(int deviceId, localPreprocessor = std::make_unique(); localPreprocessor->buildPatterns(queriesHost); localPreprocessor->syncToDevice(nullptr); + + // The workers use independent non-blocking streams, so fully publish the + // secondary-device queries and patterns before any worker can consume them. + cudaCheckError(cudaStreamSynchronize(nullptr)); + queriesPtr = localQueries.get(); preprocessorPtr = localPreprocessor.get(); }