기술 요약: NOSA: Native 및 Offloadable Sparse Attention

문제 정의

대규모 언어 모델(LLM)의 디코딩 처리량(throughput)은 주로 KV 캐시(Key-Value cache)의 저장 요구량으로 인한 GPU 메모리 제한에 의해 점점 더 제약을 받고 있습니다. KV 캐시 오프로딩(중복된 컨텍스트를 CPU에 보관하고 일부를 GPU로 가져오는 방식)은 메모리 압박을 완화하지만, 기존 방식들은 다음과 같은 결정적인 트레이드오프를 겪습니다:

1. 훈련이 필요 없는 오프로딩(Training-Free Offloading): ShadowKV나 InfLLM과 같은 방법들은 훈련 없이 희소한 부분 집합(sparse subsets)을 선택합니다. 이는 **훈련-추론 불일치(training-inference mismatch)**를 유발하며, 훈련 시 보지 못한 희소 패턴이 성능 저하를 일으켜 특히 긴 생성 시나리오에서 선택 오류가 누적되는 결과를 초래합니다.
2. 학습 가능한 희소 어텐션(Trainable Sparse Attention): 희소 패턴을 학습하는 방법(예: InfLLMv2)은 이러한 불일치를 피할 수 있지만, 효율적인 오프로딩과 호환되지 않습니다. 제약 조건이 없다면 이 방식들은 CPU로부터 임의의 KV 블록을 가져올 수 있으며, 이는 처리량 이득을 상쇄할 만큼 막대한 PCIe 통신 오버헤드를 발생시킵니다.

현재의 솔루션들은 높은 디코딩 효율성(오프로딩을 통한)과 높은 생성 품질(훈련 일관성을 통한)을 동시에 달一个하지 못하고 있습니다.

방법론

저자들은 KV 캐시 오프로딩을 위해 명시적으로 설계된 학습 가능한 희소 어텐션 메커니즘인 NOSA(Native and Offloadable Sparse Attention)와 특화된 추론 시스템인 NOSI를 제안합니다.

1. 알고리즘: NOSA

NOSA는 디코딩 단계 전반에서 희소 어텐션 패턴이 안정적으로 유지되도록 하여 CPU-GPU 데이터 전송을 최소화하기 위해, 훈련 중에 **국소성 제약(locality constraint)**을 도입합니다.

• 이중 선택 메커니즘(Dual-Selection Mechanism): NOSA는 KV 선택을 두 가지 구성 요소로 분해합니다:
  • 쿼리 인식 선택 ($\Gamma_q$): 쿼리를 사용하여 특정하고 잠재적으로 먼 정보를 검색하여 재현율(recall)을 보존합니다.
  • 쿼리 무관 선택 ($\Gamma_e$): DMA에서 영감을 받은 증발 헤드(eviction head)를 사용하여 쿼리와 독립적으로 중요도 점수를 할당합니다. 이 구성 요소는 엄격한 **증발 정책(eviction policy)**을 강제합니다: 만약 어떤 토큰이 단계 $t$에서 선택되지 않았다면, 그 토큰은 이후의 단계 $t' > t$에서도 선택될 수 없습니다.
• 국소성 보장(Locality Guarantee): 쿼리 무관 선택에 고정된 예산 $k_e$를 할당함으로써, NOSA는 수학적으로 단계 $t$에서 선택된 KV 블록의 대다수가 단계 $t-1$에서도 선택되었음을 보장하여 $\gamma(t) \geq k_e/k$의 최소 국소성을 보장합니다. 이는 선택된 KV 블록들이 GPU에 계속 캐싱될 수 있도록 합니다.
• 블록 단위 구현(Block-Wise Implementation): PCIe 대역폭 활용을 극대화하기 위해, NOSA는 개별 토큰이 아닌 블록 단위로 동작합니다. 선택 점수를 계산하기 위해 2단계 풀링 전략(서브 블록 내 평균 풀링 후 블록 레벨에서의 최대 풀링)을 사용합니다.
• ED-DMA (Exp-Delayed DMA): 증발 헤드의 수치적 정밀도 문제를 처리하기 위해, NOSA는 지수 연산을 어텐션 계산 단계까지 지연시켜 지수 적용 전의 점수를 기반으로 블록을 선택합니다.

2. 시스템: NOSI

저자들은 vanilla 구현체(예: Hugging Face Transformers)가 실제 통신 병목 현상을 가리는 과도한 커널 런칭 오버헤드를 발생시킨다는 점을 관찰했습니다. NOSI는 이러한 병목을 드러내고 최적화하기 위해 구축된 맞춤형 추론 시스템입니다:

