RAFT: 用于向量搜索和更多功能的可重用加速函数和工具

[!IMPORTANT] RAFT中的向量搜索和聚类算法正在迁移到一个新的名为cuVS的向量搜索专用库。我们将继续在此迁移期间支持RAFT中的向量搜索算法，但将在RAPIDS 24.06（6月）发布后不再更新它们。我们计划在RAPIDS 24.08（8月）发布前完成迁移。

内容

1. 有用资源
2. 什么是RAFT？
3. 用例
4. RAFT适合我吗？
5. 入门
6. 安装RAFT
7. 代码库结构和内容
8. 贡献
9. 参考文献

有用资源

• RAFT参考文档: API文档。
• RAFT入门: 开始使用RAFT。
• 构建和安装RAFT: 安装和构建RAFT的说明。
• 示例笔记本: 示例Jupyter笔记本
• RAPIDS社区: 获取帮助、贡献和协作。
• GitHub仓库: 下载RAFT源代码。
• 问题跟踪器: 报告问题或请求功能。

什么是RAFT？

RAFT包含机器学习和信息检索中广泛使用的算法和原语。这些算法是CUDA加速的，并形成构建高性能应用程序的构建块。

通过采用基于原语的算法开发方法，RAFT

• 加速算法构建时间
• 通过最大化跨项目重用，减少维护负担，
• 集中核心可重用计算，使未来的优化能惠及所有使用它们的算法。

以下一般类别虽然不是详尽的，但有助于总结RAFT中的加速函数

RAFT是一个仅包含C++头文件的模板库，包含一个可选的共享库，

1. 可以加快常见模板类型的编译时间，
2. 并提供主机可访问的“运行时”API，使用时不需要CUDA编译器

除了是C++库之外，RAFT还提供了2个Python库

• pylibraft - RAFT的可访问“运行时”API的轻量级Python包装器。
• raft-dask - 用于在GPU上使用Dask构建分布式算法的多节点多GPU通信基础设施。

用例

向量相似度搜索

RAFT包含GPU上近似最近邻搜索（ANNS）算法的最新实现，例如

• 暴力搜索。不使用索引执行暴力最近邻搜索。
• IVF-Flat和IVF-PQ。使用倒排文件索引结构将内容映射到其位置。IVF-PQ还使用产品量化来减少向量的内存使用。这些方法最初由FAISS库推广。
• CAGRA（Cuda Anns GRAph-based）。使用优化的快速ANNS图构建和搜索实现，适用于GPU。CAGRA在大批量查询、单个查询和图构建时间上优于最先进的CPU方法（即HNSW）。

使用RAFT ANNS算法加速向量搜索的项目包括：Milvus、Redis和Faiss。

请参阅示例Jupyter笔记本，以开始使用Python中的RAFT进行向量搜索。

信息检索

RAFT包含一组可重用原语，用于组合需要快速邻域计算的算法，例如

1. 计算向量之间的距离和计算核语法矩阵
2. 执行球半径查询以构建epsilon邻域
3. 聚类点以划分空间以进行更小和更快的搜索
4. 从密集向量构建邻域“连通性”图

机器学习

RAFT的原语被用于几个RAPIDS库，包括cuML、cuGraph和cuOpt，以构建许多端到端的机器学习算法，这些算法跨越了广泛的不同的应用程序，包括

• 数据生成
• 模型评估
• 分类和回归
• 聚类
• 流形学习
• 降维。

RAFT还用于流行的协同过滤库implicit，用于推荐系统。

RAFT适合我吗？

RAFT 包含用于加速应用程序和工作流程的低级原语。数据源提供者和应用程序开发者可能会发现特定的工具，如 ANN 算法，非常有用。RAFT 并不打算被数据科学家直接用于探索和实验。有关数据科学工具，请参阅 RAPIDS 网站。

入门指南

RAPIDS 内存管理器 (RMM)

RAFT 严重依赖 RMM，这使得配置全局范围内使用它的库的不同分配策略变得容易。

多维数组

RAFT 的 API 接受 mdspan 多维数组视图，用于表示类似于 Numpy Python 库中的 ndarray 的高维数据。RAFT 还包含相应的拥有 mdarray 结构，这简化了主机和设备（GPU）内存中多维数据的分配和管理。

mdarray 在 RMM 上形成了一个便利层，并可以使用许多不同的辅助函数在 RAFT 中构建。

#include <raft/core/device_mdarray.hpp>

int n_rows = 10;
int n_cols = 10;

auto scalar = raft::make_device_scalar<float>(handle, 1.0);
auto vector = raft::make_device_vector<float>(handle, n_cols);
auto matrix = raft::make_device_matrix<float>(handle, n_rows, n_cols);

C++ 示例

RAFT 中的大多数原语都接受一个 raft::device_resources 对象来管理创建代价高昂的资源，例如 CUDA 流、流池以及其他 CUDA 库（如 cublas 和 cusolver）的句柄。

以下示例演示了创建 RAFT 句柄，并使用 device_matrix 和 device_vector 来分配内存、生成随机簇以及计算成对欧几里得距离。

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdarray.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

