Long Context 구현 관점 정리: Sparse Attention, RoPE Scaling, KV Cache

Long Context를 구현할 때, '최대 context 길이를 128K로 변경'하는 하나의 설정만으로는 전체 그림을 알 수 없습니다. 적어도 다음 세 가지는 별도로 설계해야 합니다.

• Attention 연결 그래프 (Attention connection graph): 어떤 query-key 쌍을 계산할 것인가 -
• 위치 매핑 (position mapping): 토큰 위치를 RoPE에 어떻게 전달할 것인가 -
• KV 상태 (KV state): 디코딩 중에 몇 개의 토큰, 몇 개의 헤드 분량의 Key/Value를 유지할 것인가

Sparse Attention은 1번, RoPE Scaling은 2번, GQA/MQA 또는 KV 양자화는 3번을 주로 변경합니다. FlashAttention은 1번 연결을 바꾸지 않으면서 dense attention을 I/O 효율적으로 계산합니다.

이 글에서는 Longformer, Position Interpolation, YaRN을 중심으로 Long Context를 텐서 형태와 계산량 관점에서 정리합니다.

Dense Attention에서 늘어나는 텐서

batch와 head를 생략하고 시퀀스 길이를

Attention score는 다음입니다.

4096 tokenと8192 tokenのAttention matrixを面積とpair数で比較し、二乗増加を示す図

1 head, 1 layer당 스코어 쌍의 수는 다음과 같이 증가합니다.

| | |
|---|---|---|
| 4,096 토큰 대비 | | 1배 |
| ... | | |

Causal mask로 미래 쪽을 무효화하더라도, dense한 연결 그래프를 갖는 것은 변하지 않습니다. 구현이 상삼각행렬(upper triangle)을 물리적으로 계산/저장하는지에 따라 커널(kernel)에 차이가 있지만, 시퀀스 길이를 2배로 했을 때 Attention 계산이 무거워지는 이유는

Sparse Attention은 마스크가 아닌 연결 그래프로 생각해야 한다

Attention mask를

Sparse Attention의 비용은 개념적으로 행렬 전체의

하지만, dense matrix를 만든 후 불필요한 셀을

Longformer의 local + global 패턴

Longformer는 sliding-window attention과 task 의존적인 global attention을 결합합니다.

LongformerのAttention matrixで局所windowが対角bandを作り、global tokenが行と列を横断する構造

각 토큰
global 토큰 위치 집합을
Dense Attention의

global token은 라우팅 결과가 아닌 설계 입력

global attention은 '중요해 보이는 토큰을 모델이 자동으로 선택한다'는 의미가 아닙니다. Longformer 원 논문에서는 task에 따라 global position을 결정합니다. 분류(classification)라면 맨 앞의 특별한 토큰, question answering이라면 question token과 같은 방식입니다.

이는 구현상 다음 trade-off로 이어집니다.

| global token 수 | 원거리 연결 | 비용 (cost) | 설계상의 리스크 |
|---|---|---|
| 적음 | 제한적 | 작다 | 필요한 정보가 local path에만 머무르게 됨 |
| 많음 | 늘어남 | |
| sparsity의 장점이 약해짐 |

Longformer는 long-document encoder로 제안 및 평가된 model입니다. 이 패턴을 decoder-only model에 적용할 경우에는, causal 제약과 global edge의 방향을 별도로 설계해야 합니다.

local window의 receptive field

1 layer에서 각 측

하지만, 도달 가능성(reachability)과 직접 비교하는 것은 다릅니다. 거리

FlashAttention은 edge를 줄이지 않는다

FlashAttention은

| 방법 | 연결 edge | 산술적 쌍의 수 | 주로 줄이는 것 |
|---|---|---|
| Longformer형 Sparse Attention | local + global | 계산하는 pair | |
| ... | | |

실제 Long Context model에서는, sparse pattern과 I/O-aware kernel을 결합할 여지가 있습니다. 다만, 임의의 sparse mask가 accelerator 위에서 빠르다고는 할 수 없습니다. 규칙적인 block, tile size, load balance 등을 포함하여 측정해야 합니다.

RoPE는 위치 차이를 회전각으로 바꾼다

RoFormer의 RoPE (Rotary Position Embedding: 회전에 의한 위치 임베딩)는 query와 key의 2차원 성분을 position

위치

RoPEでqueryとkeyを位置ごとに回転し、内積に位置差の角度が残る仕組み

절대(absolute) position을 사용하여 회전시키면서, Attention score에는 상대적인 위치 차이가 나타납니다.

문제는 학습 시 최대 position을 초과했을 때입니다. position index만 크게 한다고 해서, 학습 중에 경험하지 못한 회전 패턴이나 상대 거리를 외삽(extrapolate)할 수는 없습니다. RoPE를 채택한 것 자체가 임의 길이로의 외삽 보장(guarantee)은 아닙니다.

Position Interpolation은 index를 압축한다

원래 학습 길이를

32768 tokenのpositionを4096 tokenの学習済み範囲へ8分の1に圧縮するPosition Interpolation

4,096에서 32,768로 8배 확장하는 경우, 대표 position의 대응은 다음과 같습니다.

