k-means는 데이터를 k개 그룹으로 나누는 대표적인 클러스터링 알고리즘입니다. 각 데이터 포인트를 가장 가까운 중심점에 할당하고, 같은 클러스터에 속한 점들의 평균으로 중심점을 갱신하는 과정을 반복합니다.1
Flash-KMeans는 이 표준 Lloyd k-means를 GPU 환경에 맞게 다시 구현한 오픈소스 패키지입니다. Triton GPU 커널 기반의 IO-aware 배치 k-means 구현체이며, PyPI에서 pip install flash-kmeans로 설치할 수 있습니다.2
핵심은 알고리즘 자체를 새로 만든 것이 아니라, GPU에서 k-means가 느려지는 지점을 메모리 관점에서 다시 설계했다는 점입니다.
병목은 거리 계산보다 거리 행렬이었다
일반적인 GPU k-means 구현은 모든 데이터 포인트와 모든 중심점 사이의 거리를 계산한 뒤, N × K 크기의 거리 행렬을 GPU 메모리에 저장합니다. 이후 이 행렬을 다시 읽어 각 점에 가장 가까운 중심점을 찾습니다.
문제는 이 중간 거리 행렬이 매우 크다는 것입니다. 데이터 포인트 수 N과 클러스터 수 K가 커질수록, 실제 연산보다 거리 행렬을 쓰고 읽는 메모리 이동 비용이 더 큰 병목이 될 수 있습니다.
Flash-KMeans의 FlashAssign은 이 과정을 바꿉니다. 거리 계산과 최솟값 탐색(argmin)을 하나의 커널 안에서 처리해 중간 거리 행렬을 구체화하지 않습니다.3 필요한 것은 모든 거리값이 아니라, 각 점에 대해 가장 가까운 중심점과 그 거리입니다. 따라서 거대한 임시 행렬을 만들지 않고도 같은 Lloyd k-means 할당 단계를 수행할 수 있습니다.
이 접근은 FlashAttention 계열의 아이디어와 닮아 있습니다. 빠른 GPU에서는 “얼마나 많이 계산하느냐”만큼이나 “얼마나 불필요한 메모리 이동을 줄이느냐”가 중요합니다.
중심점 갱신에서는 atomic 경합을 줄인다
k-means의 두 번째 병목은 중심점 갱신 단계에서 나타납니다. 할당이 끝나면 같은 클러스터에 속한 점들을 모두 더해 평균을 내야 합니다. 단순한 GPU 구현에서는 여러 스레드가 같은 중심점에 동시에 값을 더하려고 하면서 atomic write 경합이 발생할 수 있습니다.
Flash-KMeans는 Sort-Inverse Update로 이 문제를 줄입니다.3 할당 결과를 클러스터 기준으로 정렬한 뒤, 흩어진 쓰기를 보다 규칙적인 세그먼트 단위 합산으로 바꿉니다.
즉, Flash-KMeans의 개선점은 “다른 k-means”가 아니라 “GPU가 잘 처리하는 데이터 흐름으로 재배치한 k-means”에 가깝습니다. 표준 Lloyd k-means의 구조는 유지하되, 메모리 접근과 쓰기 경합을 줄이는 데 초점을 둡니다.
성능 수치는 조건과 함께 읽어야 한다
Flash-KMeans는 NVIDIA H200에서 FAISS 대비 200배 이상 빠른 성능을 보고했습니다.3 PyPI에 공개된 벤치마크도 H200 GPU, FP16 정밀도, 128차원 데이터 조건에서 클러스터 수 k, 데이터 포인트 수 n, 배치 크기 b를 바꿔 비교했다고 설명합니다.2
비교 대상에는 fast_pytorch_kmeans, fastkmeans의 Triton 구현, fastkmeans의 PyTorch fallback, 순진한 batched torch k-means 등이 포함됐습니다.2
다만 이런 숫자는 그대로 모든 환경에 적용할 수 있는 보편 성능표가 아닙니다. GPU 종류, 데이터 차원, 정밀도, N과 K의 크기, 배치 처리 방식에 따라 병목은 달라질 수 있습니다. Flash-KMeans를 평가할 때는 “내 작업이 거리 행렬 생성과 중심점 갱신에서 실제로 병목을 겪는가”를 먼저 확인하는 편이 좋습니다.
어디에 유용한가
Flash-KMeans는 특히 k-means가 파이프라인 내부에서 반복 호출되는 작업에 적합합니다. 예를 들어 대규모 배치 클러스터링, 벡터 양자화, 희소 라우팅처럼 같은 유형의 클러스터링을 여러 번 수행하는 경우에는 메모리 이동 감소 효과가 누적될 수 있습니다.
사용 측면에서도 일반적인 GPU 텐서 워크플로에 맞춰져 있습니다. PyPI 문서에 따르면 batch_kmeans_Euclid 함수로 CUDA 텐서 기반 배치 유클리드 k-means를 실행할 수 있고, FAISS 및 scikit-learn과 유사한 API 인터페이스도 제공합니다.2 또한 Sparse VideoGen2의 공식 k-means 구현체로 제공됩니다.2
대형 입력에 대한 고려도 있습니다. Flash-KMeans는 GPU 메모리에 한 번에 들어가지 않는 데이터를 위해 CPU pinned memory에서 데이터를 생성하고, 청크 단위로 GPU에 전송하는 벤치마크를 제공합니다. 이 대형 텐서 벤치마크는 FP32, 128차원 데이터, N=2^18부터 2^28까지의 규모를 사용했습니다.2
결론: 오래된 알고리즘도 병목을 다시 보면 달라진다
Flash-KMeans가 보여주는 교훈은 분명합니다. k-means 자체는 오래된 알고리즘이지만, GPU에서의 성능은 구현 방식에 크게 좌우됩니다. 특히 대규모 데이터에서는 거리 행렬을 만들고 읽는 비용, 중심점 갱신에서의 쓰기 경합이 핵심 병목이 될 수 있습니다.
따라서 Flash-KMeans는 “더 똑똑한 k-means”라기보다 “GPU 메모리 구조를 제대로 이해한 k-means”에 가깝습니다. k-means를 대규모 파이프라인에서 반복적으로 사용한다면, 알고리즘 선택만큼이나 데이터 이동을 줄이는 구현 선택이 중요하다는 점을 보여주는 사례입니다.
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.