기술 요약: Spherical KV

1. 문제 정의

대규모 언어 모델(LLM)의 롱 컨텍스트 추론은 산술 처리량(arithmetic throughput)보다는 키-값(Key-Value, KV) 캐시에 의해 점점 더 제약을 받고 있습니다. 컨텍스트 길이($T$)가 증가함에 따라, 상주 메모리 점유율(resident memory footprint)은 $T$, 배치 크기($B$), 레이어 수($L$), 헤드 수($H$)에 따라 선형적으로 증가합니다. PagedAttention과 같이 페이지드 메모리를 사용하는 프로덕션 서빙 환경에서, 지배적인 병목 현상은 부동 소수점 연산(FLOPs)에서 고대역폭 메모리(HBM) 트래픽으로 전환됩니다. 구체적으로, 디코딩 과정은 장치 메모리에서 연산 유닛으로 역사적 KV 페이지를 반복적으로 스트리밍하는 작업에 의해 제한됩니다.

기존 솔루션들은 다음 세 가지 레버를 통해 이를 완화하려고 시도합니다:

1. 유지/제거 (Retention/Eviction): 더 적은 수의 토큰을 저장함 (예: StreamingLLM, H2O).
2. 양자화/압축 (Quantization/Compression): 더 적은 비트를 저장함 (예: KVQuant, KIVI).
3. 오프로딩 (Offloading): KV를 다른 곳에 저장함 (예: CPU/Storage).

그러나 이러한 방법들은 결정적인 경로(critical-path)의 병목을 완전히 해결하지 못하는 경우가 많습니다. 많은 압축 기술은 어텐션 로짓(attention logits)을 계산하기 전에 압축된 표현으로부터 밀집 벡터(dense vectors)를 재구성해야 하는 "역구조 후 내적(decompress-then-dot-product)" 경로를 요구합니다. 이러한 재구성은 HBM 트래픽을 다시 유발하며 커널 퓨전(kernel fusion)을 방해하여, 실제 페이지드 서빙 기질(paged serving substrates)에서의 이론적인 메모리 절감 효과를 무효화합니다. 또한, 단순한 제거 정책은 저효용 상태를 버리면서도 중요한 정보(예: 검색된 증거)를 보존하는 정밀도(granularity)가 부족한 경우가 많습니다.

2. 방법론: Spherical KV

본 논문은 KV 할당을 어텐션 기하학(attention geometry)에 기반한 레이트-왜곡(rate–distortion) 문제로 취급하는 롱 컨텍스트 추론 방법인 Spherical KV를 제안합니다. 이는 페이지드 및 러기드(ragged) 서빙 기질 내에서 작동하도록 설계된 두 가지 직교하는 구성 요소인 **Angle-Domain Attention (ADA)**와 **Rate–Distortion Retention (RDR)**으로 구성됩니다.

2.1 Angle-Domain Attention (ADA)

ADA는 표준 "역구조 후 내적" 프리미티브를 컴팩트한 구형 코드(spherical codes)에서 직접 작동하는 컴퓨트 네이티브(compute-native) 방식으로 대체합니다.

• 구형 파라미터화 (Spherical Parameterization): 키(Keys)는 스칼라 반경($r = \|k\|$)과 방향 각도($\phi = \Phi(\hat{k})$)로 분해됩니다.
• 압축 도메인 유사도 (Compressed-Domain Similarity): ADA는 밀집 벡터를 캐싱하는 대신, 양자화된 반경($\tilde{r}$)과 패킹된 각도 코드($c_\theta$)를 저장합니다. 디코드 핫 루프(decode hot loop) 동안, 어텐션 커널은 이 코드들로부터 직접 로짓을 계산합니다:
  $$\hat{\ell}(q, k) \approx \|q\| \cdot \tilde{r} \cdot f_b(c_\theta^q, c_\theta^k)$$
  여기서 $f_b$는 각도 코드로부터 코사인 유사도를 계산하는 퓨즈드 커널(fused kernel) 연산입니다.
• 재구성 불변성 (No Reconstruction Invariant): 핵심 설계 제약 조건은 디코드 커널이 HBM 내에 밀집 키 벡터($\tilde{k} \in \mathbb{R}^d$)를 절대 실체화(materialize)하지 않는다는 것입니다. 이를 통해 저장된 KV 바이트의 감소가 생성된 토큰당 HBM 트래픽의 직접적인 감소로 이어지도록 보장합니다.

2.2 Rate–Distortion Retention (RDR)

RDR은 엄격한 메모리 예산 하에서 어떤 상태를 유지할지, 그리고 어떤 정밀도(티어)로 저장할지를 공동으로 최적화하는 제어 메커니즘입니다.

• 공동 최적화 (Joint Optimization): 제거와 양자화를 별개로 취급하는 방식과 달리, RDR은 이산적인 레이트-왜곡 할당 문제를 해결합니다. 각 토큰/헤드/레이어 상태 $i$에 대해, RDR은 이진 잔류 결정($z_i \in \{0, 1\}$)과 정밀도 티어($b_i \in \mathcal{B}$)를 선택합니다.
• 왜곡 대리물 (Distortion Surrogate): 컨트롤러는 ADA의 에러 분석에서 유도된 기계론적 왜곡 대리물을 사용합니다. 이는 로짓 에러가 쿼리와 키의 크기에 의해 스케일링된 각도 근사 에러($\epsilon_\theta$)에 의해 지배된다고 추정합니다.
• 보호 제약 조건 (Protection Constraints): 중요한 정보(예: 검색된 증거, 안전 구간, 아웃라이어)에 대한 치명적인 실패를 방지하기 위해, RDR은 강력한 보호 제약 조건($I_{prot}(i)=1 \implies z_i=1, b_i=b_{max}$)을 적용합니다.
• 티어 동질 페이지 (Tier-Homogeneous Pages): 컨트롤러는 유지된 상태들을 모든 항목이 동일한 정밀도 티어를 공유하는 페이지로 그룹화합니다. 이를 통해 디코드 커널이 불규칙한 게더(gather) 연산을 피하고, 헤더 가이드 스트리밍 읽기를 단일 패스로 수행할 수 있게 합니다.

