给 NdArray 装上 CUDA 的轮子
原文中文,约9400字,阅读约需23分钟。发表于: 。Ndarry是Rust编程语言中的一个高性能多维、多类型数组库。它提供了类似 numpy 的多种多维数组的算子。与 Python 相比 Rust 生态缺乏类似 CuPy, Jax 这样利用CUDA 进行加速的开源项目。虽然 Hugging Face 开源的 candle 可以使用 CUDA backend 但是 candle 项瞄准的是大模型的相关应用。...
Ndarry是Rust编程语言中的高性能多维、多类型数组库,类似numpy。作者使用cublas库实现了在CUDA上运行的矩阵乘法函数,并通过FFI连接到Rust中。作者还实现了一个trait来给NdArray数组添加cuda_dot方法,实现了1D和2D矩阵之间的点乘。cublas在高维度矩阵乘法场景下展现出巨大优势。