기술 요약: LOCKS – 효율적인 롱 컨텍스트 디코딩을 위한 페이지 로컬 압축 키 요약(Page-Local Compact Key Summaries)

1. 문제 정의

롱 컨텍스트 길이에 대해 대규모 언어 모델(LLM)을 서빙하는 것은 점점 더 KV 캐시, 구체적으로는 메모리 용량이 아닌 **KV 읽기 대역폭(KV-read bandwidth)**에 의해 병목 현상이 발생하고 있습니다. 고가용성 동시 요청이 발생하는 프로덕션 환경에서 모델 가중치는 배치(batch) 전체에 걸쳐 분산되지만, 모든 요청은 자신만의 고유한 KV 캐시를 읽어야 합니다. Llama-3.1-8B와 같은 모델의 경우, 128K 컨텍스트는 요청당 약 16 GiB의 KV 캐시를 필요로 하며, 이는 모델 가중치 크기를 초과합니다. 결과적으로 가속기가 처리할 수 있는 요청 수를 제한하는 것은 캐시입니다.

어텐션(attention)은 전역적으로는 높은 랭크(high-rank)를 갖지만, **로컬로는 낮은 랭크(low-rank)**를 가집니다. 즉, 특정 디코드 단계에서 몇 개의 페이지(토큰 블록)가 쿼리의 어텐션 질량(attention mass) 대부분을 차지합니다. 희소 디코딩(Sparse decoding)은 이러한 핵심 페이지에만 어텐션을 수행하는 것을 목표로 합니다. 그러나 기존의 선택 방식들은 다음과 같은 트레이드오프에 직면해 있습니다:

• 정확한 선택(Exact selection) (모든 키를 읽는 방식)은 전체 대역폭 비용을 발생시키므로 희소성의 목적을 무색하게 만듭니다.
• 상주 요약(Resident summaries) (페이지의 압축된 표현)은 키를 읽지 않기 위해 필요하지만, 기존 방식들은 공격적인 예산 제약 하에서 "캐리어(carrier)" 페이지(정답을 포함하고 있는 페이지)를 식별하는 데 필요한 특정 기하학적 구조를 보존하는 데 종종 실패합니다. 페이지 간 공유된 투영(Shared projections)은 페이지별 콘텐츠에 대해 이론적으로 무지(blind)하며, 이는 검색 및 추론 시나리오에서 태스크 붕괴(task collapse)로 이어집니다.

2. 방법론: LOCKS

LOCKS(Page-Local Compact Key Summaries)는 전체 KV 캐시를 상주시키되 디코딩 중 대부분의 키를 읽는 것을 피함으로써 이 문제를 해결합니다. 이는 훈련이 필요 없고 vLLM에 플러그인 형태로 바로 적용 가능한 **페이지 로컬 스펙트럼 요약(page-local spectral summary)**을 기반으로 작동합니다.

핵심 메커니즘

1. 페이지 로컬 스펙트럼 분해 (Page-Local Spectral Decomposition):

   • 공유된 전역 또는 시퀀스 수준의 기저(basis)를 사용하는 대신, LOCKS는 KV 캐시의 각 페이지에 대해 고유한 저차원 기저를 구축합니다.
   • 페이지가 확정될 때(prefill 단계 또는 페이지 교차 리프레시 시), 해당 페이지의 키 편차(centered keys)에 대해 작은 고유값 분해(eigendecomposition)를 수행합니다.
   • 요약본은 상위 $r$개의 고유벡터(기저)와 각 페이지 내 토큰에 대한 대응 계수(coefficients)를 저장합니다.
   • 양자화(Quantization): 기저는 int4로, 계수 및 중심점(centroids)은 int8로 저장됩니다. 랭크 $r=8$, 페이지 크기 $B=16$일 때, 요약본은 원래 페이지 KV 바이트의 약 10% 수준입니다.

2. 재구성 및 선택 (Reconstruction and Selection):

   • 매 디코드 단계마다 LOCKS는 압축된 요약을 사용하여 **페이지 내 로짓(within-page logits)**을 재구성합니다:
     $$\hat{s}_j(q) = \log \sum_{i \in P_j} \exp\left( \frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}} \right)$$
     여기서 $\mu_j$는 페이지 중심점, $V_j$는 페이지 로컬 기저, $c_i$는 계수입니다.
   • 이 재구성된 로짓의 **Log-Sum-Exp (LSE)**를 계산하여 페이지의 어텐션 질량을 추정합니다.
   • 선택은 이 추정된 질량을 기준으로 페이지 순위를 매겨 수행됩니다. 결정적으로, 이 선택 단계 동안 어떠한 후보 키(candidate keys)나 값(values)도 읽지 않습니다.

3. 그룹 공유 선택 (Group-Shared Selection, GQA):

   • Grouped-Query Attention (GQA)의 경우, LOCKS는 그룹 내 쿼리 헤드들 사이에서 추정된 질량 공유를 평균하여 단일 공유 선택 테이블을 형성합니다. 이를 통해 모든 헤드가 동일한 페이지 세트에 어텐션을 수행하도록 하여 트래픽 파편화를 방지합니다.

