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

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

왜 중요한가

핵심 기여

FlashAssign 기법 도입

거리 계산과 최솟값 인덱스 추출(argmin)을 하나로 합쳐 중간 거리 행렬 생성을 생략하고 메모리 트래픽을 획기적으로 줄였다.

Sort-Inverse Update 구현

데이터 업데이트 시 발생하는 하드웨어 수준의 쓰기 충돌을 해결하기 위해, 데이터를 클러스터 ID별로 정렬하여 순차적으로 처리하는 방식을 도입했다.

알고리즘-시스템 공동 설계

GPU 메모리를 초과하는 대용량 데이터를 처리하기 위한 파이프라인 설계와 하드웨어 특성에 맞춘 자동 최적화 컴파일러 기법을 적용했다.

압도적인 벤치마크 성능

NVIDIA H200 GPU에서 기존 최적화 라이브러리인 cuML 대비 33배, FAISS 대비 200배 이상의 속도 향상을 달성했다.

핵심 아이디어 이해하기

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 환경에서 별도의 수동 튜닝 없이 최적의 성능을 보장한다.

실무 활용

대규모 벡터 검색, LLM의 KV 캐시 압축, 비디오 생성 모델의 토큰 관리 등 실시간성이 중요한 최신 AI 파이프라인에 즉시 적용 가능한 고성능 클러스터링 엔진이다.

LLM 추론 시 긴 컨텍스트를 효율적으로 관리하기 위한 KV 캐시 압축 및 클러스터링
대규모 벡터 데이터베이스의 인덱싱 및 검색 속도 가속화
생성형 비디오 모델에서 의미론적 토큰 그룹화 및 연산 최적화
실시간 데이터 스트림의 이상 탐지 및 패턴 분류

코드 공개 여부: 공개

코드 저장소 보기

키워드

K-Means(K-평균 클러스터링)GPU Optimization(GPU 최적화)Memory Efficiency(메모리 효율성)FlashAssign(플래시 어사인)Vector Quantization(벡터 양자화)

용어 해설

고대역폭 메모리(HBM): — GPU 칩 옆에 배치되어 매우 높은 데이터 전송 속도를 제공하는 메모리 기술이다. K-Means 연산 시 거대한 거리 행렬을 이곳에 쓰고 읽는 과정에서 병목 현상이 발생하여 전체 성능을 저하시키는 원인이 된다.
원자적 연산(Atomic Operation): — 여러 스레드가 동시에 같은 메모리 주소의 값을 수정할 때 데이터 무결성을 보장하기 위해 순차적으로 처리하는 방식이다. K-Means의 업데이트 단계에서 많은 데이터가 동일한 클러스터 중심을 수정하려 할 때 심각한 대기 시간을 유발한다.
타일링(Tiling): — 거대한 데이터를 작은 블록(타일) 단위로 나누어 GPU 내부의 빠른 SRAM에 올린 뒤 연산하는 기법이다. 메모리 계층 구조를 효율적으로 활용하여 외부 메모리(HBM) 접근 횟수를 최소화하는 데 필수적이다.
최솟값 인덱스 추출(Argmin): — 함수나 배열에서 가장 작은 값을 가지는 위치(인덱스)를 찾는 연산이다. K-Means에서는 각 데이터 포인트가 어떤 클러스터 중심과 가장 가까운지 결정할 때 사용된다.
정적 랜덤 액세스 메모리(SRAM): — GPU 내부 연산 유닛 바로 옆에 위치한 매우 빠르고 작은 저장 공간이다. HBM보다 용량은 작지만 대역폭이 훨씬 높아, 연산 중 발생하는 중간 데이터를 일시적으로 보관하고 처리하는 데 사용된다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

FlashAttention이 거대한 어텐션 행렬을 메모리에 저장하지 않고 계산하는 원리.

CUDA Agent, CUDA 커널 최적화 성능 입증

KV cache를 검색 가능한 벡터 공간으로 라우팅

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

FlashAttention이 거대한 어텐션 행렬을 메모리에 저장하지 않고 계산하는 원리.

CUDA Agent, CUDA 커널 최적화 성능 입증

KV cache를 검색 가능한 벡터 공간으로 라우팅