RAFT：可重用加速函数和工具，用于向量搜索等

[!重要] 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 notebook 以开始使用 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（可选）Python库，用于在Dask集群中部署使用RAFT raft::comms抽象层的多节点多GPU算法。
• 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++模板项目，该项目包含可直接插入到您的项目中并针对安装的RAFT开发工件构建的CMake配置。

通过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的更全面版本的从源代码构建RAFT C++和Python部分。

您可以在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}
}