Implementation:Rapidsai Cuml KNN Sparse

Knowledge Sources	Rapidsai_Cuml
Domains	Machine_Learning, Nearest_Neighbors
Last Updated	2026-02-08 12:00 GMT

Overview

Provides a GPU-accelerated brute-force k-nearest neighbors search for sparse CSR-format input data, supporting batched computation to manage GPU memory for large sparse datasets.

Description

The knn_sparse.hpp header declares a single function, ML::Sparse::brute_force_knn, which performs exact nearest neighbor search when both the index and query datasets are stored in sparse CSR (Compressed Sparse Row) format.

Key features:

• Accepts separate CSR components (indptr, indices, data) for both the index and query datasets.
• Supports configurable batch sizes for both the index and query to control GPU memory usage when processing large sparse matrices. The default batch size is 2^16 (65536) elements.
• Supports multiple distance metrics via the ML::distance::DistanceType enum, with L2Expanded as the default.
• Outputs dense arrays of KNN indices and distances.

This is particularly useful for text processing, graph-based data, or any domain where feature matrices are naturally sparse and converting to dense format would be prohibitively expensive.

Usage

Use this function when performing nearest neighbor search on sparse data. This is common in NLP (TF-IDF matrices), recommendation systems (user-item matrices), or any domain producing high-dimensional sparse feature vectors. The batch size parameters allow tuning memory usage for very large datasets.

Code Reference

Source Location

• Repository: Rapidsai_Cuml
• File: cpp/include/cuml/neighbors/knn_sparse.hpp

Signature

namespace ML {
namespace Sparse {

constexpr int DEFAULT_BATCH_SIZE = 1 << 16;

void brute_force_knn(raft::handle_t& handle,
                     const int* idx_indptr,
                     const int* idx_indices,
                     const float* idx_data,
                     size_t idx_nnz,
                     int n_idx_rows,
                     int n_idx_cols,
                     const int* query_indptr,
                     const int* query_indices,
                     const float* query_data,
                     size_t query_nnz,
                     int n_query_rows,
                     int n_query_cols,
                     int* output_indices,
                     float* output_dists,
                     int k,
                     size_t batch_size_index = DEFAULT_BATCH_SIZE,
                     size_t batch_size_query = DEFAULT_BATCH_SIZE,
                     ML::distance::DistanceType metric = ML::distance::DistanceType::L2Expanded,
                     float metricArg = 0);

} // namespace Sparse
} // namespace ML

Import

#include <cuml/neighbors/knn_sparse.hpp>

I/O Contract

Inputs

Name	Type	Required	Description
handle	raft::handle_t&	Yes	RAFT handle for GPU resource management
idx_indptr	const int*	Yes	Device pointer to CSR row pointer array for the index matrix, size [n_idx_rows + 1]
idx_indices	const int*	Yes	Device pointer to CSR column indices for the index matrix, size [idx_nnz]
idx_data	const float*	Yes	Device pointer to CSR values for the index matrix, size [idx_nnz]
idx_nnz	size_t	Yes	Number of non-zero entries in the index matrix
n_idx_rows	int	Yes	Number of rows in the index matrix
n_idx_cols	int	Yes	Number of columns in the index matrix
query_indptr	const int*	Yes	Device pointer to CSR row pointer array for the query matrix, size [n_query_rows + 1]
query_indices	const int*	Yes	Device pointer to CSR column indices for the query matrix, size [query_nnz]
query_data	const float*	Yes	Device pointer to CSR values for the query matrix, size [query_nnz]
query_nnz	size_t	Yes	Number of non-zero entries in the query matrix
n_query_rows	int	Yes	Number of rows in the query matrix
n_query_cols	int	Yes	Number of columns in the query matrix
k	int	Yes	Number of nearest neighbors to return
batch_size_index	size_t	No	Batch size for index processing (default: 65536)
batch_size_query	size_t	No	Batch size for query processing (default: 65536)
metric	ML::distance::DistanceType	No	Distance metric (default: L2Expanded)
metricArg	float	No	Metric argument (default: 0)

Outputs

Name	Type	Description
output_indices	int*	Device array of nearest neighbor indices, size [n_query_rows x k]
output_dists	float*	Device array of nearest neighbor distances, size [n_query_rows x k]

Usage Examples

#include <cuml/neighbors/knn_sparse.hpp>

raft::handle_t handle;

// Index sparse matrix in CSR format (on device)
int* d_idx_indptr;    // [n_idx_rows + 1]
int* d_idx_indices;   // [idx_nnz]
float* d_idx_data;    // [idx_nnz]
size_t idx_nnz = 100000;
int n_idx_rows = 5000;
int n_idx_cols = 10000;

// Query sparse matrix in CSR format (on device)
int* d_query_indptr;    // [n_query_rows + 1]
int* d_query_indices;   // [query_nnz]
float* d_query_data;    // [query_nnz]
size_t query_nnz = 2000;
int n_query_rows = 100;
int n_query_cols = 10000;

int k = 10;

// Pre-allocated output arrays on device
int* d_output_indices;    // [100 x 10]
float* d_output_dists;    // [100 x 10]

ML::Sparse::brute_force_knn(handle,
                            d_idx_indptr, d_idx_indices, d_idx_data,
                            idx_nnz, n_idx_rows, n_idx_cols,
                            d_query_indptr, d_query_indices, d_query_data,
                            query_nnz, n_query_rows, n_query_cols,
                            d_output_indices, d_output_dists, k);

Related Pages

• Environment:Rapidsai_Cuml_CUDA_GPU