在許多數據分析和機器學習算法中，計算瓶頸往往來自控制端到端性能的一小部分步驟。這些步驟的可重用解決方案通常需要低級別的基元，這些基元非常簡單且耗時。

NVIDIA 制造 RAPIDS RAFT 是為了解決這些瓶頸，并在為多維數據構建算法時最大限度地重用，例如機器學習和數據分析中經常遇到的問題。

RAPIDS 是 GPU 上的一套用于數據科學和機器學習的加速庫：

pandas 類數據結構的 cuDF

cuGraph 用于圖形數據

cuML 用于機器學習

高度優化的 RAFT 計算模式構成了一個豐富的模塊化嵌入式加速器目錄，為您提供了強大的元素來組成新的算法或加速現有的庫。

這僅僅是一個開始：隨著新 GPU 架構的發布， RAFT 組件將繼續優化，確保您始終從硬件中獲得最佳性能。

RAFT 使您能夠花時間設計和開發應用程序，而不必擔心您是否能充分利用 GPU 硬件。

在這篇文章中，我討論了 RAFT 在開發人員工具箱中的位置，使用它的環境，更重要的是，在需要時使用 RAFT 的權力。

消除常見瓶頸

NVIDIA 構建 RAFT ，為開發者提供基本元素作為構建塊。

最近的鄰居就是一個很好的例子。它很常見，很有用，而且計算量很大。鄰域方法涵蓋了 machine learning 的大部分算法，如聚類和降維。支持線性代數、稀疏矩陣運算、采樣、優化和統計矩等應用數學的核心工具。

事實上， RAFT 構成了 RAPIDS cuML 中幾乎所有的算法，包括但不限于流行的 HDBSCAN 、 TSNE 、 UMAP 以及所有其他用于聚類和可視化的算法。

表 1 。 RAFT 包含幾個不同類別的構建塊，包括數據格式、稀疏和密集運算、最近鄰等空間運算、基本聚類、迭代求解器和統計。

如果你想在絕對最快的作品基礎上創作新作品，那么你現在只會受到創造力的限制。除了所有的基本方法外，向量搜索領域的最佳方法是在搜索大型語言模型（ LLM ）和推薦系統方面取得令人興奮的進展。

圖 1 。 NVIDIA H100 GPU 上的新 RAFT IVF-PQ 算法與 CPU 上的 HNSW 算法的比較

圖 1 重點介紹了一種新的 RAFT 最先進的近似最近鄰算法，它是一類被稱為 inverted file indexes （ IVF ）的算法中 product quantization （ PQ ）的變體。 IVF 將訓練數據劃分為一組集群，并將查詢減少到僅最近集群的子集。這是 PQ 算法的 IVF 版本，稱為 IVF-PQ 。

RAFT IVF-PQ 實現已針對生產向量相似性搜索系統所需的小批量查詢進行了調整。這類針對小批量調整的算法中目前最先進的技術被稱為分層可導航小世界圖（ HNSW ）算法。它尤其以在 CPU 上速度快而聞名。新的 RAFT IVF-PQ 算法的訓練速度比 HNSW 快 95 倍，在執行小批量查詢時快 3 倍。

打造始終如一的體驗

RAFT 在其核心是一個只有頭部的 C ++模板庫，它需要一組最小的依賴項。它主要依賴于 CUDA 工具包附帶的庫。這使您能夠靈活地根據需要專門化所需數據類型的模板（圖 2 ）。

圖 2:RAFT 技術堆棧

圖 2 顯示了堆棧從一個僅包含頭的模板庫開始。在純頭庫的基礎上構建了一個共享庫，其中包含預編譯的模板專用化和主機可訪問的運行時 API 。 PyLibRAFT 為 Python 用戶提供了強大的 RAFT 構建塊。

一個可選的共享庫可以通過預編譯常見類型的模板專用化來減少編譯時間。共享庫中還提供了一個運行時 API ，可以鏈接到您的構建中。它使 RAFT API 能夠以類似于從 C 源文件調用 cuBLAS 、 cuSOLVER 和 cuSPARSE API 的方式從常規 C ++源文件調用。

