feat: add CUDA ball-query backend for smooth lDDT loss

[heathcliff233]

Summary

Adds a CUDA ball-query backend for the smooth lDDT loss, opt-in via config and an [smooth-lddt-kernel] install extra. The dense reference implementation in smooth_lddt_loss remains the default — nothing changes for users who don't enable the new path.

For training on long targets, the existing dense smooth lDDT computes an [N, N] pairwise-distance tensor for both predicted and ground-truth coordinates, then keeps [N, N]-shaped autograd intermediates through the elementwise scoring stage (sqrt → abs → 4× sigmoid → mask → reduce). This O(N²) term dominates both memory and runtime once n_atom grows past ~1500 atoms. The ball-query backend replaces the dense pairwise scan with a sparse [N, K] neighbor list and pushes both the predicted-distance computation and its gradient into CUDA, so memory grows with K rather than N. With reservoir sampling on the kernel side, neighbor selection within the radius is uniform-random, matching the in-radius statistics that the dense backend would compute.

Changes

• New CUDA extension under openfold3/core/kernels/smooth_lddt/ball_query_ext/:
  • BallQueryKernelCoopWithPred — warp-cooperative (W=8 lanes/atom) ball query with reservoir sampling, finds GT neighbors and emits predicted squared distances inline. Eliminates the [N, K, 3] x_j gather entirely.
  • BallQueryPredBackwardKernel — dedicated backward that scatters 2 · grad_dists · (pred_i − pred_j) into grad_pred via atomicAdd. Saves only (pred, idx) for backward — the [N, K, 3] diff tensor is never materialized. Pattern adapted from PyTorch3D's _ball_query + knn_points_backward (BSD-3-Clause; attribution preserved in headers).
• New Python module openfold3/core/kernels/smooth_lddt/:
  • ball_query.py — JIT loader (torch.utils.cpp_extension.load) + thin wrappers around the CUDA entry points.
  • __init__.py — _BallQueryWithPredDist (custom autograd.Function) and ball_query_smooth_lddt_loss. The remaining elementwise scoring is wrapped in torch.utils.checkpoint(use_reentrant=False) so all [N, K] intermediates are freed between forward and backward.
• openfold3/core/loss/diffusion.py — backend dispatch in the diffusion loss; new keys smooth_lddt_backend (default "dense"), smooth_lddt_top_k. The existing internal chunk_size/run_low_mem_loss_fn path is incompatible with ball-query and is rejected with a clear error.
• openfold3/projects/of3_all_atom/config/model_config.py — adds the two new keys with safe defaults; existing runner YAMLs are unchanged.
• pyproject.toml — new optional extra [smooth-lddt-kernel] (just ninja); not part of the default install. Adds a benchmark pytest marker.
• examples/example_runner_yamls/smooth_lddt_ball_query.yml — example training override that enables the backend.
• Tests: openfold3/tests/test_diffusion_loss.py adds dispatch + numerical-equivalence + gradient-equivalence + bf16 + chunk-size-error tests for the ball-query path. openfold3/tests/test_smooth_lddt_benchmark.py is a --benchmark-only sweep across n_atom, top_k, and dtype.

---

Theoretical memory: dense vs ball-query

For batch B, atoms N, neighbors K, scoring dtype s bytes (fp32 = 4 B; bf16/fp16 = 2 B):

Dense backend (smooth_lddt_loss) — keeps [B, N, N]-shaped tensors live for autograd. The dominant terms are:

Tensor (forward / saved)	Shape	Bytes	At N=7680, B=2, fp32
(x_i − x_j)² reduction inputs	[B, N, N, 3] fp32	12·B·N²	~1.4 GB
dx, dx_gt − dx, 4× sigmoid, c, c·e (autograd saves)	[B, N, N] × ~8	8s · B · N²	~3.8 GB
GT-distance / mask intermediates	[B, N, N] × few	4·B·N²	~470 MB

Total ≈ (16 + 8s)·B·N² bytes; O(B · N²) regardless of locality. The dense path cannot exploit the r ≤ 30 Å smooth lDDT cutoff because the cutoff is applied after the full pairwise scan.

Ball-query backend (ball_query_smooth_lddt_loss) — keeps [B, N, K]-shaped tensors. The dominant terms are:

