cuVS: библиотека для векторного поиска и кластеризации на GPU

cuVS (CUDA Vector Search) — это библиотека от NVIDIA, которая предоставляет современные реализации алгоритмов для поиска приближённых ближайших соседей и кластеризации на графических процессорах (GPU). Она может быть интегрирована с различными базами данных и библиотеками, упрощая использование GPU для поиска по векторным данным и кластеризации.

Основная задача cuVS — упростить использование графических процессоров для задач поисковой обработки векторов и кластеризации, что становится особенно актуальным в свете растущей необходимости в семантическом поиске и работе с многими видами неструктурированных данных. Векторы используются в таких задачах, как кластеризация, визуализация и взаимодействие с графами, открывая доступ к алгоритмам анализа графов.

Технологический стек

cuVS построен на библиотеке RAPIDS RAFT, которая включает высокоэффективные примитивы машинного обучения. Библиотека обеспечивает все необходимые функции для реализации поиска по векторным данным и кластеризации на GPU.

Установка cuVS

cuVS доступен в виде предварительно собранных пакетов, которые можно установить через:

• conda
• pip
• tarball

Обрати внимание, что сборки для CUDA 13 имеют меньший размер по сравнению с предыдущими версиями, и команду разработки стремится уменьшать размеры всех поддерживаемых версий. Если размер бинарников является критически важным, рекомендуется скомпилировать cuVS из исходников или использовать `libcuvs-static` пакет.

Начало работы

Библиотека поддерживает несколько языков программирования. Примеры кода для создания индекса с использованием алгоритма CAGRA представлены ниже:

Python API

from cuvs.neighbors import cagra

dataset = load_data()
index_params = cagra.IndexParams()

index = cagra.build(index_params, dataset)

C++ API

#include 

using namespace cuvs::neighbors;

raft::device_matrix_view dataset = load_dataset();
raft::device_resources res;

cagra::index_params index_params;

auto index = cagra::build(res, index_params, dataset);

C API

#include 

cuvsResources_t res;
cuvsCagraIndexParams_t index_params;
cuvsCagraIndex_t index;

DLManagedTensor *dataset;
load_dataset(dataset);

cuvsResourcesCreate(&res);
cuvsCagraIndexParamsCreate(&index_params);
cuvsCagraIndexCreate(&index);

cuvsDataset_t dataset_view;
cuvsDatasetMakeStandardView(res, dataset, &dataset_view);

cuvsCagraBuild(res, index_params, dataset_view, index);

cuvsDatasetDestroy(dataset_view);
cuvsCagraIndexDestroy(index);
cuvsCagraIndexParamsDestroy(index_params);
cuvsResourcesDestroy(res);

Rust API

use cuvs::distance::DistanceType;
use cuvs::neighbors::cagra::{Index, IndexParams, SearchParams};
use cuvs::{AsDlTensor, AsDlTensorMut, Resources};

// cuVS не зависим от местоположения ваших векторов: `build` и `search` 
// принимают любые типы, реализующие `AsDlTensor` (входные данные) / 
// `AsDlTensorMut` (выходные данные). Запакуйте вашу собственную 
// память GPU, реализовав эти трейты.
struct GpuTensor;
// Реализация трейтов здесь...

fn main() -> Result<(), Box> {
    let res = Resources::new()?;
    // Создание индекса CAGRA...
    Ok(())
}

Ограничения

cuVS ориентирован на использование с GPU, поэтому для эффективной работы нужна соответствующая аппаратная платформа. Библиотека постоянно обновляется, чтобы поддерживать совместимость с новыми архитектурами NVIDIA и версиями CUDA, что обеспечивает высокую производительность.