NVIDIA/raft
RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
What it solves
RAFT 提供一套 CUDA 加速的基础算法与原语,供机器学习和数据挖掘使用。它通过提供可复用的构建块,解决核心计算重复实现的问题,减少开发时间和维护负担,同时将优化集中于高性能应用。
How it works
RAFT 以 C++ 头文件仅模板库实现,另提供可选的共享库以缩短编译时间并提供运行时 API。它利用 RAPIDS Memory Manager (RMM) 进行内存分配,并使用 mdspan 与 mdarray 来表示多维数据。为了让不同环境都能使用,RAFT 提供两个 Python 包:pylibraft 为运行时 API 的轻量封装,raft-dask 则提供通过 Dask 的多节点、多 GPU 分布式算法基础设施。
Who it’s for
此库面向数据源提供者和应用开发者,需要低层原语来加速工作流。它并非面向进行发现和实验的数据科学家。
Highlights
- Broad Algorithmic Support: 包含密集和稀疏线性代数、求解器、统计以及数据格式转换的加速函数。
- Interoperability:
pylibraft支持__cuda_array_interface__,允许与 CuPy、PyTorch、JAX、TensorFlow 和 RAPIDS cuDF 零拷贝集成。 - Distributed Computing: 提供使用 Dask 在多节点和多 GPU 上构建分布式 GPU 算法的基础设施。
- Flexible C++ Design: 仅头文件模板架构使其易于集成到 C++ 项目中。