Tensor (forward / saved)	Shape	Bytes	At N=7680, B=2, K=256, fp32
idx (saved by autograd)	[B, N, K] int64	8·B·N·K	~31 MB
dists_pred (saved by autograd)	[B, N, K] s	s·B·N·K	~16 MB
dists_gt (forward only)	[B, N, K] fp32	4·B·N·K	~16 MB
pred (saved by autograd)	[B, N, 3] s	3s·B·N	~180 KB
Scoring intermediates	freed by checkpoint	0	0

Total ≈ (8 + 2s)·B·N·K bytes; O(B · N · K).

Asymptotic ratio (dropping constants): memory_dense / memory_bq ≈ N / K. At N=7680, K=256 that is ~30×; at N=4000, K=256 it is ~16×. The same N/K factor governs runtime in the regime where the kernel is bandwidth-bound by the saved/scored tensors. The advantage disappears once K ≥ N.

---

Measured runtime + memory (RTX 4090, fwd+bwd peak)

Synthetic protein-density structures (~0.07 atoms / ų), B = 2 diffusion samples, fwd+bwd:

fp32

n_atom	backend	time	peak mem	speedup vs dense	memory savings
4000	dense	23.5 ms	1856 MB	1.00×	1.00×
	bq K=256	6.4 ms	146 MB	3.7×	12.7×
	bq K=512	6.8 ms	290 MB	3.4×	6.4×
	bq K=1024	8.8 ms	575 MB	2.7×	3.2×
6000	dense	54.3 ms	4177 MB	1.00×	1.00×
	bq K=256	7.9 ms	219 MB	6.9×	19.1×
	bq K=512	8.1 ms	435 MB	6.7×	9.6×
	bq K=1024	12.4 ms	864 MB	4.4×	4.8×
7680	dense	85.6 ms	6843 MB	1.00×	1.00×
	bq K=256	7.8 ms	286 MB	11.0×	23.9×
	bq K=512	9.2 ms	552 MB	9.3×	12.4×
	bq K=1024	16.5 ms	1096 MB	5.2×	6.2×

bf16

n_atom	backend	time	peak mem	speedup vs dense	memory savings
4000	dense	19.6 ms	1572 MB	1.00×	1.00×
	bq K=256	7.1 ms	138 MB	2.8×	11.4×
6000	dense	43.2 ms	3531 MB	1.00×	1.00×
	bq K=256	7.9 ms	207 MB	5.5×	17.1×
7680	dense	70.4 ms	5784 MB	1.00×	1.00×
	bq K=256	7.9 ms	269 MB	8.9×	21.5×

At the upper end of typical AF3 training samples (N≈7680), ball-query (K=256) is ~10× faster and uses ~22× less GPU memory than the dense baseline. The savings grow further for any N beyond that — the dense path scales as N², the ball-query path as N.

For the four small minimal-training samples in the unit-test suite (N = 448 / 635 / 1262 / 3329):

sample	n_atom	dense bwd	bq K=256 bwd	bq K=512 bwd	bq K=1024 bwd
17ra	448	23.4 MB	16.4 MB	27.9 MB	27.9 MB
134d	635	46.9 MB	24.1 MB	45.6 MB	56.2 MB
102m	1262	184.9 MB	46.1 MB	93.0 MB	181.1 MB
12e8	3329	1286.0 MB	123.5 MB	246.1 MB	474.9 MB

Crossover is around N≈600; ball-query starts winning materially around N≈1500. Reproduce with pytest --benchmark-only openfold3/tests/test_smooth_lddt_benchmark.py (sweeps K over {128, 256, 512, 768, 1024, 2048} × {fp32, bf16}).

---

Numerical equivalence

The relevant smooth-lDDT cutoff is 30 Å. For protein-density structures, the number of atoms within 30 Å of any given atom concentrates well below 1000. As long as K covers the full in-radius set, the ball-query backend is bit-equivalent up to floating-point reorder to the dense path:

