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 位元組數。