TriAttention: 정확도를 유지하면서 기하학적 트릭으로 LLM 메모리 사용량을 10배 절감하는 방법

긴 문맥 추론 (Long-context reasoning)은 현대 LLM 추론에서 가장 메모리를 많이 소모하는 워크로드 중 하나입니다. 모델이 32,000개의 토큰으로 구성된 사고의 사슬 (chain-of-thought)을 생성할 때, 이전 모든 어텐션 (attention) 단계에서 저장된 키와 값인 KV 캐시 (KV cache)는 수십 기가바이트의 GPU 메모리를 소비할 수 있습니다. MIT, NVIDIA, Zhejiang University의 연구진은 최근 TriAttention을 발표했습니다. 이는 전체 어텐션 (full attention)의 정확도를 유지하면서도 메모리 점유율을 최대 10.7배까지 줄이는 압축 방법입니다. 핵심 통찰은 휴리스틱 (heuristic)이 아닌 기하학적 (geometric)인 방식에 있습니다. 즉, 실행 시점에 어떤 토큰이 중요한지 추측하는 대신, TriAttention은 모델 자체 가중치 공간 (weight space)의 안정적인 속성으로부터 중요도를 예측합니다.

기존 KV 캐시 압축이 실제 운영 환경에서 실패하는 이유

KV 캐시 압축의 표준적인 접근 방식은 토큰 제거 (token eviction)입니다. 각 캐시된 토큰이 받은 어텐션 (attention) 점수를 매긴 다음, 점수가 낮은 토큰을 삭제하는 방식입니다. SnapKV 및 H2O와 같은 방법들이 이 패턴을 따릅니다. 이러한 방식은 연구 환경에서는 잘 작동하지만, 실제 운영 시스템에 배포될 때는 두 가지 구체적인 문제에 직면합니다.

FlashAttention 가시성 문제. 실제 운영 추론은 FlashAttention에 의존합니다. FlashAttention은 계산을 SRAM 내부에서 타일링 (tiling)하며, 전체 N×N 어텐션 점수 행렬을 GPU 메모리에 쓰지 않습니다. 대부분의 제거 (eviction) 방법은 어떤 토큰을 유지할지 결정하기 위해 해당 점수들이 필요합니다. 이 점수들이 없으면 시스템은 더 느린 "eager" 어텐션으로 돌아가야 하며, 이는 애초에 압축을 통해 얻으려 했던 성능 이점을 상쇄해 버립니다.

페이지 메모리 파편화 (paged memory fragmentation) 문제. vLLM과 같은 서빙 프레임워크는 페이지 할당기 (paged allocator)를 사용하여 고정된 크기의 물리적 블록 단위로 GPU 메모리를 관리합니다. 블록은 완전히 비워졌을 때만 해제됩니다. 표준적인 제거 (eviction) 전략은 소수의 "생존" 토큰들을 여러 블록에 흩어놓게 되어, 할당기가 블록을 회수할 수 없게 만듭니다. 메모리 절감 효과가 이론상으로는 존재하지만 실제로는 나타나지 않는 것입니다.

TriAttention은 NVIDIA의 KV 캐시 압축 인프라에 관한 연구 블로그에 상세히 설명된 바와 같이, 이러한 두 가지 실패 모드를 모두 피하기 위해 특별히 설계되었습니다.

핵심 아이디어: Pre-RoPE 공간에서의 Q/K 집중 (Q/K Concentration)

이 방법은 트랜스포머 (Transformer) 어텐션 헤드 (Attention heads)가 내부적으로 어떻게 동작하는지에 대한 경험적 관찰에서 시작됩니다. RoPE (Rotary Position Embedding)가 시퀀스 위치에 따라 쿼리 (Query) 및 키 (Key) 벡터를 회전시키기 전인 Pre-RoPE 표현 공간 (Pre-RoPE representation space)에서, 약 90%의 어텐션 헤드에 대한 Q 및 K 벡터는 안정적이고 0이 아닌 중심점 주변에 조밀하게 클러스터링됩니다. 저자들은 이를 **Q/K 집중 (Q/K concentration)**이라고 부릅니다.

이것이 중요한 이유는 RoPE 회전이 Post-RoPE 쿼리를 불안정하게 만드는 원인이기 때문입니다. 위치 1,000에 있는 쿼리 벡터는 RoPE에 의해 회전되었기 때문에, 위치 5,000에 있는 동일한 의미의 쿼리와는 다른 방향을 가리킵니다. 이러한 불안정성은 토큰 중요도를 추정할 때 얼마나 멀리 과거를 볼 수 있는지를 제한합니다. Pre-RoPE 벡터는 이러한 문제가 없습니다. 이들의 중심은 위치나 입력 컨텍스트에 관계없이 고정된 상태를 유지합니다.