4. 구현:

   • 수정되지 않은 vLLM을 위한 플러그인으로 구현되었습니다.
   • 전체 디코드 단계(스코어링, 선택 및 희소 어텐션)는 전체 CUDA 그래프(full CUDA graphs) 내에서 실행되어, Python 디스패치 오버헤드를 제로화하고 효율적인 커널 퓨전(kernel fusion)을 가능하게 합니다.
   • 요약 생성은 디코드 크리티컬 패스(critical path)와 분리되어, 페이지가 확정될 때만 발생합니다.

3. 주요 기여

이론적 발견

• 구조의 국소성 (Locality of Structure): 본 논문은 페이지 선택에 중요한 기하학적 구조가 시퀀스나 캐시 범위가 아닌 페이지 범위에 집중되어 있음을 입증합니다. 페이지 자체의 랭크-8 고유 기저는 대부분의 키 에너지를 포착하는 반면, 공유된 기저는 (크기가 같더라도) 페이지별 방향을 포착하는 데 실패합니다.
• 공유 투영의 불가능성 (Impossibility of Shared Projections): 저자들은 고정된 쿼리 독립적 공유 저차원 선형 투영이 페이지 내부의 $(d-\rho)$ 차원 패밀리 차이에 대해 콘텐츠 무지(content-blind) 상태임을 증명했습니다(Proposition 1). 즉, 공유 기저는 특정 콘텐츠 방향에서만 다른 페이지를 구별할 수 없으며, 이는 선택 오류를 야기합니다.
• 보존 보장 (Retention Guarantees): 본 논문은 어텐션 출력의 오차가 페이지 로컬 재구성의 스펙트럼 테일 에너지(spectral tail energy)에 의해 직접 제어됨을 보여주는 결합된 이론적 경계(Theorem 1)를 제공합니다. 페이지 로컬 기저는 직교 랭크-$r$ 재구성 중 잔차 최적(residual-optimal)임이 증명되었습니다.

실증적 결과

• 검색 및 QA (LongBench-v1, RULER): LOCKS는 가장 작은 예산에서도 "모든 키를 읽는" 오라클(exact LSE selection)의 성능을 추적합니다. 검색 밀도가 높은 작업(RULER)에서, LOCKS는 기존의 최첨단 셀렉터들(Quest, KVzip, ShadowKV, RocketKV)보다 캐리어 유지(carrier retention)(정답을 포함하는 페이지를 유지하는 능력) 측면에서 크게 앞섭니다. 이는 태스크 성공의 핵심 동력입니다.
• 장문 추론 (AIME26, MATH-500): 모델이 수천 개의 토큰을 생성하는 추론 작업에서, 베이스라인 셀렉터들은 타이트한 예산에서 성능이 붕sk히는 경우가 많습니다. LOCKS는 약 2%의 토큰에만 어텐션을 수행하면서도 FullKV 성능에 근접하며 높은 정확도를 유지합니다.
• 효율성:
  • 1M 컨텍스트 길이에서, LOCKS는 Dense Attention(FlashAttention-3/FlashInfer) 대비 토큰당 레이턴시를 2.0배 향상시킵니다.
  • 1M 컨텍스트에서 디코드 단계당 읽는 KV 바이트를 9.8배 감소시킵니다.
  • 처리량(Throughput) 향상은 배치 크기에 따라 스케일링되며, 큰 배치 환경의 256K 컨텍스트에서 1.80배에 달합니다.

4. 의의 및 주장

본 논문은 **페이지 로컬 스펙트럼 집중(page-local spectral concentration)**이 효율적인 희소 디코딩을 위한 누락된 연결 고리라고 주장합니다. 표현의 범위를 글로벌/시퀀스에서 로컬(페이지별)로 전환함으로써, LOCKS는 모델을 훈련하거나 아키텍처를 수정하지 않고도 태스크 완수에 필요한 "캐리어" 페이지를 보존합니다.

• 드롭인 통합 (Drop-in Integration): 모델 포크(fork)나 훈련 게이트를 요구하는 많은 선행 연구들과 달리, LOCKS는 수정되지 않은 vLLM 및 표준 GQA 모델과 함께 작동합니다.
• 이론-실제 정렬 (Theoretical-Practical Alignment): 이 방법은 재구성 오차에 대한 이론적 경계를 엔드투엔드 태스크 품질 및 캐리어 생존과 직접 연결한 드문 사례 중 하나로, "스펙트럼 집중에서 태스크 성능으로 이어지는 체인"이 유효함을 입증합니다.
• 확장성 (Scalability): 이 접근 방식은 메모리 대역폭이 주요 병목인 극단적인 컨텍스트(1M+ 토큰)로 확장 가능하며, CPU 오프로딩이나 계층형 메모리 없이도 롱 컨텍스트 모델을 서빙할 수 있는 실질적인 경로를 제시합니다.

저자들은 한계점에 대해 겸허한 입장을 유지합니다: 배치된 양자화된 요약의 보존력은 엄격하게 단계별로 경계가 정해지기보다 경험적으로 측정되었으며, 현재 고정된 페이지 크기를 가정하고 있습니다. 그러나 결과는 훈련이 필요 없는 페이지 로컬 스펙트럼 요약이 공격적인 희소 예산에서도 오라클 선택을 추적하기에 충분하며, 실질적인 롱 컨텍스트 가속을 가능하게 한다는 점을 입증합니다.