기술 요약: GPU 에서의 희소 Transformer 추론 가속화

문제 제기

대규모 언어 모델 (LLM) 은 Multi-Head Attention(MHA) 이 핵심 계산 구성 요소인 Transformer 아키텍처에 크게 의존합니다. 계산량을 줄이기 위해 희소성을 유도하는 마스크 레이어 (예: 인과적, 슬라이딩 윈도우, 무작위) 가 널리 사용되지만, GPU 에서 다음과 같은 심각한 성능 문제를 야기합니다:

1. 분편화된 계산: 희소 마스크는 분산된 데이터 접근 패턴을 생성하여 메모리 대역폭 병목 현상을 악화시키고 GPU 자원의 효율적 활용을 방해합니다.
2. 유연하지 않은 최적화: 기존 고성능 MHA 구현체 (예: FlashAttention) 는 주로 인과적 또는 슬라이딩 윈도우와 같은 특정 연속 마스크 패턴만 지원합니다. Bigbird 나 무작위 어텐션과 같은 이산적이거나 비구조화된 희소성에는 어려움을 겪거나, 계산 후 점수 행렬을 재설정하는 비효율적인 방법으로 회피합니다.
3. 정적 연자 융합: 현재 딥러닝 프레임워크의 연자 융합 (operator fusion) 방식은 대부분 정적이거나 규칙 기반입니다. 다양한 애플리케이션 시나리오, 가변적인 시퀀스 길이, 서로 다른 모델 하이퍼파라미터에 적응하지 못합니다. 또한, 융합된 연자를 위한 최적 구성은 개별적으로 튜닝된 연자들의 합과 크게 달라 순차적 튜닝은 최적의 결과를 내지 못합니다.

방법론: STOF 프레임워크

저자들은 유연한 마스크 지원과 적응형 연자 융합을 통합하여 GPU 에서 희소 Transformer 추론을 가속화하기 위해 STOF 프레임워크를 제안합니다. STOF 는 두 가지 주요 모듈로 구성됩니다:

1. 통합 MHA 모듈

이 모듈은 다양한 마스크 패턴에 걸쳐 희소 어텐션을 효율적으로 계산하는 과제를 해결합니다.

• 이중 계층 저장 형식: 임의의 마스크 패턴 (구조화 및 비구조화 모두) 을 처리하기 위해 STOF 는 Block Compressed Sparse Row(BSR) 와 비트맵을 결합한 새로운 저장 형식을 도입합니다.
  • 외부 타일 (OT): 전역적으로 건너뛰는 블록을 나타냅니다. "전체 (Full)"OT(모든 요소 유효) 와 "부분 (Part)"OT(혼합된 유효/무효) 는 포인터 배열을 통해 추적됩니다.
  • 내부 타일 (IT): 블록 내 요소 분포를 나타냅니다. "부분"OT 내의 유효 요소는 64 비트 비트맵 (8x8 블록당 하나의 uint64) 으로 인코딩되어 전체 희소 행렬을 물리적으로 생성하지 않고도 정밀한 마스킹을 가능하게 합니다.
• 이중 커널 전략: STOF 는 분석 모델을 기반으로 두 가지 다른 커널 유형을 구현합니다.
  • 행 단위 커널 (Row-wise Kernel): 짧은 시퀀스를 가진 고도로 희소한 시나리오에 최적화되어 있습니다. 쿼리 ($Q$) 를 행으로 슬라이스하고, 워프 수준 셔플을 활용하며, 데이터 지역성을 극대화하기 위해 워프 간 동기화를 제거합니다.
  • 블록 단위 커널 (Block-wise Kernel): 더 넓은 시나리오를 위한 범용 커널입니다. $Q$, $K$, $V$를 서브블록으로 분할하고 공유 메모리 (SMEM) 와 비동기 데이터 복사를 활용하여 메모리 전송과 계산을 중첩시킵니다.
• 커널 선택: 분석 모델이 마스크 희소도와 시퀀스 길이를 기반으로 행 단위 커널 또는 블록 단위 커널 중 어느 것을 사용할지 결정하여 가장 효율적인 실행 경로를 보장합니다.

2. 연자 융합 모듈

이 모듈은 MHA 레이어를 넘어선 적응형 융합의 필요성을 해결합니다.

