RAFT: Reusable Accelerated Functions and Tools

Contents

---

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

---

有用资源

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

RAFT 是什么？

RAFT 包含机器学习和数据挖掘中广泛使用的基础算法和原语（primitives）。这些算法经过 CUDA 加速，构成了编写高性能应用程序的基础模块。

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

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

虽然不全面，但以下一般类别有助于总结 RAFT 中的加速函数：

类别	RAFT 中的加速函数
数据格式	稀疏与稠密、转换、数据生成
稠密操作	线性代数、矩阵和向量运算、归约、切片、范数、分解、最小二乘法、SVD 及特征值问题
稀疏操作	线性代数、特征值问题、切片、范数、归约、分解、对称化、组件与标记
求解器	组合优化、迭代求解器
统计	采样、矩与汇总统计、指标、模型评估
工具与实用程序	用于开发 CUDA 应用的通用工具和实用程序、多节点多 GPU 基础设施

RAFT 是一个 C++ 仅头文件模板库，带有一个可选的共享库，该库

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

除了作为 C++ 库之外，RAFT 还提供了 2 个 Python 库：

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

RAFT 适合我吗？

RAFT 包含用于加速应用程序和工作流的底层原语。数据源提供者和应用开发人员可能会发现特定工具非常有用。RAFT 并不打算供数据科学家直接用于发现和实验。对于数据科学工具，请参阅 RAPIDS 网站。

入门指南

RAPIDS 内存管理器 (RMM)

RAFT 严重依赖 RMM，它减轻了在使用它的库中全局配置不同分配策略的负担。

多维数组

RAFT 中的 API 接受 mdspan 多维数组视图来表示更高维度的数据，类似于 Python 库 Numpy 中的 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 流、流池以及像 cublas 和 cusolver 等其他 CUDA 库的句柄。

下面的示例展示了创建一个 RAFT 句柄，并将其与 device_matrix 和 device_vector 一起使用以分配内存、生成随机簇，并使用 NVIDIA cuVS 库计算成对欧几里得距离：

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

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 = cuvs::distance::DistanceType::L2SqrtExpanded;
cuvs::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);

Python 示例

pylibraft 包包含了用于 RAFT 算法和原语（primitives）的 Python API。pylibraft 非常轻量且依赖极少，能够很好地集成到其他库中，并支持任何实现了 __cuda_array_interface__ 接口的对象，例如 CuPy 的 ndarray。此包中暴露的 RAFT 算法数量正在持续增加。

下面的示例演示了如何使用 NVIDIA cuVS 库计算 CuPy 数组之间的成对欧氏距离。注意，CuPy 不是 pylibraft 的必要依赖项。

import cupy as cp

from cuvs.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 的零拷贝（zero-copy）转换成为可能。

下面是将输出 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")  # 所有计算 API 将返回 cupy 数组
pylibraft.config.set_output_as("torch") # 所有计算 API 将返回 torch 张量

您还可以指定一个可调用对象（callable），它接受 pylibraft.common.device_ndarray 并执行自定义转换。以下示例将所有输出转换为 numpy 数组：

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

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

import cupy as cp

from cuvs.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

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

• libraft-headers C++ 头文件
• pylibraft（可选）Python 库
• raft-dask（可选）用于在 Dask 集群中部署使用 RAFT raft::comms 抽象层的多节点多 GPU 算法的 Python 库。

根据您的 CUDA 版本使用以下命令通过 conda 安装所有 RAFT 软件包（将 rapidsai 替换为 rapidsai-nightly 以安装更新但稳定性稍差的夜间构建版本）。推荐使用 mamba 命令而非 conda。

# CUDA 13
mamba install -c rapidsai -c conda-forge raft-dask pylibraft cuda-version=13.1

# CUDA 12
mamba install -c rapidsai -c conda-forge raft-dask pylibraft cuda-version=12.9

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

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

# CUDA 13
mamba install -c rapidsai -c conda-forge libraft libraft-headers cuda-version=13.1

# CUDA 12
mamba install -c rapidsai -c conda-forge libraft libraft-headers cuda-version=12.9

通过 Pip 安装 Python

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

# CUDA 13
pip install pylibraft-cu13
pip install raft-dask-cu13

# CUDA 12
pip install pylibraft-cu12
pip install raft-dask-cu12

这些软件包静态构建了 RAFT 的预编译实例，因此 C++ 头文件不会直接可用于您的代码中。

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

贡献

如果您对为 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}
}

RAFT 快速上手指南

RAFT (Reusable Accelerated Functions and Tools) 是一个用于机器学习和数据挖掘的 CUDA 加速算法与基础工具库。它提供 C++ 头文件模板库及轻量级 Python 封装，旨在加速算法构建并提高性能。

环境准备

• 硬件要求: NVIDIA GPU (支持 CUDA)。
• 操作系统: Linux / Windows / macOS (推荐 Linux)。
• 软件依赖:
  • CUDA Toolkit (版本 12.x 或 13.x)。
  • Conda 或 Mamba (推荐使用 Mamba 以加快安装速度)。
  • Python (3.8+)。
• 前置知识: 熟悉 Python 编程及基本的 CUDA 概念。

安装步骤

推荐使用 mamba 进行安装，比 conda 更快且更稳定。请根据您的 CUDA 版本选择对应的命令。

通过 Conda/Mamba 安装 (推荐)

此方式将安装 C++ 头文件、核心库以及 Python 接口。

# CUDA 13 环境
mamba install -c rapidsai -c conda-forge raft-dask pylibraft cuda-version=13.1

# CUDA 12 环境
mamba install -c rapidsai -c conda-forge raft-dask pylibraft cuda-version=12.9

注意: 上述命令会自动安装 libraft-headers 和 libraft。如需仅安装 C++ 头文件和预编译库，可移除 raft-dask 和 pylibraft。

通过 Pip 安装 (实验性)

如果您偏好使用 pip，可以使用以下命令安装实验性包：

# CUDA 13
pip install pylibraft-cu13
pip install raft-dask-cu13

# CUDA 12
pip install pylibraft-cu12
pip install raft-dask-cu12

基本使用

RAFT 提供了丰富的 Python 接口 (pylibraft)，支持与 CuPy、PyTorch、RAPIDS cuDF 等库无缝互操作。

Python 示例：计算成对距离

以下示例展示了如何使用 RAFT 基础设施（通过 cuvs 调用）计算两个 CuPy 数组之间的欧几里得距离。

import cupy as cp

from cuvs.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")

配置输出类型

pylibraft 允许您全局配置计算结果的返回类型，以便与其他库（如 PyTorch 或 NumPy）无缝集成。

import pylibraft.config

# 设置所有计算 API 返回 CuPy 数组
pylibraft.config.set_output_as("cupy") 

# 设置所有计算 API 返回 PyTorch Tensor
pylibraft.config.set_output_as("torch") 

手动转换数据类型

您也可以手动将 RAFT 的输出转换为其他格式：

# 转换为 CuPy 数组
cupy_array = cp.asarray(output)

# 转换为 PyTorch Tensor
import torch
torch_tensor = torch.as_tensor(output, device='cuda')

# 转换为 NumPy 数组
numpy_array = output.copy_to_host()