auto input = raft::make_device_matrix<float, int>(handle, n_samples, n_features);
auto labels = raft::make_device_vector<int, int>(handle, n_samples);
auto output = raft::make_device_matrix<float, int>(handle, n_samples, n_samples);

raft::random::make_blobs(handle, input.view(), labels.view());

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);

还可以创建 raft::device_mdspan 视图，以使用原始指针和形状信息调用相同的 API。

#include <raft/core/device_resources.hpp>
#include <raft/core/device_mdspan.hpp>
#include <raft/random/make_blobs.cuh>
#include <raft/distance/distance.cuh>

raft::device_resources handle;

int n_samples = 5000;
int n_features = 50;

float *input;
int *labels;
float *output;

...
// Allocate input, labels, and output pointers
...

auto input_view = raft::make_device_matrix_view(input, n_samples, n_features);
auto labels_view = raft::make_device_vector_view(labels, n_samples);
auto output_view = raft::make_device_matrix_view(output, n_samples, n_samples);

raft::random::make_blobs(handle, input_view, labels_view);

auto metric = raft::distance::DistanceType::L2SqrtExpanded;
raft::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);

Python 示例

pylibraft 包包含 RAFT 算法和原语的 Python API。pylibraft 通过非常轻量级、依赖性最小以及接受任何支持 __cuda_array_interface__ 的对象（如 CuPy 的 ndarray）很好地与其他库集成。这个包中公开的 RAFT 算法数量在版本发布中持续增长。

以下示例演示了计算 CuPy 数组之间的成对欧几里得距离。请注意，CuPy 不是 pylibraft 的必需依赖项。

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)

output = pairwise_distance(in1, in2, metric="euclidean")

上述示例中的 output 数组类型为 raft.common.device_ndarray，它支持 cuda_array_interface，这使得它可以与其他支持它的库（如 CuPy、Numba、PyTorch 和 RAPIDS cuDF）互操作。CuPy 支持 DLPack，这也使得从 raft.common.device_ndarray 到 JAX 和 Tensorflow 的零拷贝转换成为可能。

以下是将输出 pylibraft.device_ndarray 转换为 CuPy 数组的示例

cupy_array = cp.asarray(output)

以及转换为 PyTorch 张量

import torch

torch_tensor = torch.as_tensor(output, device='cuda')

或者转换为 RAPIDS cuDF 数据框

cudf_dataframe = cudf.DataFrame(output)

当相应的库已安装并可在您的环境中使用时，此转换也可以通过设置全局配置选项由所有 RAFT 计算 API 自动执行

import pylibraft.config
pylibraft.config.set_output_as("cupy")  # All compute APIs will return cupy arrays
pylibraft.config.set_output_as("torch") # All compute APIs will return torch tensors

您还可以指定一个接受 pylibraft.common.device_ndarray 并执行自定义转换的 callable。以下示例将所有输出转换为 numpy 数组

pylibraft.config.set_output_as(lambda device_ndarray: return device_ndarray.copy_to_host())

pylibraft 还支持写入预分配的输出数组，因此任何支持 __cuda_array_interface__ 的数组都可以就地写入

import cupy as cp

from pylibraft.distance import pairwise_distance

n_samples = 5000
n_features = 50

in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
output = cp.empty((n_samples, n_samples), dtype=cp.float32)

pairwise_distance(in1, in2, out=output, metric="euclidean")

安装

RAFT 的 C++ 和 Python 库都可以通过 Conda 安装，Python 库也可以通过 Pip 安装。

通过 Conda 安装 C++ 和 Python

安装 RAFT 最简单的方法是通过 conda，并提供了几个包。

• libraft-headers C++ 头文件
• libraft （可选）包含预编译模板实例化和运行时 API 的 C++ 共享库。
• pylibraft （可选）Python 库
• raft-dask（可选）用于在 Dask 集群中部署使用 RAFT raft::comms 抽象层的多节点多 GPU 算法的 Python 库。
• raft-ann-bench（可选）用于轻松生成基准测试的工具，可以将 RAFT 的向量搜索算法与其他最先进的实现进行比较。
• raft-ann-bench-cpu（可选）与上述类似的可重复基准测试工具，但不需要在机器上安装 CUDA。可以用于在具有强大 CPU 的环境中进行测试。

根据您的 CUDA 版本，使用以下命令使用 conda 安装所有 RAFT 软件包（将 rapidsai 替换为 rapidsai-nightly 以安装更新但稳定性较低的夜间包）。首选 mamba 命令而不是 conda 命令。

# for CUDA 11.8
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=11.8

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia raft-dask pylibraft cuda-version=12.5

注意，上述命令还会安装 libraft-headers 和 libraft。

您也可以使用上述 mamba 命令单独安装 conda 软件包。例如，如果您想安装 RAFT 的头文件和预编译共享库以用于您的项目

# for CUDA 12.5
mamba install -c rapidsai -c conda-forge -c nvidia libraft libraft-headers cuda-version=12.5

