TriAttention: Efficient Long Reasoning with Trigonometric KV Compression

KV Cache가 뭔지부터

Transformer는 토큰을 하나씩 생성할 때마다 이전 모든 토큰의 K, V를 다시 계산해야 한다. 이걸 피하려고 이전에 계산한 K, V를 메모리에 저장해두는 게 KV Cache다. 매 step마다 새 토큰의 K, V만 추가하면 되니까 중복 계산을 없앨 수 있다.

문제는 메모리다.

메모리 = 2 × L × H × d_head × T × bytes_per_element

 

각 항의 의미는 이렇다. 2는 K와 V를 각각 저장하기 때문이고, L은 레이어 수, H는 레이어당 헤드 수, d_head는 헤드 하나의 차원, T는 시퀀스 길이, bytes_per_element는 dtype에 따른 바이트 수(FP16이면 2, BF16이면 2, FP32이면 4)다. d_model = H × d_head로 대입하면 다음처럼 간소화된다.

메모리 = 2 × L × d_model × T × bytes_per_element

 

구체적인 수치로 확인해보면 규모가 바로 와닿는다. LLaMA 2 7B(L=32, d_model=4096, FP16)에서 시퀀스 길이 2048짜리 단일 요청 하나를 처리하면 KV cache가 딱 1GB다.

2 × 32 × 4096 × 2048 × 2 = 1,073,741,824 bytes ≈ 1GB

 

배치 크기 8, 시퀀스 길이 4096으로 늘리면 16GB가 된다.

8 × 2 × 32 × 4096 × 4096 × 2 = 17,179,869,184 bytes ≈ 16GB

 

참고로 LLaMA 2 7B 모델 가중치 자체가 FP16 기준 약 14GB다. 즉 긴 컨텍스트 요청 8개를 동시에 처리하면 KV cache가 모델 가중치보다 더 많은 메모리를 차지한다. 시퀀스가 길어지거나 배치가 커질수록 이 격차는 더 벌어진다. 배치 크기까지 곱하면 서빙 환경에서 완전히 bottleneck이 되는 이유다.

---

KV Compression: 다 저장 안 해도 되는 거 아님?

실제로 attention을 보면 모든 토큰에 균등하게 집중하지 않는다. 소수의 토큰만 높은 attention score를 받는다. 중요한 것만 남기고 나머지는 버리자는 게 KV compression의 출발점이다.

버린 토큰은 나중에 다시 계산하지 않는다. 영구 삭제다. 그래서 핵심 가정이 하나 있다.

"버리는 토큰은 어차피 attention에 거의 영향 없는 토큰이다"

 

이 가정을 얼마나 잘 지키냐가 각 방법의 차별점이다.

---

기존 방법들의 흐름

가장 단순한 StreamingLLM은 초반 "sink" 토큰 몇 개와 최근 window만 남긴다.

Softmax는 attention score 합이 반드시 1이 되도록 강제하기 때문에, 현재 query가 어떤 토큰과도 연관이 없더라도 어딘가에 score를 할당해야 한다. 자기회귀 모델 특성상 모든 후속 토큰에 항상 노출되는 초기 토큰이 이 역할을 맡게 되는데, 이게 attention sink다. 그래서 초기 토큰을 버리면 softmax 분모의 상당 부분이 사라져 attention 분포가 무너지고 성능이 급락한다. 단순하지만 중간에 중요한 토큰이 있어도 그냥 버린다는 한계가 있다.

 

Figure 4 : the KV cache of Streaming LLM

 

이를 보완하려고 H2O는 attention score를 누적해서 "heavy-hitter" 토큰을 찾는다. 여러 번 attention 받은 토큰 = 중요한 토큰이라는 직관이다. 매 decoding step마다 누적 score를 업데이트하고, budget 초과 시 score 낮은 토큰을 evict한다. heavy-hitter만 남기는 게 아니라 recent 토큰도 일부 강제 보존한다.

