fix all neighbors output in cuda-12

docs/release-notes/0.17.0.md

@@ -9,6 +9,7 @@
 * Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` losing the neighbors that cross cluster boundaries on multiple GPUs; ``overlap_factor`` now scales with the cluster count {pr}`769` {smaller}`S Dicks`
 * Fix {func}`~rapids_singlecell.pp.neighbors` with ``metric="inner_product"`` failing in the connectivity step for every ``algorithm``. The metric is no longer forwarded to ``cuml``'s ``fuzzy_simplicial_set``, which rejects the metrics it does not know and ignores it anyway when the neighbors are precomputed {pr}`769` {smaller}`S Dicks`
 * Fix {func}`~rapids_singlecell.pp.neighbors` with ``algorithm="all_neighbors"`` and ``algorithm_kwds={"n_clusters": 2}`` always raising; the derived ``overlap_factor`` is now capped at ``n_clusters - 1`` {pr}`769` {smaller}`S Dicks`
+* Fix batched ``all_neighbors`` returning incomplete or inaccurate results on multi-GPU CUDA 12 systems. {pr}`782` {smaller}`S Dicks`
 
 ```{rubric} Features
 ```

src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py

@@ -4,7 +4,9 @@
 from typing import TYPE_CHECKING
 
 import cupy as cp
+import cuvs
 import numpy as np
+from packaging.version import parse as parse_version
 
 from rapids_singlecell.preprocessing._neighbors._helper import _compute_nlist
 
@@ -18,6 +20,9 @@
     from rapids_singlecell.preprocessing._neighbors import _Metrics
 
 
+_CUVS_HOST_OUTPUT_MIN_VERSION = parse_version("26.08")
+
+
 def _default_overlap_factor(n_clusters: int) -> int:
     """Overlap needed to hold recall as the dataset is split into more clusters."""
     if n_clusters <= 1:
@@ -66,14 +71,18 @@ def _all_neighbors_knn(
         )
     algo = algorithm_kwds.get("algo", "nn_descent")
     n_clusters, overlap_factor = _all_neighbors_batching(algorithm_kwds)
-    if n_clusters == 1:
-        from cuvs.common import Resources
-
-        res = Resources()
-    else:
+    use_host_output = (
+        n_clusters > 1
+        and parse_version(cuvs.__version__) >= _CUVS_HOST_OUTPUT_MIN_VERSION
+    )
+    if use_host_output:
         from cuvs.common import MultiGpuResources
 
         res = MultiGpuResources()
+    else:
+        from cuvs.common import Resources
+
+        res = Resources()
     cuvs_metric = "sqeuclidean" if metric == "euclidean" else metric
     if algo == "ivf_pq" or algo == "ivfpq":
         from cuvs.neighbors import ivf_pq
@@ -111,8 +120,11 @@ def _all_neighbors_knn(
         ivf_pq_params=ivf_pq_params,
         nn_descent_params=nn_descent_params,
     )
-    neighbors = cp.zeros([X.shape[0], k], dtype=np.int64)
-    distances = cp.zeros([X.shape[0], k], dtype=np.float32)
+    # Host outputs use cuVS's synchronized merge path. Older releases only accept
+    # device outputs, so keep batching on one GPU to avoid concurrent writes.
+    output_module = np if use_host_output else cp
+    neighbors = output_module.zeros([X.shape[0], k], dtype=np.int64)
+    distances = output_module.zeros([X.shape[0], k], dtype=np.float32)
 
     all_neighbors.build(
         dataset=X,
@@ -122,7 +134,8 @@ def _all_neighbors_knn(
         distances=distances,
         resources=res,
     )
-    neighbors = neighbors.astype(np.int32)
+    neighbors = cp.asarray(neighbors, dtype=np.int32)
+    distances = cp.asarray(distances)
     if metric == "euclidean":
         distances = cp.sqrt(distances)
     return neighbors, distances