• 커널 퓨전(Kernel Fusion): 증발 헤드 계산, QKV 분할, 풀링 연산을 단일 Triton 커널로 융합하여 런칭 오버헤드를 줄입니다.
• 최적화된 메모리 레이아웃: KV 캐시를 최소 블록 크기 단위로 구성하고, 연속적인 메모리 요구 사항을 피하면서 블록 스와핑을 위한 집합 차이(set-differences)를 효율적으로 계산하기 위해 커스텀 CUDA 커널을 사용합니다.
• 고대역폭 통신: 통합 가상 주소 지정(Unified Virtual Addressing, UVA)을 사용하여 커스텀 호스트-투-디바이스 전송 커널을 구현함으로써, 표준 PyTorch 구현의 파편화 문제를 극복하고 근사적인 피크 PCIe 대역폭을 달성합니다.

주요 기여

1. 타당성 분석: 본 논문은 학습 가능한 희소 어텐션(InfLLMv2와 같은)이 본질적인 국소성을 가지고 있지만, 명시적인 제약 없이는 PCIe 대역폭 한계를 극복하기에 불충분하다는 점을 식별했습니다.
2. NOSA 알고리즘: 증발 정책을 통해 국소성 제약을 강제함으로써 CPU-GPU KV 전송량을 명시적으로 제한하는 새로운 학습 가능한 희소 어텐션 메커니즘을 제안합니다.
3. NOSI 시스템: 소프트웨어 오버헤드를 제거하여 통신을 진정한 병목 지점으로 드러내고 고처리량 오프로딩을 가능하게 하는, NOSA에 최적화된 추론 시스템입니다.
4. 포괄적 평가: 1B, 3B, 8B 모델에 대한 광범적인 벤치마크를 통해 NOSA가 훈련이 필요 없는 오프로딩 베이스라인 및 다른 학습 가능한 희소 어텐션 방법들보다 우수함을 입증했습니다.

실험 결과

저자들은 일반 작업, 긴 컨텍스트 입력, 긴 생성 추론 작업을 대상으로 1B, 3B, 8B 모델에 대해 NOSA를 평가했습니다.

• 성능: NOSA는 긴 컨텍스트 벤치마크(LongBench, HELMET)에서 훈련이 필요 없는 오프로딩 베이스라인(ShadowKV, InfLLM, ArkVale) 및 학습 가능한 베이스라인(InfLLMv2, DMA)보다 일관되게 우수한 성능을 보입니다. 훈련-추론 불일치로 인해 훈련이 필요 없는 방법들이 크게 저하되는 지점에서도 높은 재현율과 추론 능력을 유지합니다.
• 처리량(Throughput):
  • NOSA는 오프로딩을 통해 가능해진 더 큰 배치 크기 덕분에 FullAttn 대비 최대 5.04배 높은 디코딩 처리량을 달성합니다.
  • 국소성 제약으로 인한 통신 감소 덕분에 InfLLMv2보다 최대 1.92배 높은 처리량을 달성합니다.
  • 희소 패턴 처리의 효율성과 시스템 최적화 덕분에 ShadowKV보다 최대 1.83배 높은 처리량을 달성합니다.
• 긴 생성(Long-Generation): 긴 생성 추론 작업에서 훈련이 필요 없는 오로딩 방법들은 오류 누적과 퍼플렉시티(perplexity) 폭발을 보이는 반면, NOSA는 낮은 퍼플렉시티를 유지합니다.
• 절제 연구(Ablation): 연구는 블록 단위 선택 패턴이 PCIe 효율성에 결정적이라는 점과 ED-DMA 증발 헤드 구현이 최상의 성능을 낸다는 것을 확인했습니다.

의의 및 주장

본 논문은 NOSA가 "훈련이 필요 없는" 방식이나 "제약 없는 학습 가능한" 방식에서 **알고리즘-시스템 공동 설계(algorithm-system co-design)**로의 필수적인 전환을 나타낸다고 주장합니다. 훈련 중에 KV 전송량을 명시적으로 제한함으로써, NOSA는 메모리 효율성(오프로딩)과 생성 품질(훈련 일관성)의 간극을 메웁니다.

저자들은 NOSA가 긴 컨텍스트 능력을 갖춘 LLM을 생성 품질의 저하 없이 메모리가 제한된 하드웨어에 배포할 수 있게 한다고 밝힙니다. 또한, 디코딩 효율성을 개선하여 컴퓨팅 및 에너지 비용을 낮춤으로써 LLM 애플리케이션의 광범한 배포를 촉진할 수 있다고 강조합니다. 코드 또한 네이티브 오프로딩 메커니즘에 대한 추가 연구를 장려하기 위해 공개되었습니다.