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, **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

  • High accuracy: Approximate k-NN with >95% recall

  • Multi-GPU ready: Supports extreme scale processing

  • GPU-accelerated PCA: cuML PCA/SVD for initialization

Automatic Fallback

Falls back to PyTorch backend if cuVS is not available, ensuring compatibility across different environments.

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_index_type:

Type of 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 *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).

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).

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:

# cuVS with L1 distance for k-NN computation
reducer = DiReCuVS(
    use_cuvs=True,
    metric='(x - y).abs().sum(-1)',  # L1/Manhattan distance
    n_neighbors=32,
    cuvs_index_type='ivf_flat'
)

embedding = reducer.fit_transform(X)

Notes

Requirements: - RAPIDS cuVS: Follow the installation instructions at https://docs.rapids.ai/install/ - CUDA-capable GPU with compute capability >= 6.0

Index Selection Guidelines: - < 50K points: ‘flat’ (exact search) - 50K-500K points: ‘ivf_flat’ - 500K-5M points: ‘ivf_pq’ - > 5M points: ‘cagra’ (if dimensions <= 500)

Memory Considerations: - cuVS requires float32 precision (no FP16 support) - Index building requires additional GPU memory - ‘cagra’ uses more memory but provides best performance for huge datasets

__init__(*args, use_cuvs=None, use_cuml=None, cuvs_index_type='auto', cuvs_build_params=None, cuvs_search_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.

  • **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.

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"