直接使用共享庫中的運行時 API 的是一個名為pylibraft的 Python API ，它包含 C ++ API 的輕量級包裝器。

管理設備資源

RAFT 使用raft::resources對象來管理不同的資源，如 CUDA 流、流池和各種 CUDA 庫的句柄，如 cuBLAS 和 cuSOLVER 。raft::device_resources實例是配置和管理 GPU 特定資源以調用 RAFT API 的最簡單方法。

#include 
raft::device_resources resource_handle;

下面是 Python 中的相同示例：

from pylibraft.common.handle import DeviceResources
resource_handle = DeviceResources()

使用 CUDA 構建復雜的端到端算法傳統上需要低水平的專業知識和關于每個 GPU 體系結構的能力的高級知識，以持續保持硬件繁忙。

像 Cub 、 Thrust 和 CUTLASS 這樣的庫使編寫 CUDA 應用程序變得更加容易。它們將較低級別的 API 抽象為較高級別的原語，這些原語可用于開發各種算法。

RAFT 提供了一個類似的抽象層，但它特別關注 ML 和數據分析的稍微高級的原語。如果你已經熟悉深度神經網絡的嵌入式加速器 CUDA 庫 cuDNN ，你可以說 RAFT 與 RAPIDS cuML 等庫的關系與 cuDNN 與 TensorFlow 等庫相似。

處理多維數據

雖然開始出現與 GPU 交互的標準，但 RAFT C ++ API 基于 C ++ 23 STL 標準中的 mdspan （多維非擁有視圖）。

mdspan是一種非所有權視圖結構，具有高度的表達能力、靈活性和自文檔化。它有一個靈活的 API 來表示多維數據，在精神上類似于 NumPy 的ndarray，但在 C ++中。mdspan提供了干凈和一致的 API 體驗，因為它包裹了任何現有的指針，無論是在主（主機）內存還是設備內存中。

#include 
#include 

std::vector my_floats(10);

… populate vector …

auto my_mdspan = raft::make_host_vector(my_floats.data(), 10);

下面是 Python 中的相同示例：

#include 

float *my_floats;

… allocate and populate device memory …

auto my_mdspan = raft::make_device_vector(my_floats, 10);

為了簡化具有多維表示的內存分配，mdarray標準提供了一個 RAII 兼容的內存，與 RAFT 也采用的mdspan相對應。通過使用mdarray 來分配和包含內存，可以使示例變得更容易。

#include 
#include 

raft::device_resources handle;

int n_rows = 10;
int n_cols = 10;

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

下面是 Python 中的相同示例：

#include 

int n_rows = 10;
int n_cols = 10;

auto scalar = raft::make_host_scalar(1.0);
auto vector = raft::make_host_vector(n_cols);
auto matrix = raft::make_host_matrix(n_rows, n_cols);

在mdarray實例中分配并包含內存后，創建mdspan 視圖以使用.view方法調用 RAFT API 。

// Scalar mdspan on device
auto scalar_view = scalar.view();

// Vector mdspan on device
auto vector_view = vector.view();

// Matrix mdspan on device
auto matrix_view = matrix.view();

使用 C ++和 Python API

例如，使用新的 RAFT 近似最近鄰 API 在 C ++和 Python 中構建和查詢索引。以下代碼示例使用 RAFT IVF-PQ 算法構建索引。

#include 

raft::device_resources handle;
raft::neighbors::ivf_pq::index_params idx_params;
auto index = raft::neighbors::ivf_pq::build(handle, idx_params, dataset);

您可以使用以下代碼示例來查詢新建的索引。

#include 

raft::neighbors::ivf_pq::search_params search_params;

uint32_t n_query_rows = 1000;
uint32_t k = 5;
…

auto out_inds = raft::make_device_matrix(handle, n_query_rows, k);
auto out_dists = raft::make_device_matrix(handle, n_query_rows, k);

raft::neighbors::ivf_pq::search(handle, search_params, index, queries,                                     
            out_inds.view(), out_dists.view());

