플래시케이맥스 빠르고 메모리 효율적인 정확한 K 평균

초록

플래시케이맥스는 GPU 환경에서 기존 K‑means 구현이 겪는 거리 행렬 물리화와 원자적 업데이트 충돌을 근본적으로 없애는 두 핵심 커널(FlashAssign와 Sort‑Inverse Update)을 제안한다. 스트리밍 방식으로 거리 계산과 argmin을 결합해 N×K 행렬을 메모리에 저장하지 않으며, 클러스터별 정렬 후 구간 단위 축소로 원자적 쓰기 경쟁을 제거한다. 추가로 청크 기반 스트림 오버랩과 캐시‑인식 컴파일 히어스틱을 적용해 대규모·동적 워크로드에서도 높은 처리량과 낮은 레이턴시를 달성한다. 실험 결과, H200 GPU에서 엔드‑투‑엔드 속도가 기존 최적 구현 대비 최대 17.9배, cuML·FAISS 대비 각각 33배·200배 이상 향상된다.

상세 분석

플래시케이맥스는 현대 GPU가 제공하는 대역폭과 텐서 코어 연산을 최대한 활용하면서, 메모리·동기화 병목을 최소화하는 시스템‑알고리즘 공동 설계 접근법을 채택한다. 첫 번째 핵심인 FlashAssign는 기존 구현이 거리 행렬 D∈ℝ^{N×K}를 HBM에 기록하고 다시 읽어오는 과정을 완전히 제거한다. 입력 포인트와 센트로이드를 블록 단위로 스트리밍하면서, 각 포인트에 대해 실시간으로 최소 거리와 해당 클러스터 인덱스를 업데이트한다. 이 과정은 거리 계산을 행렬 곱으로 구현하고, 결과를 즉시 argmin 연산에 전달하므로 메모리 트래픽이 O(NK)에서 O(N) 수준으로 감소한다. 두 번째 핵심인 Sort‑Inverse Update는 클러스터 할당 벡터를 클러스터 ID 기준으로 정렬한 뒤, 정렬된 구간을 순차적으로 축소한다. 이렇게 하면 원자적 add 연산이 동일 클러스터에 집중되는 상황을 피하고, 대신 고대역폭 세그먼트 축소 연산으로 대체된다. 결과적으로 원자적 충돌이 크게 감소하고, 실제 집계 단계에서 2배 이상의 메모리 대역폭을 활용한다. 시스템 차원에서는 청크 단위 스트림 오버랩을 도입해 PCIe 전송과 연산을 겹치게 함으로써 대규모 데이터(예: 10억 포인트) 처리 시 통신 오버헤드를 거의 무시할 수준으로 낮춘다. 또한 캐시‑인식 컴파일 히어스틱은 워크로드 특성(N, K, d, 배치 크기)에 따라 최적 스레드 블록 크기와 메모리 접근 패턴을 자동 선택하도록 설계돼, 수동 튜닝 없이도 최적에 근접한 성능을 보장한다. 종합적으로 볼 때, 플래시케이맥스는 알고리즘 복잡도 자체를 변화시키지 않으면서도 구현 레벨에서 발생하는 IO와 동기화 비용을 근본적으로 제거함으로써, 기존 GPU 기반 K‑means 구현이 직면한 한계를 뛰어넘는다.