• test_smooth_lddt_ball_query_matches_dense: top_k = n_atom − 1 → forward agrees with dense at atol = rtol = 1e-5 (fp32).
• test_smooth_lddt_ball_query_gradient_matches_dense: same setup → gradient w.r.t. x agrees at atol = 1e-3, rtol = 5e-3. The looser tolerance comes from atomicAdd reordering and the [N, K+1] reduction order vs. [N, N] in the dense path.
• test_smooth_lddt_ball_query_bf16_path: bf16 input → bf16 output; loss agrees with fp32 at atol = rtol = 5e-3, gradient at atol = rtol = 5e-2.
• test_smooth_lddt_ball_query_chunk_size_error: enabling both chunk_size and ball-query is rejected.

When K is smaller than the true in-radius count, ball-query returns a uniform-random size-K subsample of the qualifying neighbors per atom — an unbiased estimator of the same lDDT mean, with variance that vanishes as K → in_radius_count.

---

Warp-cooperative search with reservoir sampling

Each query atom is owned by W = 8 cooperating threads. The candidate atoms are partitioned across the W lanes; each lane scans its stride and tests ‖p1_i − p2_j‖² ≤ r². Qualifying neighbors race for output slots via a shared atomicAdd counter, so the kernel never makes a global pass over all N² pairs and the work is balanced across the warp.

Two-phase sampling:

1. Fill phase (seen < K): the first K qualifying neighbors are written into output slots in order. Each slot is written by exactly one thread.
2. Reservoir phase (seen ≥ K): each newly-qualifying neighbor is admitted with probability K / total_qualified by Vitter's Algorithm R. A per-lane hash-RNG (hash_rng(seed ⊕ n·1000003 ⊕ i·997 ⊕ lane·31)) replaces a uniformly random slot.

To recover bit-exact agreement with dense, set top_k ≥ max_neighbors_in_radius (in practice top_k = n_atom − 1): both phases vanish, every qualifying neighbor lands in a fill-phase slot, and the output set is exactly the in-radius set. This is what the equivalence tests above use.

The sampling itself is a pure function of (seed, n, i, lane) — pass a fixed seed (e.g. the training step number) for reproducible survivors. Two remaining sources of nondeterminism:

1. Slot order within a fill batch. When multiple lanes hit atomicAdd(counter) in the same cycle, the order of slot assignment depends on the GPU scheduler. Survivor set is identical, but its layout in the [K] axis can permute between runs. The loss is a sum over K so this does not affect the value, but if bit-exact idx is needed we can switch to a per-lane interleaved layout (each lane owns slots lane, lane+W, lane+2W, …).
2. Backward atomicAdd reordering. BallQueryPredBackwardKernel accumulates fp32 grads via atomicAdd (matches PyTorch3D's knn_points_backward and PyTorch's index_add_). With torch.use_deterministic_algorithms(True) this kernel emits a "nondeterministic op" warning; a sort-and-segment-reduce variant is straightforward to add as an opt-in flag if needed (~2× backward time).

---

Related Issues

None.

Testing

• pytest openfold3/tests/test_diffusion_loss.py -q — 14/14 pass (5 new ball-query tests + dispatch error + dense baseline).
• pytest --benchmark-only openfold3/tests/test_smooth_lddt_benchmark.py — runs without regression; numbers above (RTX 4090, torch 2.7.1+cu126).
• ruff format and ruff check clean on all touched files.

Other Notes

• Default config (smooth_lddt_backend: "dense") is unchanged; default install does not pull ninja. Three explicit opt-ins are required: install extra (pip install -e ".[smooth-lddt-kernel]"), config (smooth_lddt_backend: ball_query plus smooth_lddt_top_k), and CUDA availability.
• The CUDA extension JITs on first use via torch.utils.cpp_extension.load; set OPENFOLD3_SMOOTH_LDDT_VERBOSE=1 for the full nvcc log.
• File headers credit PyTorch3D (BSD-3-Clause) for the autograd / scatter-backward pattern with links to the upstream repo.
• See openfold3/core/kernels/smooth_lddt/README.md for the kernel layout and a more detailed view of the autograd boundary.

[jandom]

Big fan of this @heathcliff233 !

[jnwei]

Thank you for this addition @heathcliff233 ! Overall this is a really well done PR, with extensive benchmarking and tests.

I have personally tested this on a DGX Spark and an A100, the new diffusion tests pass on both of these machines. I was unable to run the benchmark tests. It would be good if you provide a script to generate the mini pdb dataset.

