기술 요약: 계층적 글로벌 어텐션 (Hierarchical Global Attention, HGA)

문제 정의

사전 학습된 롱 컨텍스트(Long-context) 거대 언어 모델(LLM)을 배포할 때 발생하는 주요 병목 현상은 단순히 밀집 인과적 어텐션(Dense Causal Attention)의 $O(n^2)$ 계산 복잡성 때문만이 아니라, 키/값(Key/Value, K/V) 캐시의 메모리 요구량 때문입니다. 밀집 어텐션에서는 모든 과거 토큰의 K/V 캐시가 가속기 메모리(VRAM)에 상주해야 합니다. 양자화된 대규모 모델(예: FP8)의 경우, 모델 가중치 자체가 VRAM 용량의 대부분을 차지하여 롱 컨텍스트(예: 32K–64K 토큰)를 위한 K/V 캐시를 저장할 충분한 공간이 남지 않게 됩니다. 이러한 제한은 재학습이나 구조적 변경 없이 기존의 사전 학습된 체크포인트를 소비자용 또는 단일 GPU 하드웨어에서 "즉시(out-of-the-box)" 추론하는 것을 방해합니다.

기존의 희소 어텐션(Sparse Attention) 방식은 종종 처음부터 다시 학습해야 하거나, 보정(Calibration) 파라미터를 도입하거나, 어텐션 출력 자체를 근사하여 학습된 표현의 품질을 저하시킬 가능성이 있습니다. 따라서 기존의 사전 학습된 체크포인트를 활용하면서도, 최종 어텐션 계산의 정확성을 유지하며 희소하고 메모리 효율적인 검색을 수행할 수 있는 드롭인(Drop-in) 메커니즘이 필요합니다.

방법론: 계층적 글로벌 어텐션 (HGA)

HGA는 사전 학습된 트랜스포머의 밀집 인과적 어텐션을 대체하기 위해 설계된 시스템 레벨의 패치입니다. 이는 원래의 체크포인트 파라미터($W_Q, W_K, W_V, W_O$)를 보존하며, 재학습이나 보정을 요구하지 않습니다.

핵심 메커니즘

HGA는 장거리 검색 문제를 밀집된 토큰 간 매질화(Materialization)에서 2단계 계층적 라우팅 프로세스로 전환한 뒤, 이후 정밀한 토큰 레벨 어텐션을 수행합니다.

1. 분할(Partitioning): 입력 시퀀스는 고정된 크기의 청크(예: $C=64$ 토큰)로 나뉩니다. 각 청크는 다시 그룹으로 세분화됩니다.
2. 계층적 라우팅(Hierarchical Routing):
   • 레벨 1 (청크 라우팅): 현재 청크는 이전 청크들로부터 유도된 압축된 "청크 키 요약(Chunk Key Summaries)" 테이블을 쿼리합니다.
   • 레벨 2 (그룹 라우팅): 선택된 청크 내에서, 쿼리는 선택을 정교화하기 위해 특정 "그룹 키 요약"으로 라우팅됩니다.
   • 결정론적 가시성(Deterministic Visibility): 라우팅 정책은 콘텐츠 기반 검색과 고정된 가시성 윈도우를 결합합니다. 즉, 현재 청크, 설정된 수의 초기 "싱크(Sink)" 청크, 그리고 최근의 로컬 청크는 항상 가시적 상태로 유지됩니다.
3. RoPE 인식 요약(RoPE-Aware Summaries): 회전 위치 임베딩(Rotary Position Embeddings, RoPE)과의 호환성을 보장하기 위해, HGA는 혼합 전략을 사용하여 요약을 구성합니다. 고주파 RoPE 쌍은 평균을 내기 전에 토큰별로 회전시키고, 저주파 쌍은 원시 키 공간에서 평균을 낸 후 해당 청크의 대표 위치에서 회전시킵니다. 이러한 요약은 새로운 선형 레이어 없이 기존의 투영된 키(Projected Keys)로부터 계산됩니다.
4. 정밀 토큰 어텐션(Exact-Token Attention): 관련 청크나 그룹이 선택되면, 라우터는 호스트 RAM(또는 NVMe)에서 원래의 토큰 레벨 K/V 벡터를 VRAM으로 가져옵니다. 최종 어텐션 출력은 이 제한된 정밀한 토큰 집합에 대해 표준 소프트맥스를 사용하여 계산됩니다. 합성된 요약 값은 최종 출력에 사용되지 않습니다.

메모리 아키텍처

HGA는 계층적 K/V 저장 추상화를 도입합니다.

• Hot (고온): 모델 가중치, 청크 요약, 항상 가시적인 윈도우는 VRAM에 상주합니다.
• Warm (온도): VRAM 내의 제한된 LRU 캐시가 최근에 라우팅된 청크들을 보유합니다.
• Cold (저온): 전체 이력 토큰 K/V 캐시는 호스트 RAM(또는 NVMe)에 상주합니다.
  이를 통해 GPU 메모리 소비를 전체 컨텍스트 길이로부터 분리합니다. VRAM 사용량은 주로 모델 가중치와 라우팅된 작업 집합(Working Set)에 따라 결정되며, 전체 시퀀스 길이에 의존하지 않습니다.

