GPU geometry for PyTorch point clouds. Fast k-NN, farthest-point sampling, interpolation, and ray tracing.

Documentation · API reference · Benchmarks

Performance

25.3x faster k-NN than torch_cluster GPU and 9.2x faster approximate FPS than fpsample CPU, averaged across all 20 benchmark workloads (arithmetic mean of per-workload ratios).

torchbvh k-NN/MLS rows were refreshed on 2026-10-07; the 160 competitor, grid_sample and FPS rows retain their 2026-09-23 measurements. See the benchmark protocol and data and MLS numerical behavior.

About

torchbvh provides CUDA operations for 2D and 3D float32 point clouds, including fixed-size batches:

• k-NN: BVH-accelerated nearest-neighbor search.
• FPS: exact and bucketed approximate farthest-point sampling.
• Interpolation: moving least squares (MLS) over scattered point features.
• Ray tracing: segment and triangle intersections.

Installation

Install CUDA-enabled PyTorch, a matching CUDA toolkit with NVCC, and a supported C++ compiler, then install torchbvh:

pip install torchbvh

Starting with 0.3.3, installation packages the CUDA sources and leaves your PyTorch installation unchanged. Importing torchbvh does not compile anything. The first operation compiles the native extension against your installed PyTorch and CUDA toolkit; later operations and Python sessions reuse the cached binary. Changing the source or build environment selects a separate cache.

See the build guide for compiler and source-install details.

Quickstart

import torch
import torchbvh as tb

points = torch.rand(10_000, 3, device="cuda")
queries = torch.rand_like(points)

with tb.BVH(points) as bvh:
    neighbors, distances_sq = bvh.knn(queries, k=4)

sample = tb.fps(points, target_tokens=2_500, mode="approx_bucketed")
features = torch.rand(10_000, 64, device="cuda")
values = tb.mls_interpolate(points, queries, features, k=4)

See the examples for exact FPS, ray tracing, and gradients.

Share point geometry within a forward

other_queries = queries + 0.01
with tb.PointGeometry(points) as geometry:
    values = tb.mls_interpolate(points, queries, features, k=4, geometry=geometry)
    other = tb.mls_interpolate(points, other_queries, features, k=4, geometry=geometry)
# Saved tensors remain available for backward after the geometry closes.

Create a new geometry object each model forward. The original tensor must stay unchanged while the object is open. Conditional MLS accepts true_geometry and false_geometry; multihead MLS accepts geometry. See prepared geometry for lifecycle rules and automatic dispatch eligibility.

References

The BVH implementation builds on Chitalu, Dubach, and Komura (2020) and ImplicitBVH.jl. Released under the MIT license.