如果您正在安装 C++ API，请参阅 使用 libraft 了解有关使用预编译共享库的更多信息。您还可以参考 示例 C++ 模板项目，该模板项目提供了一个可插入您项目的 CMake 配置，并针对上面安装的 RAFT 开发工件进行构建。

通过 Pip 安装 Python

pylibraft 和 raft-dask 都有实验性软件包，可以通过 pip 安装

pip install pylibraft-cu11 --extra-index-url=https://pypi.nvidia.com
pip install raft-dask-cu11 --extra-index-url=https://pypi.nvidia.com

这些软件包会静态构建 RAFT 的预编译实例，因此 C++ 头文件和预编译共享库不会立即可用于您的代码。

构建说明中的 构建说明 包含有关从源代码构建 RAFT 并将其包含到下游项目中的更多详细信息。您还可以在构建说明的 从源代码构建 RAFT C++ 和 Python 部分找到上述 CPM 代码片段的更全面版本。

您可以在 cpp/template 目录中找到示例 RAFT 项目模板，该模板演示了如何使用 RAFT 构建新应用程序或将 RAFT 添加到现有 CMake 项目中。

贡献

如果您有兴趣为 RAFT 项目做出贡献，请阅读我们的 贡献指南。有关开发指南、工作流程和原则的详细信息，请参阅 开发者指南。

参考文献

在一般引用 RAFT 时，请考虑引用此 Github 项目。

@misc{rapidsai,
  title={Rapidsai/raft: RAFT contains fundamental widely-used algorithms and primitives for data science, Graph and machine learning.},
  url={https://github.com/rapidsai/raft},
  journal={GitHub},
  publisher={Nvidia RAPIDS},
  author={Rapidsai},
  year={2022}
}

如果引用稀疏成对距离 API，请考虑使用以下 bibtex

@article{nolet2021semiring,
  title={Semiring primitives for sparse neighborhood methods on the gpu},
  author={Nolet, Corey J and Gala, Divye and Raff, Edward and Eaton, Joe and Rees, Brad and Zedlewski, John and Oates, Tim},
  journal={arXiv preprint arXiv:2104.06357},
  year={2021}
}

如果引用单链接层次聚类 API，请考虑以下 bibtex

@misc{nolet2023cuslink,
      title={cuSLINK: Single-linkage Agglomerative Clustering on the GPU},
      author={Corey J. Nolet and Divye Gala and Alex Fender and Mahesh Doijade and Joe Eaton and Edward Raff and John Zedlewski and Brad Rees and Tim Oates},
      year={2023},
      eprint={2306.16354},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

如果引用 CAGRA，请考虑以下 bibtex

@misc{ootomo2023cagra,
      title={CAGRA: Highly Parallel Graph Construction and Approximate Nearest Neighbor Search for GPUs},
      author={Hiroyuki Ootomo and Akira Naruse and Corey Nolet and Ray Wang and Tamas Feher and Yong Wang},
      year={2024},
      series = {ICDE '24}
}

如果引用 k-选择例程，请考虑以下 bibtex

@proceedings{10.1145/3581784,
    title = {Parallel Top-K Algorithms on GPU: A Comprehensive Study and New Methods},
    author={Jingrong Zhang, Akira Naruse, Xipeng Li, and Yong Wang},
    year = {2023},
    isbn = {9798400701092},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    location = {Denver, CO, USA},
    series = {SC '23}
}

如果引用最近邻下降 API，请考虑以下 bibtex

@inproceedings{10.1145/3459637.3482344,
    author = {Wang, Hui and Zhao, Wan-Lei and Zeng, Xiangxiang and Yang, Jianye},
    title = {Fast K-NN Graph Construction by GPU Based NN-Descent},
    year = {2021},
    isbn = {9781450384469},
    publisher = {Association for Computing Machinery},
    address = {New York, NY, USA},
    url = {https://doi.org/10.1145/3459637.3482344},
    doi = {10.1145/3459637.3482344},
    abstract = {NN-Descent is a classic k-NN graph construction approach. It is still widely employed in machine learning, computer vision, and information retrieval tasks due to its efficiency and genericness. However, the current design only works well on CPU. In this paper, NN-Descent has been redesigned to adapt to the GPU architecture. A new graph update strategy called selective update is proposed. It reduces the data exchange between GPU cores and GPU global memory significantly, which is the processing bottleneck under GPU computation architecture. This redesign leads to full exploitation of the parallelism of the GPU hardware. In the meantime, the genericness, as well as the simplicity of NN-Descent, are well-preserved. Moreover, a procedure that allows to k-NN graph to be merged efficiently on GPU is proposed. It makes the construction of high-quality k-NN graphs for out-of-GPU-memory datasets tractable. Our approach is 100-250\texttimes{} faster than the single-thread NN-Descent and is 2.5-5\texttimes{} faster than the existing GPU-based approaches as we tested on million as well as billion scale datasets.},
    booktitle = {Proceedings of the 30th ACM International Conference on Information \& Knowledge Management},
    pages = {1929–1938},
    numpages = {10},
    keywords = {high-dimensional, nn-descent, gpu, k-nearest neighbor graph},
    location = {Virtual Event, Queensland, Australia},
    series = {CIKM '21}
}