기술 요약: ThriftAttention

문제 정의

장기 컨텍스트 작업에서 대규모 언어 모델 (LLM) 의 효율적인 추론은 현재 어텐션 메커니즘의 2 차 계산 비용과 메모리 트래픽에 의해 병목 현상을 겪고 있습니다. NVIDIA 의 Blackwell 아키텍처는 FP16 대비 4 배의 산술 처리량을 제공하는 네이티브 FP4 텐서 코어를 도입했으나, SageAttention3 와 같은 균일 양자화를 통해 어텐션에 이를 활용하면 시퀀스 길이가 증가함에 따라 특히 심한 품질 저하가 발생합니다. 반면, KV 블록을 완전히 제거함으로써 이를 완화하려는 기존 희소성 방법은 FP4 지연 시간과 일치시키기 위해 종종 75% 이상의 블록을 제거하는 공격적인 희소성 비율을 요구하며, 이로 인해 되돌릴 수 없는 성능 손실이 발생합니다. 장기 컨텍스트 설정에서는 추론 효율성과 출력 품질 사이에 근본적인 긴장 관계가 존재합니다.

방법론: ThriftAttention

저자들은 FP4 추론 효율성으로 FP16 에 근접한 품질을 제공하도록 설계된 훈련이 필요 없는 선택적 혼합 정밀도 어텐션 메커니즘인 ThriftAttention을 제안합니다. 이 방법을 이끄는 핵심 통찰력은 양자화 오차가 균일하게 분포되어 있는 것이 아니라, 오차의 영향이 매우 비균일하며 어텐션 점수가 가장 높고 최종 출력 분포에 기능적으로 가장 관련성이 높은 소수의 쿼리 - 키 블록 쌍에 집중된다는 점입니다.

알고리즘은 두 단계로 진행됩니다:

1. 휴리스틱 블록 선택:

   • 쿼리 ($Q$) 및 키/밸류 ($K, V$) 텐서는 블록으로 분할됩니다.
   • 경량의 휴리스틱이 토큰 평균의 내적을 사용하여 각 쿼리 - 키 블록 쌍 $(i, j)$에 대한 중요도 점수를 계산합니다: $\hat{S}_{ij} = \bar{q}_i \cdot \bar{k}_j$. 여기서 $\bar{q}_i$와 $\bar{k}_j$는 해당 블록의 평균 벡터입니다.
   • 점수가 가장 높은 상위-$k$개의 블록 쌍이 고정밀도 계산을 위해 식별됩니다.

2. 혼합 정밀도 계산 및 병합:

   • FP16 경로: 선택된 상위-$k$개의 블록 쌍은 전체 FP16 정밀도로 계산됩니다.
   • FP4 경로: 나머지 블록은 Blackwell GPU 에서 지원되는 NVFP4 마이크로스케일링 포맷을 활용하여 FP4 양자화를 사용하여 계산됩니다.
   • 온라인 병합: 두 경로 모두 계산되어 FlashAttention-2 와 유사한 온라인 소프트맥스 절차를 통해 병합되어 단일 출력을 생성합니다. FP4 경로는 SageAttention3 에서 적응된 확률 블록에 대한 2 단계 양자화 방식을 사용하여 누적 중 수치적 안정성을 유지합니다.

구현은 단일 퓨즈드 CUDA 커널로 이루어집니다. 레지스터 압력을 관리하기 위해 커널은 먼저 선택되지 않은 블록을 FP4 경로를 통해 처리한 다음, 승격된 블록을 위해 별도의 FP16 헬퍼 루틴을 활용하며, 필요한 경우에만 FP16 쿼리 타일을 로드합니다. 상위-$k$ 블록을 포함하지 않는 워프는 불필요한 메모리 로드를 방지하기 위해 FP16 경로를 완전히 우회합니다.

주요 기여