The cuda kernel required by the ball query smooth lddt feature presents an interesting packaging challenge. I am calling on @sdvillal , our resident packaging expert for advice here.

In short, in order for us to continue to make openfold3 distributable on pypi, we cannot include any cuda / non python code in the pypi recipe. The custom cuda kernel that was present in the original OpenFold repository was the main blocker to why we cannot offer the original openfold on pypi.

My understanding (please correct me if I'm wrong @sdvillal ) is that we have two options for the ball query kernel:

• Option 1: We make a separate package for the ball query kernel that is separately compiled. OpenFold3 can then install the ball query kernel as an extra dependency, similar to the approach used for installing cuequivariance. In this case, it would be good if @heathcliff233 can maintain the ball query package. This may involve creating packages that are suitable for multiple versions of cuda / pytorch.

• Option 2: We keep the ball query source code in openfold3, but exclude it from the pypi recipe. A user who wishes to use the ball query source code will need to install from source.

For both Option 1 and Option 2, changes will need to be made to the base OpenFold3 code to try / except installation of the ball query code. I believe that @heathcliff233 already wrote most of the installation guards required, but we should explicitly test this with the different installation paths.

I think of these options, Option 1 is preferred. While Option 2 might be possible, it could create too much overhead for us to maintain in the long run. But I will defer to @sdvillal for his advice.

[sdvillal]

I would not create two new environments, but simply add the extension to the existing ones (including the pypi envs)

[jnwei]

Hi @heathcliff233 after chatting with @sdvillal offline, I think we can go forward with Option 2, e.g. include the smooth ball query code in OpenFold3, but do not include for OpenFold3 packaging to avoid having to distribute CUDA code.

If that sounds good to you, we can move forward with the PR. Do you need any help with the script to fetch a minimal PDB set for the benchmarking tests?

[heathcliff233]

Thanks @jnwei! Option 2 sounds good to me.

I've pushed a commit that addresses the benchmark data request:

• scripts/download_smooth_lddt_benchmark_data.py — downloads the preprocessed NPZ samples from s3://openfold3-data (OF3 public training data), or symlinks from an existing local training set via --training-set-dir. Usage:

  python scripts/download_smooth_lddt_benchmark_data.py
  

• Added two larger samples (1tii: 5,667 atoms, 1xfy: 7,533 atoms) alongside the original four, to better show the scaling behavior. Restructured the benchmark test to print a per-sample comparison table.

To run the correctness tests (numerical equivalence + gradient parity vs dense + bf16 + error handling):

pytest openfold3/tests/test_diffusion_loss.py -v

To run the benchmark (requires downloading the data first):

pytest openfold3/tests/test_smooth_lddt_benchmark.py -s -v

[jnwei]

Thanks for including the script to download PDB examples @heathcliff233 ! Now I am able to run the benchmark tests. It's a very nice presentation of the speedups.

This is getting really close. I have a few more general change requests:

• At this point, building OpenFold3 without compiling the smooth ball query kernel results in an import error. We will need to resolve these issues in order to ensure users who are using the pypi distribution can still run OpenFold3 withoug error. I will provide more detailed instructions for how you can test a build locally, and where you might introduce a fix.

• Benchmark file locations: Could you pleas move both the test_smooth_lddt_benchmark.py and the download_smooth_lddt_benchmark_data.py to a sub directory under scripts (e.g. scripts/benchmarks/ball-query-lddt-benchmark). At this point, we aren't well equipped to run these benchmark tests as part of our continuous integration / testing harness.

• Documentation: Could you add some documentation to the OpenFold3 docs about how to install the ball query kernel? I have reformatted the kernel documentation page slightly and added a stub for the ball query documentation.

[heathcliff233]

Thanks @jnwei , addressed in 03ddfa7

1. For install-error, is_smooth_lddt_kernel_installed is now separated from is_smooth_lddt_kernel_available. Deferred the torch.utils.cpp_extension import. get_smooth_lddt_loss_fn("ball_query") now raises two distinct RuntimeErrors pointing at the missing piece.
2. For benchmarks. They are now under scripts/benchmarks/ball-query-lddt-benchmark/ as suggested.
3. For the kernels related docs, I added something about the ball query kernel on the installation, target usage and limitations.

[jnwei]

Thanks for the documentation and the fixes!

I will aim to test these changes locally by early next week. Update I have now run the tests locally.

@heathcliff233 please merge main into this PR. I think are almost finished PR, and some of the tests will fail without the fixes from the main branch.

Once main has been merged in, all test seemsto pass with the pip install build using the [dev, deepspeed] extra options, following the deepspeed installation instructions

I have not tested the pip install openfold3[smooth-lddt-kernel] option, as I believe this option should be removed.

A few further documentation suggestions and I think we wrap things up. Thanks again for the hard work @heathcliff233 !

[jnwei]

I think it is cleaner to tell users to install the ball query kernel via the just the pixi route. The extra CUDA toolchain requirements as noted in the documentation means that installation will not be tightly contained and will be a struggle for users to debug.

@sdvillal would you agree with this approach?

[jnwei]

This runner.yaml example should be expanded to include the full model_update sections, as you have in the example runner.yml file, e.g.

model_update:
  custom:
    architecture:
      loss_module:
        diffusion:
          chunk_size: null
          smooth_lddt_backend: ball_query
          smooth_lddt_top_k: 128

[heathcliff233]

Thanks @jnwei (and @sdvillal) — addressed the remaining review items in three commits:

1. 7aaf9a47 — merged main into this branch and regenerated pixi.lock after the lock conflict.
2. c9a09ccc — folded the kernel into existing pixi envs: restored env name openfold3-cuda13 (dropped openfold3-cuda13-smooth-lddt), kept the smooth-lddt-kernel feature on openfold3-cuda12 / openfold3-cuda13 / pypi CUDA envs only, and removed the openfold3[smooth-lddt-kernel] pip optional extra.
3. f777cca4 — docs/install strings: pixi-first install path in docs/source/kernels.md, full model_update.custom.architecture.loss_module.diffusion runner YAML example, and matching install/error/skip messages.

Local check with the of3bq env: pytest openfold3/tests/test_diffusion_loss.py -q → 14/14 passed.

[jnwei]

Overall, the changes are great, I have tested building the sdist package locally with [dev, deepspeed] optional dependencies and can verify that everything passes.

I see that the last commit adds a triton version of the smooth_lddt_ball_query kernel. That is very exciting; we can consider including the triton versions in the default openfold3 package as it does not require additional non-python build artifacts.

Could I ask you to please make a new PR with the triton kernel? The code changes for the kernel are quite large and deserve a separate review. Additionally, we will want to reorganize the pixi environments to select the triton build when CUDA is not available.

Thank you very much for your hard work and patience. Once the triton commit is moved and the main branch is merged in, I think this PR is good to merge.

[heathcliff233]

Overall, the changes are great, I have tested building the sdist package locally with [dev, deepspeed] optional dependencies and can verify that everything passes.

I see that the last commit adds a triton version of the smooth_lddt_ball_query kernel. That is very exciting; we can consider including the triton versions in the default openfold3 package as it does not require additional non-python build artifacts.

Could I ask you to please make a new PR with the triton kernel? The code changes for the kernel are quite large and deserve a separate review. Additionally, we will want to reorganize the pixi environments to select the triton build when CUDA is not available.

Thank you very much for your hard work and patience. Once the triton commit is moved and the main branch is merged in, I think this PR is good to merge.

Thank you for the feedback. @Lim-ZQ and I are trying to migrate the original CUDA code to triton to get rid of the dependency headache. It took us sometime but we think the current triton version can now deliver similar optimizations and could serve as a replacement for the CUDA version (and potentially on AMD platform). We will revert this commit and create a new PR for the triton kernel recently.

[heathcliff233]

hi @jnwei , the latest commit is reverted and moved to PR #351 . We believe that it might be a better way to just migrate the triton one in #351 as that requires no additional dependencies and is compatible with all other types of devices. Thanks again for your effort on it.

[jnwei]

Hi @heathcliff233

Certainly, If you and @Lim-ZQ prefer, we can merge the triton PR#351 in lieu of this one. I see that you have already provided benchmark run statistics in the PR description there.

If this is the case, could you please close this PR?

We currently have a large number of PRs to process, but if possible, I will try to review #351 later this week.

Thank you again for your contributions!

[heathcliff233]

Sure. Thank you very much for your help!