Skip to content

v2.10.7

Choose a tag to compare

View all tags
@volgin volgin released this 28 May 20:00
· 4 commits to main since this release
v2.10.7
5103619

Highlights

  • Fixed affinity prediction running ~5× slower than necessary at default settings. Upstream hardcoded max_parallel_samples=1 for the affinity diffusion path, which was silently a no-op under upstream's buggy chunking math. Our earlier diffusion fix made that hardcode actually take effect, forcing N sequential single-sample forward passes per affinity record. The affinity path now honors the user's --max_parallel_samples, capped at --diffusion_samples_affinity.

Details

Background — how this regression slipped in

Upstream's chunking math in diffusionv2.py was:

sample_ids_chunks = sample_ids.chunk(multiplicity % max_parallel_samples + 1)

tensor.chunk(n) splits a tensor into n chunks, not into chunks of size n. For affinity (multiplicity=5, max_parallel_samples=1) this collapsed to chunk(1) — one chunk of 5 samples, all batched into a single forward pass. The hardcoded =1 upstream put on the affinity path never actually took effect.

Our v2.10.x diffusion fix replaced that with the obviously correct:

sample_ids_chunks = sample_ids.split(max_parallel_samples)

That fix was independently correct — it made --max_parallel_samples behave as documented in divisible cases — but it changed the semantics under the affinity caller, which was unknowingly tuned against the broken behavior. With the fix in place, the hardcoded 1 started doing what it said: 5 sequential single-sample passes per affinity record.

What changed

  • New helper _resolve_affinity_max_parallel_samples(max_parallel_samples, diffusion_samples_affinity) in src/boltz/main.py:
    • Honors the CLI --max_parallel_samples for affinity, matching the behavior already used for structure prediction.
    • Caps at diffusion_samples_affinity so we never claim to batch more affinity samples than diffusion will actually run.
    • Treats max_parallel_samples=None the same way the diffusion code does ("use multiplicity") — relevant for programmatic callers of predict.callback(...) that pass None.
  • predict_affinity_args calls the helper instead of hardcoding 1.

Behavior matrix

--max_parallel_samples --diffusion_samples_affinity Resolved affinity max Affinity behavior
5 (default) 5 (default) 5 1 batched forward pass (matches upstream's effective pre-fix behavior)
2 5 2 User override respected
50 5 5 Capped — no inflated work
1 5 1 Explicit single-sample preserved
None 5 5 Falls back to diffusion sample count

Tests

Six new unit tests in tests/test_cli.py::TestAffinityMaxParallelSamples:

  • Default CLI setting batches all samples.
  • User overrides below the cap are honored.
  • User overrides above the cap are capped.
  • Explicit 1 is preserved.
  • None falls back to diffusion_samples_affinity (regression guard against TypeError: '<' not supported between 'NoneType' and 'int').
  • Source regression guard: scans main.py for the predict_affinity_args block and fails if anyone re-hardcodes \"max_parallel_samples\": 1 there or removes the helper.

Total fork test count is now 283.

What did not change

  • The chunking primitive itself is correct; this only fixes the call site that was implicitly relying on the old broken behavior.
  • Structure prediction was already routing --max_parallel_samples correctly.
  • Numerical outputs are unchanged — this is purely a parallelism / wall-clock fix.

Commits since v2.10.6

  • 5103619 Release 2.10.7
Assets 2
Loading