# 설명을 위해 batch=1로 단순화한 코드
def update(self, new_key, new_value, attention_weights):
    # 기존 토큰들에 대한 attention score 누적 (heads 평균)
    attn_update = attention_weights.mean(dim=1).squeeze(-2)
    self.cumulative_attention = self.cumulative_attention + attn_update
    self.keys = torch.cat([self.keys, new_key], dim=2)

    if current_size > self.max_cache_size:
        self._evict()

def _evict(self):
    # recent 토큰 먼저 확보
    _, recent_indices = positions.topk(self.recent_size)
    recent_mask[recent_indices] = True

    # 나머지 중에서 heavy-hitter 선정
    non_recent_attention[recent_mask] = float("-inf")
    _, heavy_hitter_indices = non_recent_attention.topk(self.heavy_hitter_size)

    keep_mask = recent_mask | heavy_hitter_mask  # 둘 다 보존

 

SnapKV는 접근이 다르다. budget을 초과하는 순간 압축하는데, prefill 시점에 한 번 동작하도록 설계된 방법이다. 프롬프트 끝의 observation window query들이 prefix 전체에 attention을 계산해서 중요한 토큰을 "투표"로 선정한다. window 내 query들의 attention score를 평균내는 방식으로 각 prefix 토큰의 중요도를 집계하는 것이다. 단순히 top-k만 뽑으면 context 연속성이 깨지기 때문에 pooling으로 선택된 토큰 주변도 같이 보존한다.

def update_kv(self, key_states, query_states, value_states):
    # 마지막 window_size개 query로 window 이전 전체에 attention 계산
    attn_weights_sum = nn.functional.softmax(
        attn_weights[:, :, -self.window_size:, :-self.window_size],
        dim=-1,
        dtype=torch.float32,  # 수치 안정성을 위해 fp32로 계산
    ).mean(dim=-2)            # heads별 평균으로 집계 (투표)

    # max_pool1d로 주변 토큰도 같이 챙김 (pooling)
    attn_cache = F.max_pool1d(attn_weights_sum, kernel_size=self.kernel_size, ...)

    indices = attn_cache.topk(self.budget - self.window_size, dim=-1).indices

    # window 이전 구간에서 선택된 것 + window 전체 concat
    k_past_compress = key_states[:, :, :-self.window_size, :].gather(dim=2, index=indices)
    k_cur = key_states[:, :, -self.window_size:, :]  # window는 통째로 보존
    key_states = torch.cat([k_past_compress, k_cur], dim=2)

 

VATP는 attention score만 보면 안 된다는 문제의식에서 나왔다. attention sink 토큰은 score가 높아도 value vector의 norm이 거의 0이라 output에 기여하는 게 없다. 그래서 importance를 이렇게 재정의한다.

importance = attention_score × norm(value_vector)

 

정리하면 이런 흐름이다.

• 위치 규칙만 봄 (StreamingLLM)
• → attention 누적 (H2O)
• → prefill 시점 예측 (SnapKV)
• → score 보정 (VATP)

---

공통 한계: 계산해봐야 안다

위 방법들(H2O, SnapKV, VATP)은 전부 post-RoPE query로 attention score를 계산해서 중요도를 추정한다. 여기서 근본적인 문제가 있다.

RoPE는 query를 위치에 따라 회전시킨다. position 100의 query랑 position 10000의 query는 같은 의미라도 방향이 완전히 다르다. 그래서 최근 몇 개의 query만 가지고 score를 추정하면 대표성이 없다.

 

선행 연구(LazyEviction, Zhang et al., 2025)에서 확인된 바로는, observation window를 늘려도 성능이 개선되지 않았고 약 25개 query에서 peak를 찍은 뒤 오히려 하락했다. reasoning 모델에서는 이 문제가 더 심각하다. 어떤 토큰은 수천 토큰 뒤에야 중요해지는데, 그 전에 low score로 evict되면 chain of thought가 끊겨버린다.

 

TriAttention은 방향을 바꾼다.

"실제 attention 계산 없이, RoPE 구조만으로 어떤 토큰이 중요한지 예측할 수 있다"