학습되지 않은 거대한 index를 그대로 사용하지 않고, 학습된 범위로 보간(interpolate)하는 발상입니다. 원 논문에서는 LLaMA를 최대 32,768 token까지 확장하고, 1,000 step 미만의 fine-tuning으로 평가했습니다. 이는 해당 논문의 모델과 조건에 국한된 결과입니다.

압축하면 근거리 position 간의 각도 차이도 작아집니다. 장거리를 범위 안에 담는 대신, 단거리 position resolution을 변경하는 트레이드오프(trade-off)가 있습니다.

YaRN은 RoPE의 주파수 대역별 처리와 Attention scale을 조정하여 context extension을 안정화하는 방법입니다. 구현 시에는 'RoPE Scaling'이라는 하나의 boolean이 아니라, 다음 사항들을 확인합니다.

• 원래 학습 context 길이
• 목표 context 길이 및 scale factor
• 저주파/고주파 성분의 보간 방법
• 추가 학습 data의 길이 분포
• 단문 task와 장문 task 양쪽에서의 quality

RoPE Scaling으로 변하지 않는 것

position index를 대체하더라도,

RoPE Scaling이 직접적으로 변경하지 않는 것은 다음과 같습니다.

• Dense Attention의 pair 수
• prefill 시 계산 $QK^ op$ - decode 시 유지하는 KV token 수
• 긴 의존 관계를 학습한 data 양
• task 상의 유효 context 길이

'position error가 사라졌다'와 '계산할 수 있다', '정보를 사용할 수 있다'를 구분하여 테스트해야 합니다.

prefill과 decode를 분리하여 프로파일링한다

prefill에서는 prompt 전체

| 단계 | query 길이 | 참조하는 Key/Value | 대표적인 관측치 |
|---|---|---|
| Prefill | prompt 길이 | prompt 전체 | time to first token, peak activation |
| Decode | 보통 1 | 과거 | time per output token, KV bandwidth |

같은 128K context라도, '첫 번째 token이 느린 것'과 '생성 중인 token이 느린 것'은 원인이 다릅니다.

KV Cache를 128K token으로 추정한다

batch size 1의 KV Cache 본체는 다음 식으로 대략 계산할 수 있습니다.

약 15.6 GiB입니다. MHA에서

32 layer、128K tokenのKV CacheがGQA 8 KV headで15.6 GiB、MHA 32 headで62.5 GiBになる比較

이것은 1 sequence의 KV 본체만입니다. weight, activation, allocator overhead, padding, temporary buffer는 포함하지 않습니다. batch size나 동시 sequence 수도 곱해집니다.

GQA/MQA는

유효 context 길이를 task별로 측정한다

최대 길이까지 forward할 수 있음을 확인한 후에는, 그 길이로 정보를 이용할 수 있는지 측정합니다. RULER는 단일 Needle-in-a-Haystack뿐만 아니라, retrieval, multi-hop tracing, aggregation, question answering을 포함하는 13가지 task에서 길이와 복잡성을 변경합니다. 구현 시에는 최소한 다음 사항들을 sweep해야 합니다.

구현 체크리스트

Long Context 대응을 추가할 때는 다음 순서로 나누어 진행하면 원인을 추적하기 쉽습니다.

• 토크나이저 이후의 입력/출력 token budget 확정하기
• model이 학습한 원래 context 길이와 추가 학습 조건을 확인하기
• Attention pattern과 실제 kernel이 sparsity를 이용하는지 확인하기
• RoPE의 scale factor, 주파수 설정, position ID 기록하기
• KV Cache를 layer・KV head・dtype별로 계산하기
• prefill latency와 decode latency를 개별적으로 측정하기
• 단문 quality와 장문 quality를 동일한 변경 전후에서 비교하기
• retrieval, multi-hop, aggregation을 다른 task로서 평가하기

설정 파일에 max_position_embeddings가 늘어난 것만으로는 2부터 8까지의 완료를 의미하지 않습니다.

요약

Long Context 구현에서는 다음 세 가지 계층으로 나눕니다.

• Longformer형 Sparse Attention은 계산하는 query-key edge를 local + global로 한정합니다.
• Position Interpolation이나 YaRN은 RoPE에 전달할 position과 주파수 scale을 조정합니다.
• GQA/MQA나 KV quantization은 decode 중에 유지・전송되는 KV state를 줄입니다.

FlashAttention은 dense edge를 유지한 채 IO를 개선합니다. Context Parallelism은 sequence 영역을 device에 분산합니다. 어떤 기법도 Long Context 전체를 한 번에 해결하는 것은 아닙니다.

최종적으로는 최대 context 길이가 아니라, task별 quality, prefill/decode latency, peak memory를 동일한 조건에서 측정하여 판단합니다.

Longformer의 global attention, 128K token의 KV Cache 계산, RAG와의 선택 구분까지 포함하는 전체 그림은 개인 블로그의 Long Context 완전판에서 자세히 설명하고 있습니다.

참고 자료

• Longformer: The Long-Document Transformer
• RoFormer: Enhanced Transformer with Rotary Position Embedding
• Extending Context Window of Large Language Models via Positional Interpolation
• YaRN: Efficient Context Window Extension of Large Language Models
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• RULER: What's the Real Context Size of Your Long-Context Language Models?
• Lost in the Middle: How Language Models Use Long Contexts

토론