Sparse KV 캐시를 통한 어텐션 스케일링 비용 절감

Sparse key-value (KV) 캐시는 softmax 어텐션 (attention)의 이차 함수적 급증(quadratic blow-up)을 시퀀스 길이(sequence length)에 따라 거의 선형적으로 증가하는 비용으로 압축합니다. 각 쿼리(query)가 블록 단위의 KV 메모리 중 아주 작은 top-k 서브셋(subset)에만 어텐션하도록 함으로써, 쿼리당 작업량이 전체 컨텍스트(context) 규모에 따라 늘어나는 것을 방지합니다. 이러한 작은 변화는 초장기 시퀀스(ultra-long sequences)에 대한 확장성 곡선을 뒤집으며, 단일 GPU에서 수백 킬로바이트(multi-hundred-kilobyte) 규모의 윈도우(window)를 실용적으로 사용할 수 있게 만듭니다.

이 연구 이전의 지배적인 방식은 dense attention (밀집 어텐션)이었으며, 이는 $O(N^{2})$의 메모리 및 FLOP 예산으로 인해 컨텍스트 윈도우를 수천(few k) 토큰 수준으로 제한했습니다. Grouped Query Attention (GQA)은 캐시 재사용성을 개선했지만, 여전히 각 그룹이 모든 KV 블록을 스캔해야 했기에 이차 항(quadratic term)이 그대로 남아 있었습니다. 이러한 접근 방식들은 윈도우가 커짐에 따라 연산량을 일정하게 유지할 수 없었으며, 길이와 지연 시간(latency) 사이의 트레이드오프(trade-off)를 강요했습니다.

MSA는 100만 토큰 컨텍스트에서 토큰당 어텐션 연산량을 28.4배 절감합니다. 저자들은 "네이티브 멀티모달(native multimodal) 학습을 거친 109B 파라미터 모델에서, MSA는 1M 컨텍스트에서 토큰당 어텐션 연산량을 28.4배 줄이면서도 GQA와 대등한 성능을 보여준다"라고 보고했습니다 [1]. 이러한 감소 폭은 길이에 따라 더 심화되는데, "그림 4에서 보여주듯, MSA는 우리의 설정에서 GQA 대비 토큰당 어텐션 FLOPs를 상당히 줄이며, 컨텍스트가 길어질수록 그 감소 폭이 커진다"라고 명시되어 있습니다 [1].

KV 메모리 사용량은 **최대 50%**까지 감소하며, 퍼플렉시티(perplexity)는 dense baseline과 구별할 수 없을 정도로 유지됩니다. README에 따르면 sparse 브랜치는 "퍼플렉시티를 유지하면서 KV 메모리를 최대 50%까지 줄인다"라고 언급되어 있으며, 이는 절감된 비용이 정확도 손실로 이어지지 않음을 확인시켜 줍니다 [1]. KV 점유 면적(footprint)을 절반으로 줄이는 것은 동일한 하드웨어에서 더 큰 윈도우를 사용할 수 있음을 직접적으로 의미합니다.

실제 실행 시간(Wall-clock speedups)의 향상 또한 매우 극적입니다. H800 GPU에서 프리필(prefill)은 14.2배, 디코딩(decoding)은 7.6배 더 빠르게 실행됩니다. 논문의 벤치마크 스위트(benchmark suite)에 따르면, 공동 설계된 커널(co-designed kernel)을 사용할 때 “14.2배의 프리필 및 7.6배의 디코딩 실제 실행 시간 향상”을 보여줍니다 [1]. 이러한 이득은 지수 연산이 없는 Top-k 선택기(Top-k selector)와 블록 단위의 텐서 코어 활용(block-granular tensor-core utilization)에서 비롯됩니다.

이러한 증거는 단일 109B 파라미터 멀티모달 모델(multimodal model)과 MiniMax의 코드베이스에 긴밀하게 결합된 커널로부터 도출된 것이므로, 다른 모델 제품군이나 하드웨어 세대로의 이식성(portability)은 아직 입증되지 않았습니다. 또한, 블록 단위 희소성(blockwise sparsity)은 대부분의 관련 토큰이 고정된 인과적 지평(causal horizon) 내에 존재한다고 가정합니다. 따라서 전역적 어텐션(global attention)이 필요한 작업은 여전히 어려움을 겪을 수 있습니다. 논문 자체에서도 “선택된 블록은 최대 인과적으로 보이는 토큰들을 포함하며, 쿼리당 어텐션 비용은 감소하지만... 이는 시퀀스 길이가 증가함에 따라 고정된다” [1]라고 경고하며, 비인과적 패턴(non-causal patterns)에 대한 잠재적인 사각지대를 암시하고 있습니다.

만약 이러한 희소성 기법(sparsity tricks)이 모델 크기에 관계없이 유지된다면, 개발자들은 메모리나 연산량을 추가하지 않고도 일반 GPU에서 컨텍스트 윈도우(context windows)를 두 배 또는 세 배로 늘릴 수 있습니다. 이전에 8k 토큰으로 제한되었던 벤치마크들은 MSA가 활성화된 상태에서 다시 실행되어야 하며, 프로덕션 파이프라인(production pipelines)은 최대 길이 기본값을 안전하게 높임으로써 저장소 규모의 코드 분석(repository-scale code analysis) 및 지속적인 대화 메모리(persistent conversational memory)와 같은 새로운 유스케이스(use-cases)를 확보할 수 있습니다.

References

1. MiniMax Sparse Attention