FlashML-org/flashlib

解决的问题

FlashLib 提供了经典机器学习算子（如聚类、降维、回归）的高性能 GPU 实现。这些算子通常比现代深度学习内核更慢或效率更低。该库旨在将「Flash」风格的优化效率带入传统机器学习任务中。

工作原理

基于 Triton 和 CuteDSL 构建，该库将大量基础算子实现为顶层函数和类似 scikit-learn 的类。它包含多种专用的近似最近邻（ANN）索引，例如 IVFFlat（高召回率）、IVFPQ（内存压缩）、CAGRA（使用邻近图上的融合贪心遍历实现高吞吐搜索）。

适用人群

需要在大规模数据集上使用 GPU 运行经典机器学习算法以减少计算时间和内存开销的数据科学家和机器学习工程师。

主要亮点

• 广泛覆盖：涵盖聚类、最近邻、分解、流形学习、回归和分类等领域的 18 个高级算子。
• 优化的 ANN 搜索：提供多种索引策略（IVF-Flat、IVF-PQ、CAGRA），以在速度、召回率和内存使用之间取得平衡。
• 多精度 GEMM：提供一系列 GEMM 变体（如 TF32、BF16、FP16、Int8），构成帕累托前沿，支持灵活的精度与性能选择。
• 资源估算：flashlib.info 子模块可在无需 GPU 或重导入的情况下，仅在 CPU 上预测运行时间、FLOPs 和 HBM 字节数。