OmniStack-RS: 추천 시스템을 위한 KV 캐시 압축 및 개인화 추론 실험

실용적 조언

메모리 대역폭이 제한된 환경에서는 역양자화 로직을 어텐션 커널에 통합(Fused)하여 I/O 오버헤드를 줄여야 한다.
사용자별 어댑터가 많은 경우 Multi-LoRA 디스패치 최적화를 통해 제어 흐름 오버헤드를 최소화해야 한다.

섹션별 상세

KV 캐시 압축을 위해 INT4 Lloyd-Max 양자화와 1비트 Rademacher QJL 잔차 기법을 결합했다. 입력 데이터를 Lloyd-Max 알고리즘으로 4비트 코드북에 매핑하고, 발생하는 오차를 1비트 잔차로 보정하여 정밀도를 유지한다. 이를 통해 BF16 대비 원소당 4.75비트 수준으로 데이터를 압축하면서도 FP32 대비 최대 오차를 0.002403으로 억제했다. 압축 효율은 기존 대비 3.37배 향상되어 더 많은 동시 사용자를 수용할 수 있는 구조를 갖췄다.

역양자화와 Softmax 연산을 통합한 Fused Triton Attention 커널을 구현하여 연산 효율을 극대화했다. 커널 내부에서 압축된 데이터를 실시간으로 복원하며 어텐션 연산을 수행하여 메모리 대역폭 낭비를 줄인다. NVIDIA A10 GPU에서 측정된 P99 커널 지연 시간은 0.69ms이며, 전체 엔드투엔드 지연 시간은 1.13ms로 나타났다. 초당 1,633.93개의 쿼리를 처리할 수 있는 성능을 확보하여 실시간 추천 환경에 적합함을 증명했다.

사용자별 개인화를 위해 O(1) 복잡도의 Multi-LoRA 디스패치 시스템을 적용했다. 각 사용자마다 별도의 어댑터 상태를 가지더라도 연산 오버헤드가 사용자 수에 비례해 늘어나지 않도록 설계했다. Criteo Day 23 광고 상호작용 데이터셋을 활용한 벤치마크에서 초당 104,571명의 사용자 컨텍스트를 처리하는 처리량을 기록했다. 이는 대규모 사용자 기반의 서비스에서 개인화된 LLM 추론을 확장 가능하게 만드는 핵심 요소이다.

용어 해설

키-값 캐시(KV Cache): — LLM 추론 시 이전 토큰들의 Key와 Value 행렬을 메모리에 저장해두는 기술이다. 매 단계마다 전체 시퀀스를 재계산하지 않아도 되어 속도가 빠르지만, 사용자 수가 늘어날수록 GPU 메모리 점유율이 급격히 높아지는 병목 현상의 원인이 된다.
로이드-맥스 양자화(Lloyd-Max Quantization): — 데이터의 분포를 고려하여 양자화 오차(MSE)를 최소화하도록 코드북을 설계하는 비균일 양자화 기법이다. 이 프로젝트에서는 KV 캐시 값을 INT4로 압축하면서도 정보 손실을 줄이기 위해 사용됐다.
트라이톤 커널(Triton Kernel): — OpenAI에서 개발한 GPU 프로그래밍 언어로, CUDA보다 생산성이 높으면서도 고성능 커널 작성이 가능하다. 본문에서는 역양자화와 어텐션 연산을 하나로 합친 Fused Kernel을 구현하는 데 활용됐다.
QJL 잔차(QJL Residual): — Quantized Johnson-Lindenstrauss의 약자로, 양자화 과정에서 발생하는 오차를 보정하기 위해 추가적인 비트를 사용하는 기법이다. 여기서는 1비트 Rademacher 투영을 통해 양자화 정밀도를 회복하는 역할을 한다.

언급된 도구

Triton추천

Fused Attention 커널 및 역양자화 로직 구현

Nsight Compute추천

GPU 커널 성능 프로파일링 및 분석

언급된 리소스

GitHubOmniStack-RS GitHub Repository