UC 버클리와 UT 오스틴 연구팀이 GPU 환경에서 정확한 K-평균(K-Means) 클러스터링을 대폭 가속하는 오픈소스 라이브러리 Flash-KMeans를 공개했다. NVIDIA H200 GPU 기준으로 업계 표준 벡터 검색 라이브러리 FAISS 대비 200배 이상, NVIDIA cuML 대비 33배, 기타 최적화 기반선 대비 최대 17.9배의 처리 속도를 달성했다고 연구팀이 밝혔다. Apache 2.0 라이선스로 공개되며 `pip install flash-kmeans`로 설치할 수 있다.
Flash-KMeans의 핵심 차별점은 수학 알고리즘을 변경하거나 근사값을 쓰지 않는다는 점이다. 표준 로이드(Lloyd’s) K-평균과 동일한 연산을 수행하면서, 데이터를 GPU에서 이동시키는 방식만 재설계했다. 기존 K-평균의 두 가지 주요 병목은 할당 단계에서 HBM(고대역폭 메모리)에 N×K 거리 행렬 전체를 기록했다가 다시 읽는 것, 그리고 업데이트 단계에서 여러 스레드가 동일한 클러스터 중심을 동시에 갱신하려는 경쟁으로 발생하는 아토믹 충돌이었다. 연구팀은 FlashAttention에서 착안한 FlashAssign으로 거리 행렬 구체화를 없애고 IO 복잡도를 O(NK)에서 O(Nd+Kd)로 낮춰 할당 속도를 최대 21.2배 개선했다. 업데이트 단계에는 Sort-Inverse Update를 도입해 할당 벡터를 정렬한 뒤 연속 세그먼트 단위로 처리함으로써 아토믹 연산 충돌을 줄이고 최대 6.3배 속도를 높였다.
벤치마크는 H200 GPU에서 CUDA 12.8, FP16 데이터 기준으로 진행됐다. N=8백만, K=1024 조건에서 최고 성능 기준선 대비 17.9배 향상을 기록했고, N=1M, K=8192 조건에서 FlashAssign 커널만 21.2배를 달성했다. 10억 개 포인트(K=32768, d=128) 규모의 아웃오브코어 처리에서는 반복당 41.4초를 기록해 기존 라이브러리의 261.8초와 비교된다. 또한 커널 튜닝 오버헤드를 캐시 인식 컴파일 휴리스틱으로 최대 175배 단축하면서도 최적화 속도 대비 0.3% 이내의 성능을 유지한다고 연구팀은 설명했다.
Flash-KMeans의 실용적 가치는 훈련·추론 루프 내부에서 K-평균을 반복 호출하는 현대 AI 파이프라인에서 두드러진다. 스파스 어텐션 라우팅, KV캐시 압축, 저비트 양자화 코드북 생성, Diffusion Transformer의 토큰 클러스터링 등에서 밀리초 단위 클러스터링이 가능해진다고 연구팀은 밝혔다. FAISS가 주요 프로덕션 벡터 검색 시스템의 인덱스 구성에 폭넓게 쓰이는 상황에서, Flash-KMeans의 속도가 실서비스에서 검증될 경우 AI 인프라 인덱스 갱신 주기와 비용 구조에 직접적인 영향을 줄 수 있다.
