Flash-KMeans: 빠르고 메모리 효율적인 정확한 K-Means 알고리즘

기존 GPU 기반 K-Means는 연산량보다 메모리 입출력과 데이터 업데이트 시 발생하는 충돌 때문에 성능이 제한되었습니다. 이 논문은 메모리 구조를 재설계하여 이 병목을 제거함으로써, 대규모 데이터를 실시간으로 처리해야 하는 최신 AI 시스템에서 K-Means를 핵심 도구로 사용할 수 있게 합니다.

K-Means는 데이터를 비슷한 그룹끼리 묶는 기초적인 알고리즘으로, 각 데이터와 그룹 중심 사이의 거리를 계산하고 가장 가까운 그룹을 선택하는 과정을 반복한다. 기존 GPU 구현은 이 거리 값들을 모두 메모리에 기록한 뒤 다시 읽어와서 최솟값을 찾는데, 이 과정에서 발생하는 막대한 데이터 이동이 전체 속도를 늦추는 '메모리 벽' 문제를 야기한다.

Flash-KMeans는 FlashAttention의 아이디어를 차용하여, 거리 계산 결과가 나오는 즉시 메모리에 저장하지 않고 GPU 내부의 빠른 저장 공간(SRAM)에서 바로 최솟값을 판별한다. 이를 통해 불필요한 메모리 쓰기와 읽기를 완전히 제거하여 물리적인 한계에 가까운 속도를 낸다.

또한, 여러 데이터가 동시에 같은 그룹의 중심점을 업데이트하려고 할 때 발생하는 병목을 해결하기 위해, 데이터를 그룹별로 미리 정렬하여 충돌 없이 매끄럽게 합쳐지도록 구조를 변경했다. 결과적으로 연산 효율뿐만 아니라 하드웨어 자원 활용도를 극대화하여 대규모 AI 워크로드에 적합한 성능을 제공한다.

FlashAssign은 거리 행렬 $D \in \mathbb{R}^{N \times K}$를 메모리에 명시적으로 생성하지 않는 기법이다. [데이터 포인트 $x_i$와 클러스터 중심 $c_k$를 입력으로] → [타일링 기법을 통해 GPU 내부 레지스터에서 즉석으로 거리를 계산하고 온라인 argmin 연산을 수행하여] → [가장 가까운 클러스터 인덱스 $a_i$만 결과로 출력하며] → [중간 단계의 거대한 거리 행렬 저장 비용을 제거한다].

Sort-Inverse Update는 업데이트 단계의 원자적 연산 충돌을 방지한다. [할당된 클러스터 인덱스 벡터 $a$를 입력으로] → [인덱스 순서대로 정렬하여 역매핑을 생성한 뒤] → [동일한 클러스터에 속한 데이터들을 연속된 세그먼트로 묶어 한 번에 합산하며] → [수만 개의 스레드가 동시에 같은 메모리 주소에 쓰려고 경쟁하는 현상을 제거한다].

System Co-design은 GPU VRAM 용량을 초과하는 데이터를 처리하기 위해 Chunked Stream Overlap을 적용했다. [데이터를 청크 단위로 나누어 입력으로] → [CPU에서 GPU로 데이터를 복사하는 동안 이전 청크의 연산을 동시에 진행하는 이중 버퍼링을 수행하여] → [PCIe 대역폭 병목을 숨기고] → [최대 10억 개의 포인트까지 끊김 없이 처리 가능한 파이프라인을 구축한다].

NVIDIA H200 GPU 환경에서 수행된 실험 결과, Flash-KMeans는 기존 최적화 라이브러리인 cuML 대비 최대 33배, FAISS 대비 200배 이상의 속도 향상을 기록했다. 특히 메모리 집약적인 대규모 클러스터링 환경에서 가장 큰 성능 차이를 보였다.

커널 단위 분석에서 FlashAssign은 기존 거리 계산 커널 대비 21.2배, Sort-Inverse Update는 업데이트 커널 대비 6.3배의 속도 향상을 달성했다. 이는 알고리즘의 수학적 결과는 동일하게 유지하면서 구현 방식의 혁신만으로 얻어낸 결과이다.

10억 개의 데이터 포인트를 처리하는 Out-of-core 실험에서 기존 방식들이 메모리 부족으로 실패하거나 매우 느렸던 것과 달리, 제안된 파이프라인은 10.5배의 가속을 유지하며 안정적으로 동작했다. 또한 캐시 기반 컴파일 최적화를 통해 튜닝 시간을 175배 단축했다.

Flash-KMeans는 Lloyd 알고리즘의 수학적 정확성을 유지하면서 GPU 아키텍처에 최적화된 데이터 흐름을 재설계했다. 핵심은 연산 중심이 아닌 메모리 중심 병목을 해결하는 데 집중한 것이다. FlashAssign은 FlashAttention의 온라인 소프트맥스 기법과 유사하게, 전체 데이터를 한 번에 처리하지 않고 타일 단위로 쪼개어 SRAM 내에서 중간 결과인 거리 값을 소모한다. 이는 $O(NK)$의 메모리 복잡도를 $O(Nd + Kd)$ 수준으로 낮추는 효과를 가져온다.

Sort-Inverse Update는 비정형적인 Scatter 연산을 정형적인 Segmented Reduction으로 변환한다. 이는 하드웨어의 L1/L2 캐시 적중률을 높이고, 메모리 컨트롤러 수준에서의 직렬화를 방지하여 이론적 최대 대역폭에 근접한 성능을 낸다. 구현 측면에서는 하드웨어의 캐시 크기와 데이터 형상을 고려한 Heuristic 기반의 컴파일 최적화를 통해 다양한 GPU 환경에서 별도의 수동 튜닝 없이 최적의 성능을 보장한다.