• 선택적 혼합 정밀도 어텐션: 가장 중요한 블록 상호작용을 FP16 으로 계산하고 나머지는 FP4 로 유지하는 ThriftAttention 의 도입. 저자들은 이것이 어텐션 계산에 대해 이러한 특정 혼합 정밀도 방식으로 서브바이트 포맷을 활용한 최초의 작업이라고 지적합니다.
• 훈련이 필요 없는 접근법: 이 방법은 추론 중 휴리스틱 선택 메커니즘에만 의존하며 모델 재학습이나 파인튜닝이 필요하지 않습니다.
• 종합적 평가: 이 접근법은 다양한 모델 계열 (Llama, Qwen, Ministral) 과 장기 컨텍스트 벤치마크 (LongBench-v1, HELMET, RULER, PG-19) 에서 평가되어 다양한 아키텍처와 컨텍스트 길이에서 유효성을 입증했습니다.

결과

• 품질 회복: 5% FP16 블록 예산 (쿼리 - 키 블록의 5% 만 FP16 으로 계산) 에서 ThriftAttention 은 FP4 와 FP16 어텐션 간의 성능 격차 평균 **89.1%**를 회복합니다. 이 회복률은 10% 및 25% 예산에서 각각 91.8% 및 92.4% 로 증가합니다.
• 효율성 향상:
  • 프리필 (Prefill): FlashAttention-2(FP16) 대비 최대 1.7 배의 커널 속도 향상을 달성하며, 131k 컨텍스트에서 약 1.2 배의 일관된 엔드투엔드 프리필 개선을 보입니다.
  • 디코드 (Decode): FlashAttention-2 대비 3 배~5.5 배의 속도 향상을 달성하여, 131k 컨텍스트 길이에서 전체 FP16 어텐션 대비 거의 2 배의 엔드투엔드 생성 속도 향상을 실현합니다.
• 컨텍스트 길이 확장: ThriftAttention 의 이점은 시퀀스 길이가 증가함에 따라 커집니다. 균일 FP4 어텐션은 컨텍스트 길이가 증가함에 따라 크게 저하되는 반면 (예: Llama3.1-8B 의 경우 8k 에서 50% 였던 FP4 유지율이 131k 에서 32% 로 감소), ThriftAttention 의 FP4 대비 상대적 개선은 증가하여 131k 에서 2.2 배에 달합니다.
• 희소성과의 비교: 계산량 (FP16 등가 FLOPs) 을 일치시켰을 때, ThriftAttention 은 추론 시 희소 어텐션 베이스라인 (Quest, Sparse Top-k 등) 보다 현저히 우수한 성능을 보입니다. 저자들은 희소 방법이 블록을 완전히 폐기하여 꼬리 오차를 발생시키는 반면, ThriftAttention 은 모든 상호작용을 FP4 로 유지하고 삭제가 아닌 양자화를 통해 점진적으로 저하시킨다는 사실에 기인한다고 설명합니다.

중요성 및 주장

이 논문은 ThriftAttention 이 FP4 지연 시간에 근접하면서 FP16 에 근접한 품질을 유지하는 장기 컨텍스트 추론을 위한 실용적인 경로를 제공한다고 주장합니다. 양자화 오차가 특정 고중요도 상호작용에 집중되어 있음을 식별함으로써, 이 방법은 장기 컨텍스트에서 균일 저비트 어텐션에서 관찰되는 체계적인 품질 저하를 완화합니다.

저자들은 이 작업을 Blackwell 시대 추론에서 효율성과 품질 사이의 "근본적인 긴장"을 해결하는 방안으로 위치시킵니다. 그들은 부분 집합만 고정밀도로 계산하기 위해 공격적으로 희소화하는 대신, 저정밀도 (FP4) 에서 전체 어텐션 지원을 유지하는 것이 일치된 계산 예산에서 더 나은 결과를 산출한다고 강조합니다. 이 접근법은 네이티브 FP4 텐서 코어를 활용하기 위해 소비자 및 데이터 센터용 Blackwell GPU 를 대상으로 하지만, 저자들은 현재 KV 캐시 메모리 발자국 (이중 캐싱으로 인해 28% 증가) 에 대한 제한 사항과 훈련이 아닌 추론에 초점을 맞추고 있다는 점을 지적합니다.