dire_rapids.dire_cuvs module

DIRE with cuVS backend for GPU-accelerated k-NN at scale.

This module provides optional cuVS integration for massive datasets. Falls back to PyTorch if cuVS is not available.

Requirements:: Follow the installation instructions at https://docs.rapids.ai/install/

class dire_rapids.dire_cuvs.DiReCuVS(*args, use_cuvs=None, use_cuml=None, cuvs_index_type='auto', cuvs_build_params=None, cuvs_search_params=None, cuvs_knn_method='auto', all_neighbors_algo='nn_descent', all_neighbors_n_clusters=1, all_neighbors_overlap_factor=None, all_neighbors_device_ids=None, all_neighbors_algo_params=None, **kwargs)[source]

Bases: DiRePyTorch

RAPIDS cuVS/cuML accelerated implementation of DiRe for massive datasets.

This class extends DiRePyTorch with optional RAPIDS cuVS (CUDA Vector Search) integration for GPU-accelerated k-nearest neighbors computation and cuML integration for GPU-accelerated PCA initialization. It provides substantial performance improvements for large-scale datasets.

Performance Advantages over PyTorch/PyKeOps

10-100x faster k-NN: For large datasets (>100K points)
Massive scale support: Handles 10M+ points efficiently
Tunable accuracy: Exact brute force or approximate graph builders
Multi-GPU ready: Supports extreme scale processing
GPU-accelerated PCA: cuML PCA/SVD for initialization

Automatic Fallback

With knn_backend='auto', falls back to PyTorch backend if cuVS is not available or is not beneficial for the current data. With knn_backend='cuvs', cuVS is required and unavailable/unsupported configurations raise instead of falling back.

param use_cuvs:

Whether to use cuVS for k-NN computation. If None, automatically detected based on availability and hardware.

type use_cuvs:

bool or None, default=None

param use_cuml:

Whether to use cuML for PCA initialization. If None, automatically detected based on availability and hardware.

type use_cuml:

bool or None, default=None

param cuvs_knn_method:

cuVS graph-construction strategy. 'auto' preserves the established index-and-search policy. The purpose-built all-neighbors API is available through explicit 'all_neighbors' opt-in while its downstream embedding-quality validation remains experimental.

type cuvs_knn_method:

{‘auto’, ‘all_neighbors’, ‘index_search’}, default=’auto’

param cuvs_index_type:

Type of legacy cuVS index to build:

‘auto’: Automatically select based on data characteristics
‘ivf_flat’: Inverted file index without compression
‘ivf_pq’: Inverted file index with product quantization
‘cagra’: Graph-based index for very large datasets
‘flat’: Brute-force exact search

type cuvs_index_type:

{‘auto’, ‘ivf_flat’, ‘ivf_pq’, ‘cagra’, ‘flat’}, default=’auto’

param cuvs_build_params:

Custom parameters for cuVS index building. Overrides defaults.

type cuvs_build_params:

dict, optional

param cuvs_search_params:

Custom parameters for cuVS search. Overrides defaults.

type cuvs_search_params:

dict, optional

param all_neighbors_algo:

Local graph builder used by cuVS all-neighbors.

type all_neighbors_algo:

{‘nn_descent’, ‘brute_force’, ‘ivf_pq’}, default=’nn_descent’

param all_neighbors_n_clusters:

Paper parameter c. Values greater than one keep input on the host and process spatial partitions out-of-core.

type all_neighbors_n_clusters:

int, default=1

param all_neighbors_overlap_factor:

Paper spill factor s. None selects 0 for a single cluster and min(2, c - 1) for an out-of-core build.

type all_neighbors_overlap_factor:

int or None, default=None

param all_neighbors_device_ids:

GPU ids for multi-GPU out-of-core construction.

type all_neighbors_device_ids:

sequence of int or None, default=None

param all_neighbors_algo_params:

Overrides for the selected local graph builder’s parameter object.

type all_neighbors_algo_params:

dict, optional

param *args:

Additional arguments passed to DiRePyTorch parent class. Includes: n_components, n_neighbors, init, max_iter_layout, min_dist, spread, cutoff, neg_ratio, verbose, random_state, use_exact_repulsion, metric (custom distance function for k-NN computation), knn_backend.

param **kwargs:

Additional arguments passed to DiRePyTorch parent class. Includes: n_components, n_neighbors, init, max_iter_layout, min_dist, spread, cutoff, neg_ratio, verbose, random_state, use_exact_repulsion, metric (custom distance function for k-NN computation), knn_backend.

use_cuvs

Whether cuVS backend is enabled and available.

Type:: bool

use_cuml

Whether cuML backend is enabled and available.

Type:: bool

cuvs_index

Built cuVS index for k-NN search.

Type:: object or None

Examples

Basic usage with automatic backend selection:

from dire_rapids import DiReCuVS
import numpy as np

# Large dataset
X = np.random.randn(100000, 512)

# Auto-detect cuVS/cuML availability
reducer = DiReCuVS()
embedding = reducer.fit_transform(X)

Force cuVS with custom index parameters:

reducer = DiReCuVS(
    use_cuvs=True,
    cuvs_index_type='ivf_pq',
    cuvs_build_params={'n_lists': 2048, 'pq_dim': 64}
)

Massive dataset processing:

# 10M points, 1000 dimensions
X = np.random.randn(10_000_000, 1000)

