Apple Silicon 기반 TurboQuant 벡터 검색 구축기: 학습한 내용 공유

합의점 vs 논쟁점

합의점

ARM 환경에서 NEON SIMD 최적화는 성능 향상을 위해 필수적이다
그래프 기반 ANN 검색은 GPU 오버헤드 때문에 단일 쿼리 상황에서 CPU보다 느릴 수 있다

실용적 조언

ARM 기반 시스템에서 벡터 연산 시 NEON SIMD 지원 여부를 반드시 확인해야 한다
데이터셋 학습 과정 없이 빠른 양자화가 필요하다면 데이터 비의존적인 TurboQuant 기법이 유효한 대안이다

섹션별 상세

기존 C++ 기반 NGT 라이브러리를 Rust로 포팅하며 Apple Silicon의 NEON SIMD 명령어를 직접 적용했다. C++ 원본이 ARM에서 스칼라 연산으로 후퇴하는 문제를 해결하기 위해 NEON 내장 함수를 사용해 거리 계산 로직을 최적화했다. 실험 결과 빌드 시간은 1분 49초에서 57초로, 쿼리 지연 시간은 0.272ms에서 0.158ms로 약 1.7배 단축됐다. 이는 현대적인 언어와 하드웨어 특화 최적화의 결합이 레거시 코드보다 월등한 성능을 낼 수 있음을 입증했다.

데이터 비의존적 양자화 기법인 TurboQuant를 도입하여 별도의 k-means 학습 없이 벡터를 압축했다. Walsh-Hadamard Transform(WHT)과 무작위 부호 반전을 통해 모든 벡터 좌표를 가우시안 분포로 변환한 뒤, 미리 계산된 Lloyd-Max 코드북으로 양자화하는 방식이다. 118만 개의 벡터를 양자화하는 데 블록 기반 WHT를 사용하여 처리 시간을 6.2초에서 0.77초로 줄였다. 이 기법은 모델이나 데이터셋에 관계없이 일정한 성능을 보장하며 메모리 사용량을 2.8배 절감했다.

그래프 기반 ANN 검색을 Metal GPU로 가속하려 했으나 CPU 대비 60배 이상 느린 17ms의 지연 시간이 발생했다. 그래프 탐색의 각 단계가 이전 단계의 결과에 의존하는 순차적 구조여서 GPU의 병렬성을 활용하기 어려웠다. 특히 각 단계에서 탐색하는 이웃 노드가 10~40개에 불과해 5~10μs 수준의 GPU 커널 호출 오버헤드를 극복하지 못했다. 단일 쿼리 기반의 데스크톱 RAG 환경에서는 GPU보다 CPU의 SIMD 연산이 훨씬 효율적이라는 결론에 도달했다.

cpp

kernel void tq_batch_distance(
device const float* query_rot,
device const uchar* codes, 
device const float* norms,
device const float* centroids,
device float* distances,
...) {
// Each threadgroup = one neighbor
// Threads split dimensions
// Reduction via threadgroup shared memory
}

Metal을 이용해 TurboQuant 양자화된 벡터 간의 거리를 계산하려 시도한 커널 코드 예시

메모리 레이아웃과 캐시 구조가 양자화 성능에 미치는 영향을 확인했다. 연속적인 평면 저장소 구조와 하드웨어 프리페치를 활용하는 것이 복잡한 소프트웨어 최적화보다 더 큰 속도 향상을 제공했다. 거리 계산 시 룩업 테이블(LUT)을 사용하는 방식은 테이블 크기가 L1 캐시 용량을 초과할 경우 오히려 성능이 저하되는 현상이 나타났다. 결과적으로 하드웨어의 메모리 서브시스템 특성을 고려한 순차적 접근 패턴 설계가 최적화의 핵심임이 확인됐다.

용어 해설

근사 근접 이웃(ANN): — 데이터셋 전체를 검색하는 대신 근사치를 빠르게 찾는 기술로, 대규모 벡터 검색의 핵심이다.
단일 명령 다중 데이터 처리(SIMD): — 하나의 명령어로 여러 데이터를 동시에 처리하는 하드웨어 가속 기술로, 벡터 연산 속도를 비약적으로 높인다.
월시-하다마드 변환(Walsh-Hadamard Transform): — 벡터를 회전시켜 데이터 분포를 가우시안 형태로 정규화하는 수학적 변환으로, 양자화 효율을 극대화한다.
로이드-맥스 코드북(Lloyd-Max Codebook): — 특정 데이터 분포에 최적화된 양자화 간격을 정의한 표로, TurboQuant에서는 학습 없이 정해진 표를 사용한다.
메탈(Metal): — Apple 기기에서 GPU 자원을 직접 제어하는 로우레벨 API로, 고성능 병렬 연산 구현에 사용된다.

언급된 도구

NGT중립

Nearest Neighbor Search library

Metal비추천

GPU acceleration API

Rust추천

Programming language for implementation

언급된 리소스

논문TurboQuant: Data-Oblivious Vector Quantization