• 융합 방식 변환: STOF 는 해시 인코딩을 사용하여 융합 방식을 수치 표현으로 변환합니다. 계산 그래프의 인접 연자는 이진 배열 (융합 경계 표현) 로 매핑된 후, 특정 컴파일 템플릿(Triton 및 TileLang 을 통해 구현됨) 으로 디코딩됩니다.
• 계층적 검색 엔진: 정적 규칙 대신 STOF 는 최적의 융합 경계와 커널 파라미터를 결정하기 위해 2 단계 검색 메커니즘을 사용합니다:
  1. 초기화: 신경 해싱과 사전 정의된 규칙을 사용하여 빈번한 서브그래프 패턴을 식별하고 초기 융합 방식을 설정합니다.
  2. 2 단계 튜닝:
     • 1 단계 (융합 확장): 보상 기반 알고리즘 (확장, 점거, 경쟁) 을 사용하여 융합 경계를 반복적으로 확장하며, 이전 결과 캐시와 비교하여 성능 향상을 테스트합니다.
     • 2 단계 (파라미터 샘플링): 경계가 고정되면 커널 파라미터 (예: 블록 크기, 파이프라인 단계) 를 샘플링하여 특정 융합 방식에 대한 최적 구성을 찾습니다.

주요 기여

1. 종합적 분석: 저자들은 다양한 마스크 패턴과 추론 구성의 영향을 분석하여 MHA 계산과 하류 연자 융합 모두에서 최적화 기회를 드러냈습니다.
2. 통합 MHA 모듈: 고유한 저장 형식 (BSR + 비트맵) 과 최적화를 구현하며 분석적 선택 모델에 의해 통제되는 행 단위 및 블록 단위 커널을 구현한 새로운 모듈입니다.
3. 적응형 연자 융합: 수치 디코딩을 통해 융합 방식을 컴파일 템플릿으로 변환하고 2 단계 검색 엔진을 활용하여 동적으로 최적 설정을 결정하는 모듈입니다.
4. STOF 프레임워크: GPU 에서 유연한 마스크 패턴과 최적의 연자 융합을 가능하게 하는 완전한 추론 시스템입니다.

실험 결과

저자들은 PyTorch Native, FlashAttention-2(FA2), FlexAttention, ByteTransformer, MCFuser, Bolt 와 같은 최신 베이스라인에 대해 NVIDIA RTX 4090(Ada) 및 A100(Ampere) GPU 에서 STOF 를 평가했습니다.

• MHA 성능: STOF 는 MHA 계산에서 기존 최상위 MHA 구현체 (예: FlexAttention) 대비 최대 1.6 배의 속도 향상을 달성했습니다. 다양한 마스크 패턴 (인과적, 슬라이딩 윈도우, Longformer, Bigbird) 과 시퀀스 길이 (최대 4,096) 에서 일관된 우월성을 입증했습니다.
• 종단 간 추론: BERT, GPT-2, LLaMA, T5, ViT 와 같은 모델에서 종단 간 추론 시 최대 1.4 배의 속도 향상을 달성했습니다.
• 튜닝 효율성: STOF 튜닝 프로세스는 보상 기반 샘플링 및 캐싱 메커니즘 덕분에 대규모 시나리오에서 경쟁자 (MCFuser 및 Bolt) 보다 6.7 배에서 6.9 배 빠르게 수행되었습니다.
• 확장성: 초장 시퀀스 (최대 16k) 를 사용한 실험에서 STOF 는 PyTorch Compile 대비 최대 16.8 배의 성능 향상을 보였으며, 베이스라인이 실패한 시퀀스 길이에서도 메모리 효율성을 유지하며 Out-of-Memory(OOM) 오류를 방지했습니다.

중요성 및 주장

본 논문은 STOF 가 희소 Transformer 시나리오를 위해 유연한 마스크 패턴과 다양한 연자 융합 방식을 동시에 가능하게 하는 최초의 시스템이라고 주장합니다. 그 중요성은 손으로 튜닝된 희소 전용 커널과 자동화된 밀집 연자 융합 사이의 간극을 메우는 데 있습니다. 커스텀 MHA 커널과 생성적 컴파일 템플릿을 통합하는 완전한 스택을 제공함으로써 STOF 는 정적이거나 규칙 기반 접근법으로는 이전까지 달성할 수 없었던 더 넓은 최적화 기회를 확립합니다. 저자들은 그들의 접근 방식이 새로운 LLM 아키텍처 (예: Mixture-of-Experts) 와 차세대 GPU 로 확장 가능하며, 대규모 모델 추론에서 지속적인 성능 향상을 위한 길을 제시한다고 강조합니다.