---

핵심 관찰: pre-RoPE Q/K는 뭉쳐있다

RoPE 적용 후(post-RoPE) Q 벡터는 위치마다 방향이 달라서 대표값을 뽑기 어렵다. 그런데 RoPE 적용 전(pre-RoPE) 공간을 보면 다르다.

같은 헤드, 같은 band의 Q 벡터들이 서로 비슷한 방향을 가리키는 경향이 있다.

q_나   = [0.80, 0.21]
q_는   = [0.79, 0.23]  ← 비슷한 방향
q_학교 = [0.81, 0.20]  ← 비슷한 방향

이걸 Mean Resultant Length R로 수치화한다.

Rf = ||E[qf]|| / E[||qf||]

R = 1 → 모든 벡터가 같은 방향 → 완벽히 뭉침
R = 0 → 벡터가 사방으로 퍼짐

Qwen3-8B 기준으로 약 90%의 attention head에서 R > 0.95가 나왔다. 그리고 이 center는 입력 content나 위치가 달라져도 안정적으로 유지된다. 모델 가중치에 내재된 특성이기 때문이다.

---

 

 

왜 거리만으로 예측 가능한가

내적과 cos

두 벡터의 내적은 이렇게 표현된다.

내적 = ||v1|| × ||v2|| × cos(각도차이)

같은 방향이면 cos=1, 수직이면 cos=0, 반대면 cos=-1. attention의 logit이 바로 이 내적이다 (softmax 전 raw 값).

RoPE = 위치만큼 편각을 더하는 것

RoPE는 Q, K를 위치에 따라 회전시킨다. 복소수로 표현하면:

q̃f(p) = qf · e^(iωf·p)

e^(iθ) = cos(θ) + i·sin(θ)는 복소평면에서 각도 θ만큼 회전하는 연산이다. 길이는 항상 1이라서 (||e^(iθ)|| = 1) qf에 곱하면 크기는 그대로고 편각만 ωf·p만큼 더해진다.

내적하면 cos이 나오는 이유

복소수 곱셈에서 길이는 곱해지고 각도는 더해진다. Q (position pq)와 K (position pk)를 내적하면:

Re(q̃f(pq) · k̃f(pk)*)
= Re(qf · kf* · e^(iωf·(pq-pk)))
= ||qf|| ||kf|| cos(ωf·∆ + ϕf)

여기서 ∆ = pq - pk (거리), ϕf = arg(qf) - arg(kf) (원래 방향 차이). 모든 band에 대해 합산하면:

logit = Σf ||qf|| ||kf|| cos(ωf·∆ + ϕf)

RoPE가 "회전"이기 때문에 내적하면 자연스럽게 cos이 나온다.

Q/K가 뭉쳐있으면 ∆만 변수가 된다

Q/K가 뭉쳐있으면 center로 근사할 수 있다: qf ≈ q̄f, kf ≈ k̄f. 이러면 ||qf||, ϕf가 전부 상수가 된다. cos 합차 공식으로 전개하면:

logit(∆) = Σf [ af·cos(ωf·∆) + bf·sin(ωf·∆) ]

af =  ||q̄f|| ||k̄f|| cos(ϕf)   ← 상수
bf = -||q̄f|| ||k̄f|| sin(ϕf)   ← 상수

변수가 ∆ 하나뿐이 됐다. Q/K center만 알면 실제 attention 계산 없이 거리 ∆만으로 logit을 예측할 수 있다. 이게 TriAttention의 핵심 통찰이다.

---

이론이 실제로 맞는가: Reconstruction Correlation

여기까지는 "center로 근사하면 logit이 ∆만의 삼각급수가 된다"는 이론적 유도였다. 그런데 이 근사가 실제 attention과 정말 맞아떨어질까? 논문은 이걸 직접 검증한다.