下面是 Python 中相同的索引構建示例：

import cupy as cp
from pylibraft.neighbors import ivf_pq
n_samples = 50000
n_query_rows = 1000
index_params = ivf_pq.IndexParams(
     n_lists=1024,
     metric="sqeuclidean",
     pq_dim=10)
index = ivf_pq.build(index_params, dataset)

下面是 Python 中相同的索引搜索示例：

search_params = ivf_pq.SearchParams(n_probes=20)
k = 10
…
distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(), index,
                                     queries, k)

PyLibRAFT 可以與任何支持 __cuda_array_interface__ （ CAI ）的庫進行互操作，如 PyTorch 、 CuPy 和 Numba 。默認情況下， PyLibRAFT 函數可以接受任何符合 CAI 的數組，即使在需要時也可以將輸出寫入到位。您可以通過配置執行輸出轉換的函數，將此行為自定義到您選擇的庫中，以實現更無縫的集成。

以下代碼示例將 PyLibRAFT 配置為返回 CuPy ndarray或 PyTorch tensors。

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 函數的所有輸出轉換為主（主機）存儲器中的 NumPy 數組。

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

獲取 RAFT

以下步驟提供了有關安裝 RAFT 的不同方式的信息。 Conda 和 PIP 是最簡單的路線，尤其是在與 PyLibRAFT 交互時。由于 RAFT 是一個核心的 C ++庫，我們還簡要展示了它可以多么容易地集成到 Cmake 中。

康達牌手表

對于 C ++和 Python 使用，安裝 RAFT 的一種簡單方法是使用 conda 。將 C ++標頭安裝到您的 conda 環境中。

conda install -c rapidsai -c conda-forge -c nvidia libraft-headers

將預編譯的二進制文件也安裝到您的環境中。

conda install -c rapidsai -c conda-forge -c nvidia libraft-distance libraft-nn

當然， Python 庫也有 conda 版本：

conda install -c rapidsai -c conda-forge -c nvidia pylibraft

皮普

RAPIDS 還提供了 pip 包，可以在純 Python 環境中輕松安裝。

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

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

C 制造商

RAFT 也可以使用 CMake Package Manager （ CPM ）和 rapids-cmake 輕松集成到您的 cmake 項目中，它們還可以下載所需版本并配置raft::raft cmake 目標。

如果 RAFT 已經安裝，例如之前的 conda 命令，請使用 cmake 的find_package使raft::raft目標可用于您的項目。

find_package(raft)

raft::raft目標攜帶了許多主要 RAFT 功能所需的基本依賴項集，如 RAPIDS 內存管理器（ RMM ）和 CUDA 工具包庫。但是，如果您使用來自以下任何命名空間的 API ，則可能需要額外的依賴項：

raft::distance

raft::neighbors

raft::comms

RAFT 使用可選的 cmake 組件，這使得新的 cmakes 目標能夠將這些依賴關系傳遞到您的項目中。

find_package(raft COMPONENTS distance distributed)

此 cmake 功能可啟用目標 raft :： distance 和 raft :: distributed ，此外還有 raft ：： raft 。將它們添加到必要的構建目標中，以便傳遞包含路徑和依賴關系的鏈接庫信息。例如，可以為距離組件鏈接 CUTRASS ，或者為分布式組件鏈接 NVIDIA Collective Communications Library （ NCCL ）和 Unififed Communications X （ UCX ）。

主要收獲

RAFT 是一個用于機器學習和數據分析的高度可重用計算模式庫。集中化重要計算使您能夠在進行優化時自動獲得性能改進的好處，例如在發布新的 GPU 體系結構和功能時。

RAFT 還包含高度可重復使用的基礎設施，用于構建具有干凈靈活接口的生產質量加速庫，遵循 C ++ STL 等重要標準，以保證使用壽命。

RAFT 通過覆蓋從 C ++到 Python 的整個堆棧來簡化集成，進一步實現了與其他流行的 GPU 加速庫的互操作性，如 PyTorch 、 CuPy 和 Numba ，以及 RAPIDS 生態系統中的庫。