주요 기여

1. 드롭인 호환성: 가중치, 노름(Norm), 또는 보정 파라미터를 수정하지 않고도 기존 GQA/MoE 트랜스포머에서 밀집 어텐션을 대체할 수 있는 메커니즘을 제공합니다.
2. 정밀 토큰 검색: 어텐션 출력을 근사하는 방식과 달리, HGA는 오직 검색 단계만을 근사하여 최종 어텐션이 선택된 토큰들에 대해 정확하게 계산되도록 보장합니다.
3. 2단계 라우팅 정책: 청크 및 그룹 요약을 활용하여, 결정론적 싱크/로컬 윈도우와 콘텐츠 기반 검색을 결합한 전략을 사용합니다.
4. 계층적 저장 시스템: 전체 K/V 이력이 호스트 RAM에 존재할 수 있도록 하는 구현을 통해, 제한된 VRAM을 가진 하드웨어(예: 단일 RTX 5090)에서도 롱 컨텍스트 추론을 가능하게 합니다.
5. 제로샷 성능: 미세 조정 없이 32K–64K 컨텍스트에서 Qwen3-30B-A3B 모델을 구동할 수 있음을 입증했습니다.

실험 결과

본 논문은 40M 파라미터 규모의 SmallLM과 30B 파라미터 규모의 Qwen3 모델에 대한 벤치마크를 제시합니다.

• 품질 (손실 격차):
  • 직접적인 가중치 복사(미세 조정 없음)를 적용한 40M SmallLM에서, 라우팅된 모델은 8K 토큰에서 밀집 어텐션 대비 +0.018 nats의 손실 격차를 달로 달성했습니다.
  • 4K에서 64K 사이의 컨텍스트 길이에 걸쳐, 약 3%에서 12.5%의 희소도(Sparsity)를 사용할 때 손실 격차는 밀집 어텐션 대비 0.01–0.02 nats 이내를 유지했습니다.
  • 64K 토큰에서 6.25%의 희소도는 약 0.01 nats의 격차를 유지했습니다.
• 검색 정확도:
  • 64K 토큰 환경의 Qwen3-30B-A3B에 대한 "Needle-in-a-Haystack" 테스트에서, 그룹 레벨 라우팅과 약 1.9%의 희소도를 사용하여 미세 조정 없이 100% 통과율(3/3 깊이)을 달려냈습니다.
• 성능:
  • 40M SmallLM에서 Triton 퓨즈(Fused) 구현을 통해, 12K 토큰 기준 밀집 베이스라인 대비 학습에서 2.72배, 프리필(Prefill)에서 2.43배의 속도 향상을 달성했습니다.
  • Qwen3-30B의 경우, 패치된 모델은 단일 RTX 5090 (32GB VRAM)에서 32K 컨텍스트로 대화형 실행이 가능하며, 이는 밀집 K/V 저장이 불가능한 구성입니다.
• 미세 조정: 0.6B 모델을 라우팅 어텐션으로 미세 조정했을 때, 밀집 베이스라인 대비 미미한 라우팅 비용(+0.015 nats)과 함께 약 917 tokens/sec의 처리량을 기록했습니다 (밀집 모델의 ~570 대비).

의의 및 주장

본 논문은 HGA가 롱 컨텍스트 LLM의 메모리 병목 현상에 대한 실질적인 시스템 레벨 솔루션을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 사전 학습된 지식 보존: 미세 조정을 피하고 최종 출력을 위해 정밀 토큰 K/V를 사용함으로써, HGA는 기존 체크포인트가 학습된 표현을 저하시키지 않고도 희소 어텐션의 이점을 누릴 수 있게 합니다.
• 하드웨어 효율성: 밀집 캐싱에 필요한 VRAM 용량을 훨씬 초과하는 컨텍스트에 대해서도 소비자용 하드웨어에서 대규모 양자화 모델을 배포할 수 있게 합니다.
• 오차 귀인(Error Attribution): 저자들은 남아있는 작은 품질 격차(0.01–0.02 nats)가 컨텍스트 길이에 따라 급격히 증가하지 않는다는 점을 관찰했습니다. 이들은 이 격차가 라우팅 알고리즘 자체보다는 **롱 컨텍스트 위치 인코딩 외삽(Extrapolation)**에 의해 지배될 가능성이 높다고 제안하며, 이는 HGA가 위치 표현 개선을 통해 더욱 최적화될 수 있는 견고한 검색 메커니즘임을 시사합니다.

저자들은 HGA를 트랜스포머 아키텍처의 재설계가 아니라, 현재의 하드웨어 제약 조건 하에서 롱 컨텍스트 추론을 가능하게 하는 호환 가능한 계층적 라우팅 레이어로 정의합니다.