center에서 계산한 삼각급수 예측값과 실제 attention logit을, 로그 간격 거리(∆ = 1, 2, 4, 8, ...)에서 비교한 Pearson 상관계수를 측정한다. 이걸 Reconstruction Correlation r̄이라 부른다. 로그 간격으로 샘플링하는 이유는 가까운 거리(개수가 많음)가 상관계수를 지배하지 않게 해서 장거리 패턴까지 균형 있게 반영하기 위해서다.

 

Qwen3-8B의 첫 레이어 첫 헤드(cherry-picking을 피하려고 일부러 첫 헤드 선택)에서 r̄ = 0.72가 나왔고, Qwen3 / Qwen2.5 / Llama3 세 아키텍처 전반에서 평균 0.5 이상이 나왔다. 즉 center만으로 계산한 삼각급수가 실제 attention 패턴을 상당히 정확하게 재현한다는 뜻이다.

 

이게 중요한 이유는, "pre-RoPE가 뭉쳐있다"는 관찰과 "center로 중요도를 scoring하겠다"는 방법 사이를 잇는 다리이기 때문이다. 뭉쳐있다는 사실만으로는 부족하고, 그 center가 실제 attention 동작을 예측한다는 게 입증돼야 scoring의 타당성이 성립한다.

 

---

TriAttention Scoring

이론을 바탕으로 실제 scoring은 두 가지 신호를 합산한다.

Strig: 거리 선호도 기반 score

Strig(k, ∆) = Σf ||E[qf]|| · ||kf|| · cos(ωf·∆ + ϕf)

여기서 Q는 calibration에서 구한 center E[qf]로 대체하고, K는 실제 cached key 값을 그대로 사용한다. K를 center로 근사하지 않는 이유는 K가 이미 cache에 존재하는 고정된 값이기 때문이다. Q는 미래 위치에서 어떻게 올지 불확실하지만, K는 이미 확정된 값이라 실제 값을 쓰는 게 더 정확하다.

 

key 하나가 미래의 여러 위치에서 query될 수 있으니 geometric spacing({1, 2, 4, ..., 2^16})으로 여러 offset에서 score를 계산해 평균낸다 (기본값 기준).

실제 코드에서의 구현

한 가지 짚고 넘어갈 게 있다. vLLM은 K를 RoPE 적용 후(post-RoPE)로 저장한다. 논문은 pre-RoPE K를 가정하는데, 실제 구현에서는 K에 rotation이 이미 baked in 돼있어서 수식을 다르게 전개한다.

# K_rot = K_unrot * e^{i*p*omega} 이므로
# 먼저 invert_rope로 post-RoPE K를 pre-RoPE로 복원한 뒤 scoring
k_unrot = invert_rope(k_values, cos_table, sin_table, attention_scale, style=rope_style)

# Q_center * conj(K_rot) 로 직접 계산
prod_real = q_mean_real * k_real + q_mean_imag * k_imag
prod_imag = q_mean_imag * k_real - q_mean_real * k_imag

# query 위치 t에서의 phase만 적용
t = round_start + offset
phase = t * omega
score = freq_scale_sq * (prod_real * cos(phase) - prod_imag * sin(phase))
score = score.sum(dim=-1)  # band 합산

 

수학적으로는 논문 수식과 동일한 결과다. K에 이미 들어간 position rotation과 query side phase가 합쳐져서 결국 ωf·∆가 된다.

Snorm: norm 기반 보완 score

R이 낮은 헤드, 즉 Q/K가 덜 뭉쳐있는 헤드에서는 삼각급수 근사가 부정확하다. 이를 보완하기 위해 norm 기반 score를 더한다.

# (E[||qf||] - ||E[qf]||) = (1 - Rf) · E[||qf||] 에 해당
extra_coef = q_abs_mean - q_mean_abs  # [num_kv_heads, freq_count]
extra_term = (k_abs * extra_coef * freq_scale_sq).sum(dim=-1)

R이 높으면 q_abs_mean ≈ q_mean_abs라서 extra_coef가 거의 0이 돼 Snorm의 기여가 자연스럽게 줄어든다. R이 낮으면 차이가 커져서 Snorm이 더 많이 반영된다.

 

최종 score는:

S(k) = Strig(k) + Snorm(k)

---

실제 동작 흐름

1단계: calibration (사전 계산)

모델을 실제로 돌리기 전에, 소량의 calibration 데이터로 pre-RoPE 공간에서 각 헤드, 각 band의 Q center(E[qf])와 R을 계산해서 저장해둔다. 논문 실험에서 calibration 데이터 품질이나 양에 크게 민감하지 않았다.

 

Google 홈페이지 HTML 같은 저품질 데이터로 calibration해도 reasoning 데이터로 한 결과와 비슷했다. Q/K center가 모델 가중치에 내재된 특성이기 때문이다. 

2단계: inference 중 scoring

compression trigger 조건은 두 가지를 동시에 충족해야 한다. decode 토큰 수가 budget에 도달한 상태이고, 생성된 절대 위치(absolute_position)가 128의 배수인 타이밍일 때다. "budget 초과 후 128 step 기다린다"는 의미가 아니라, 절대 위치 기준으로 128의 배수인 순간에 trigger된다는 뜻이다.

 

scoring 대상은 decode 토큰뿐이다. 기본 설정에서 prefill 토큰(초기 프롬프트)은 항상 보존되고, decode 토큰만 남은 budget 안에서 경쟁한다. 단, allow_prefill_compression=True로 설정하면 prefill 토큰도 압축 대상에 포함된다.

 

vLLM은 K를 post-RoPE 상태로 저장한다. 그래서 scoring 전에 invert_rope로 RoPE를 역으로 풀어서 pre-RoPE K를 복원한 다음, 모든 layer의 sampled head에서 Strig + Snorm으로 score를 계산해 각 head별 점수 행렬을 만든다.

3단계: eviction (union-based selection)

단순히 score 기준 상위 B개를 뽑는 게 아니라 union-based selection을 사용한다. 동작은 다음 순서로 이루어진다.

1. 각 sampled head가 독립적으로 top-k 토큰을 선정한다.
2. 모든 head의 선정 결과를 합집합(union)으로 모은다.
3. union에서 combined score 기준으로 최종 B개를 추린다.
4. union이 B보다 작으면 나머지는 score 순으로 채운다.

이 방식의 의미는 "어떤 head 하나라도 중요하다고 판단한 토큰은 후보에 남긴다"는 것이다. 특정 head에서만 높은 점수를 받는 토큰이 global 평균에 묻혀 사라지는 걸 방지한다. 최종 선정된 B개 index로 K, V tensor를 slicing하고 나머지는 영구 삭제한다. 버린 토큰은 복구하지 않는다.

---

기존 방법과의 차이

기존 방법 (H2O, SnapKV):
  post-RoPE query로 Q · K 내적 계산
  → position마다 query 방향이 달라서 대표성 없음
  → 좁은 window에서만 scoring 가능
  → 단순 top-k로 최종 선정

TriAttention:
  pre-RoPE center로 삼각급수 계산
  → position과 무관하게 안정적
  → 실제 attention 계산 없이 모든 token scoring 가능
  → union-based selection으로 head별 중요 토큰 보존

 

---

결과

AIME25 기준으로 Full Attention 대비 처리량 2.5배 향상 또는 KV 메모리 10.7배 감소를 달성한다. 같은 메모리 budget에서 기존 방법(R-KV) 대비 정확도가 약 15%p 높다. MATH 500에서는 32K 토큰 중 1024개만 KV cache에 남겨도 Full Attention(69.6%)과 거의 동일한 정확도(68.4%)를 달성했다.

 

물론 근사 기반이라 한계는 있다. Q/K가 잘 뭉쳐있지 않은 헤드에서는 오차가 생길 수 있고, calibration이 필요하다. 기존 방법들이 "어떤 score로 중요도를 측정할까"에 집중했다면, TriAttention은 "RoPE 구조 자체를 이용해서 score 계산을 대체하자"는 방향으로 접근한 게 차별점이다.

 

 

 

TriAttention 주소

https://github.com/WeianMao/triattention