reducer = DiReCuVS(
    use_cuvs=True,
    use_cuml=True,
    cuvs_index_type='cagra',  # Best for very large datasets
    n_neighbors=32
)

embedding = reducer.fit_transform(X)

With custom distance metric:

# Custom expressions use the PyTorch k-NN fallback
reducer = DiReCuVS(
    metric='(x - y).abs().sum(-1)',  # L1/Manhattan distance
    n_neighbors=32,
)

embedding = reducer.fit_transform(X)

Notes

Requirements:

RAPIDS cuVS 26.06
CUDA 12.2–12.9 on Volta or newer, or CUDA 13.0–13.2 on Turing or newer

Legacy Index Selection Guidelines:

< 50K points: ‘flat’ (exact search)
50K-500K points, or more than 500 dimensions: ‘ivf_flat’
500K-5M points at up to 500 dimensions: ‘ivf_pq’
> 5M points: ‘cagra’ (if dimensions <= 500)

Memory Considerations:

cuVS all-neighbors requires float32 input
all_neighbors_n_clusters > 1 keeps input in host memory, but the final N x k graph must still fit on one GPU
Larger cluster counts reduce working memory; larger overlap factors improve boundary recall at additional compute cost

__init__(*args, use_cuvs=None, use_cuml=None, cuvs_index_type='auto', cuvs_build_params=None, cuvs_search_params=None, cuvs_knn_method='auto', all_neighbors_algo='nn_descent', all_neighbors_n_clusters=1, all_neighbors_overlap_factor=None, all_neighbors_device_ids=None, all_neighbors_algo_params=None, **kwargs)[source]

Initialize DiReCuVS with cuVS and cuML backend configuration.

Parameters:

*args – Positional arguments passed to DiRePyTorch parent class.
use_cuvs (bool or None, default=None) – Whether to use cuVS for k-NN computation: - None: Auto-detect based on availability and GPU presence - True: Force cuVS usage (raises error if unavailable) - False: Disable cuVS, use PyTorch backend
use_cuml (bool or None, default=None) – Whether to use cuML for PCA initialization: - None: Auto-detect based on availability and GPU presence - True: Force cuML usage (raises error if unavailable) - False: Disable cuML, use sklearn backend
cuvs_index_type ({'auto', 'ivf_flat', 'ivf_pq', 'cagra', 'flat'}, default='auto') –
Type of cuVS index to build:
- ’auto’: Automatically select optimal index based on data size/dimensionality
- ’ivf_flat’: Inverted file index without compression (good balance)
- ’ivf_pq’: Inverted file with product quantization (memory efficient)
- ’cagra’: Graph-based index (best for very large datasets)
- ’flat’: Brute-force exact search (small datasets only)
cuvs_build_params (dict, optional) – Custom parameters for cuVS index building. These override the automatically determined parameters. See cuVS documentation for index-specific parameters.
cuvs_search_params (dict, optional) – Custom parameters for cuVS search operations. These override the automatically determined parameters. See cuVS documentation for index-specific search parameters.
cuvs_knn_method ({'auto', 'all_neighbors', 'index_search'}, default='auto') – Select the purpose-built graph API or the legacy build-index and self-query implementation.
all_neighbors_algo ({'nn_descent', 'brute_force', 'ivf_pq'}, default='nn_descent') – Local graph algorithm for all-neighbors.
all_neighbors_n_clusters (int, default=1) – Number of spatial partitions (paper parameter c).
all_neighbors_overlap_factor (int or None, default=None) – Number of partitions per point (paper spill factor s).
all_neighbors_device_ids (sequence of int or None, default=None) – Device ids used by MultiGpuResources for a batched build.
all_neighbors_algo_params (dict, optional) – Local algorithm parameter overrides.
**kwargs – Additional keyword arguments passed to DiRePyTorch parent class. See DiRePyTorch documentation for available parameters including: n_components, n_neighbors, init, max_iter_layout, min_dist, spread, cutoff, neg_ratio, verbose, random_state, use_exact_repulsion, metric (custom distance function for k-NN computation).

Raises:

ImportError – If cuVS or cuML are requested but not available.
RuntimeError – If GPU is required but not available.

get_diagnostics()[source]: Return fitted pipeline diagnostics, including effective cuVS policy.

fit_transform(X, y=None)[source]

Fit the model and transform data with cuVS/cuML acceleration.

This method extends the parent implementation with intelligent backend selection and logging to inform users about the acceleration being used.

Parameters:

X (array-like of shape (n_samples, n_features)) – High-dimensional input data to transform.
y (array-like of shape (n_samples,), optional) – Ignored. Present for scikit-learn API compatibility.

Returns:

Low-dimensional embedding of the input data.

Return type:

numpy.ndarray of shape (n_samples, n_components)

Notes

Backend Selection Logic: - Uses cuVS for k-NN if dataset is large enough and cuVS is available - Uses cuML for PCA initialization if available and init=’pca’ - Falls back to PyTorch implementations automatically

Performance Benefits: - cuVS k-NN: 10-100x speedup for large datasets - cuML PCA: 5-50x speedup for high-dimensional initialization

Examples

Large dataset with cuVS acceleration:

import numpy as np
from dire_rapids import DiReCuVS

# 500K points, 1000 dimensions
X = np.random.randn(500000, 1000)

reducer = DiReCuVS(verbose=True)  # Will log backend selection
embedding = reducer.fit_transform(X)
# Output: "Using cuVS-accelerated backend for 500000 points"