FlashML-org/flashlib

해결하는 문제

FlashLib는 클러스터링, 차원 축소, 회귀와 같은 전통적 기계학습 연산자에 대한 고성능 GPU 구현을 제공합니다. 이러한 연산들은 종종 현대 딥러닝 커널보다 느리거나 효율적이지 않습니다. 이 라이브러리는 전통적인 ML 작업에 'Flash' 스타일 최적화의 효율성을 가져오려는 목표를 가지고 있습니다.

작동 방식

Triton과 CuteDSL을 기반으로 구축되었으며, 다양한 기본 연산을 상위 수준 함수와 scikit-learn 스타일 클래스로 구현합니다. 고정확도를 원하는 IVFFlat(고 재현율), 메모리 압축을 위한 IVFPQ(메모리 압축), 근접 그래프에서 융합 그리디 탐색을 사용한 고처리량 검색을 위한 CAGRA와 같은 특화된 근사 최근접 이웃(ANN) 인덱스를 포함합니다.

대상 사용자

대규모 데이터셋에서 GPU를 사용하여 전통적 ML 알고리즘을 실행해야 하는 데이터 과학자 및 ML 엔지니어.

주요 특징

• 포괄적인 커버리지: 클러스터링, 최근접 이웃, 분해, 다양체 학습, 회귀, 분류 분야에 걸쳐 18개의 고수준 기본 연산을 포함합니다.
• 최적화된 ANN 검색: 속도, 재현율, 메모리 사용량의 균형을 맞추기 위해 IVF-Flat, IVF-PQ, CAGRA와 같은 다양한 인덱싱 전략을 제공합니다.
• 다중 정밀도 GEMM: TF32, BF16, FP16, Int8 등의 GEMM 변형을 파레토 최적 경계에 따라 제공하여 정밀도와 성능의 유연한 선택을 가능하게 합니다.
• 자원 예측: GPU나 무거운 임포트 없이 CPU에서 실행 시간, FLOPs, HBM 바이트 수를 예측할 수 있는 flashlib.info 서브모듈을 제공합니다.