2.3 시스템 통합

Spherical KV는 표준 페이지드/러기드 KV 기질(vLLM 스타일 레이아웃과 호환 가능) 위에서 작동합니다. 파이프라인은 밀집 키를 구형 코드로 매핑하고, RDR을 적용하여 필터링 및 티어링을 수행한 뒤, 압축된 형식으로 저장합니다. 디코드 커널은 이러한 코드들을 직접 소비하며, "밀집 재구성 없음(no-dense-reconstruction)" 불변성을 유지합니다.

3. 주요 기여

1. Angle-Domain Attention (ADA): 밀집 키 재구성과 관련된 HBM 트래픽 페널티를 제거하고, 컴팩트한 구형 코드로부터 직접 어텐션 로짓을 계산하는 커널 레벨 프리미티브를 제공합니다.
2. Rate–Distortion Retention (RDR): 고효용 상태를 우선시하고 취약한 상태를 보호하면서, 고정된 예산 하에서 잔류 여부와 정밀도 티어를 공동으로 할당하는 통합 컨트롤러를 제공합니다.
3. 배포 지향적 평가 (Deployment-Oriented Evaluation): 본 방법론은 단순한 연속적 캐시가 아닌 실제적인 페이지드/러기드 서빙 기질에서 평가되어, 엔드 투 엔드 처리량, HBM 트래픽 및 행동 안정성을 측정합니다.
4. 파레토 프런티어 이동 (Pareto Frontier Shift): Spherical KV가 밀집 KV 및 강력한 베이스라인들과 비교하여, 더 낮은 메모리 점유율에서 더 높은 처리량을 달자함으로써 메모리-품질-처리량의 프런티어를 이동시킴을 입증합니다.

4. 실험 결과

저자들은 Llama-3.1-8B-Instruct, Qwen2.5-14B-Instruct, gpt-oss-20b를 대상으로 컨텍스트 길이 8K, 32K, 128K에 걸쳐 Spherical KV를 평가했습니다. 워크로드에는 언어 모델링(PG-19), 검색 중심 QA(LongBench: HotpotQA, 2WikiMultiHopQA), 그리고 에이전틱 툴 실행(AgentBench, ToolBench)이 포함되었습니다.

• 처리량 및 메모리 효율성: 일치된 품질 제약(iso-quality) 하에서, Spherical KV는 토큰당 상주 KV 바이트를 24–42% 줄이는 동시에 디코드 처리량을 1.55×–1.72× 개선했습니다.
• 컨텍스트 스케일링: 이러한 이득은 KV 잔류 및 HBM 트래픽이 시스템 성능을 지배하는 128K 영역에서 가장 두드러졌습니다.
• 커널 구현 효율성: 본 방법론은 생성된 토큰당 HBM 바이트를 성공적으로 감소시켰으며, 이는 메모리 절감이 재구성 오버헤드에 의해 상쇄되지 않았음을 확인시켜 줍니다.
• 안정성: 시스템은 더 타이트한 예산 하에서도 매끄러운 성능 저하를 보여주었으며, 추론 궤적과 툴 호출 시퀀스를 보존하여 하드 제거(hard-eviction) 베이스라인에서 나타나는 급격한 실패를 방지했습니다.

5. 의의 및 주장

본 논문은 Spherical KV가 롱 컨텍스트 추론의 근본적인 병목 현상인 메모리 용량과 메모리 대역폭 사이의 충돌을 해결한다고 주장합니다. KV 압축을 단순히 저장의 문제가 아니라 컴퓨트 네이티브(compute-native) 문제(즉, 압축된 표현이 커널에 의해 직접 실행 가능한 문제)로 다룸으로써, 이 방법은 운영 프런티어의 진정한 변화를 달성합니다.

저자들은 다음을 강조합니다:

• 압축은 컴퓨트 네이티브여야 한다: 디코드 경로가 밀집 재구성을 요구한다면 단순히 더 적은 비트를 저장하는 것만으로는 불충분합니다.
• 메모리는 중요한 곳에 예산을 배정해야 한다: 레이트-왜곡 제어를 통해 시스템은 저효용 상태에서는 절약하고 고영향 상태에는 정밀도를 집중할 수 있습니다.
• 평가는 배포 환경에 충실해야 한다: 실제적인 페이지드 레이아웃, 메타데이터 오버헤드 및 커널 제약 조건 속에서도 이득이 유지되어야 진정한 성과로 간주될 수 있습니다.

Spherical KV는 현대적인 LLM 서빙의 메모리 제한적(memory-bound) 환경을 타겟으로 하여, 디코드 효율성과 행동 안정성을 유지하면서 KV 잔류 및 HBM 트래픽을 줄이는 메커니즘으로 제시됩니다.