Q 및 K 벡터가 고정된 중심 주변에 집중될 때, 임의의 쿼리와 임의의 키 사이의 어텐션 로짓 (Attention logit)은 해당 중심점들과 두 토큰 사이의 위치적 거리(Positional distance)에만 의존하는 계수를 가진 **삼각 급수 (Trigonometric series)**로 근사할 수 있습니다. TriAttention은 실제 어텐션 점수 (Attention scores)를 관찰할 필요 없이 이 급수를 사용하여 키의 중요도를 점수화합니다. 집중도가 낮은 소수의 헤드에 대해서는, 측정된 집중도 지표에 의해 자동으로 가중치가 부여되는 노름 기반 (Norm-based) 신호를 혼합합니다.

그 결과, 정확하면서도 인프라 친화적인 중요도 추정기가 탄생했습니다. 이는 어텐션 점수 행렬 (Attention score matrix)을 필요로 하지 않으므로 FlashAttention과 네이티브하게 작동합니다.

메모리 파편화 문제 해결

어떤 토큰을 유지할지 예측하는 것은 문제의 절반에 불과합니다. 나머지 절반은 그 토큰들이 점유했던 메모리를 실제로 해제하는 것입니다. TriAttention은 **Forward-Packing Compaction (전방 패킹 압축)**이라 불리는 메커니즘을 도입합니다. 약 128번의 디코딩 (decode) 단계마다, 시스템은 생존한 토큰들을 물리적으로 이동시켜 가능한 한 적은 수의 메모리 블록으로 통합하며, 이를 통해 꼬리 부분의 블록(tail blocks)을 비워 페이지 할당기 (paged allocator)가 이를 회수할 수 있도록 합니다.

이는 단순한 엔지니어링 단계이지만, 메모리 절감이 이론적인 수치에 그치지 않고 실제로 구현되게 만드는 핵심 요소입니다. 압축 (compaction) 과정이 없다면, 제거 (eviction) 방식은 할당기가 건드릴 수 없는 파편화된 생존 토큰들을 남기게 됩니다.

수치적 결과

저자들은 32,000개 토큰 생성 시 KV 캐시 예산을 3,072개 토큰으로 제한하여(전체 캐시 크기에서 약 10배 감소) AIME 2025 데이터셋에서 TriAttention을 평가했습니다.

MATH 500 데이터셋에서 이 방법론은 32,768개의 토큰 중 단 1,024개만을 유지하면서도 68.4%의 점수를 기록했습니다. 이는 Full Attention의 69.6%와 비교했을 때, 32배 압축률에서 단 1.2%포인트의 차이입니다.

실질적인 결과로서, 긴 문맥 생성 (long-context generation) 중에 단일 RTX 4090에서 메모리 부족 (out of memory) 현상이 발생할 수 있는 Qwen3-32B와 같은 모델들도 TriAttention을 적용하면 성공적으로 실행할 수 있습니다. GitHub 저장소에는 vLLM 및 SGLang 통합 기능이 포함되어 있으며, HIP/ROCm을 지원하는 llama.cpp용 커뮤니티 포팅 버전도 등장했습니다.

이것이 긴 문맥 추론 (Long-Context Inference)에 의미하는 바

TriAttention의 더 넓은 의미는 실험실 수준의 압축 결과와 실제 프로덕션 배포(production deployment) 사이의 간극을 메운다는 점에 있습니다. 기존 방식들은 벤치마크에서는 좋은 성능을 보이는 경우가 많았으나, FlashAttention이나 페이지 메모리 관리(paged memory management)와 충돌하여 실제 서빙 스택(serving stacks)에서는 실패하곤 했습니다. TriAttention은 처음부터 이러한 제약 사항들을 염두에 두고 설계되었습니다.

삼각 함수 근사(trigonometric approximation) 접근 방식은 향후 연구를 위한 방향성도 제시합니다. 즉, 토큰의 중요도(token importance)를 런타임(runtime)에 측정해야 하는 대상으로 취급하는 대신, 모델이 학습한 기하학적 구조(learned geometry)로부터 예측할 수 있다는 것입니다. 관찰(observation)에서 예측(prediction)으로의 이러한 전환은 이 방법론이 FlashAttention의 가시성 문제(visibility problem)를 완전히 우회할 수 있게 해줍니다.

대규모로 긴 문맥(long-context) 워크로드를 실행하는 팀에게, 정확도 손실 없이 10배의 메모리 절감과 2.5배의 처리량(throughput) 향상을 결합한 이 기술은 검토할 가치가 충분합니다. arXiv 논문에는 구현 세부 사항이 포함되어 있으며, Hugging Face 논문 페이지를 통해 커뮤니티 토론 및 추가 리소스로 연결할 수 있습니다.

이 방법론은 MIT, NVIDIA, Zhejiang University의 저자들에 의해 ICML 2026에서 발표되었습니다.