본문으로 건너뛰기

[flashinfer] Hopper(SM90)의 잠재력을 깨우는 Attention 최적화: FlashInfer 'Cake' 백엔드 분석

PR 링크: flashinfer-ai/flashinfer#5586 상태: Merged | 변경: +12056 / -8437

들어가며

NVIDIA의 Hopper 아키텍처(SM90)는 WGMMA(Warpgroup Level Matrix Multiply-Accumulate), TMA(Tensor Memory Accelerator), 그리고 Thread Block Cluster와 같은 혁신적인 기능을 도입하여 딥러닝 연산 성능을 비약적으로 향상시켰습니다. 하지만 이러한 하드웨어 기능을 소프트웨어 레벨에서 완벽하게 활용하여 최적의 Latency를 뽑아내는 것은 매우 도전적인 과제입니다.

FlashInfer 라이브러리의 이번 PR은 가변 블록 희소 어텐션(Variable Block-Sparse Attention, VSA)을 위한 cake 백엔드의 4번째 최적화 라운드입니다. 특히 'Small-selection' 커널, 즉 처리해야 할 블록의 수가 적은 케이스에서 발생하는 오버헤드를 줄이기 위해 Two-warpgroup 병렬화, Cluster Rendezvous 지연, FP32 Partial Accumulation 등의 기법을 도입했습니다. 이 글에서는 실제 코드 변경 사항을 바탕으로 시니어 엔지니어의 관점에서 이 최적화들이 왜 효과적인지 분석해 보겠습니다.

코드 분석: 핵심 변경 사항

1. Two-Warpgroup 기반의 병렬 처리 (Small CTA 최적화)

기존의 작은 커널들은 단일 Warpgroup(128개 스레드)에서 순차적으로 블록을 처리했습니다. 이번 변경에서는 하나의 CTA 내에서 두 개의 Warpgroup을 사용하여 작업을 분담합니다.

[Before & After 개념 분석]

  • Before: 단일 Warpgroup이 $N$개의 블록을 순차적으로 처리. QK -> Softmax -> PV 체인이 길게 이어짐.
  • After: Warpgroup 0은 짝수 블록을, Warpgroup 1은 홀수 블록을 처리. 이후 Warpgroup 1이 계산한 FP32 Partial 결과(max, sum, accumulator)를 Warpgroup 0에게 전달하여 최종 병합.

이 변경은 flashinfer/cake_vsa_sm90.py의 플래너(Planner) 비용 모델에 다음과 같이 반영되었습니다.

# Before (개념적 로직)
# cost = blocks * block_latency

# After (PR에 반영된 비용 모델 변경)
# 두 개의 warpgroup이 분담하므로 체인의 길이는 ceil(blocks / 2)가 됨
# handoff 비용(약 0.35-0.46 us)이 추가되지만 전체 Latency는 감소
SPLIT_MERGE_COST = 0.3
SPLIT_FIXED_COST = 2.0

왜 좋은가? Warpgroup 간의 의존성 체인을 절반으로 줄임으로써 연산 유닛의 유휴 시간을 최소화합니다. Warpgroup 1에서 0으로 데이터를 넘기는 데 드는 비용은 Shared Memory를 통한 32KB 라운드 트립 정도로, 약 0.4us 내외입니다. 이는 전체 연산 시간이 짧은 Small Kernel에서 병렬 이득이 통신 비용을 상회하는 지점을 정확히 공략한 것입니다.

2. Deferred Cluster Rendezvous (클러스터 동기화 지연)

Hopper의 Thread Block Cluster 기능을 사용할 때, CTA 간의 데이터 공유를 위해 barrier.cluster를 사용합니다. 기존에는 커널 시작 직후(Prologue)와 종료 직전(Exit)에 동기화를 수행했습니다.

[Before & After 로직 변경]

// Before
__cluster_barrier_all(); // 시작 시 모든 CTA 대기
// ... K/V 로딩 및 연산 ...
__cluster_barrier_all(); // 종료 시 모든 CTA 대기

// After (Deferred Prologue Rendezvous)
// 1. 초기화 및 K/V 블록 로딩을 동기화 없이 시작 (Overlap)
// 2. 실제 cross-CTA 트래픽인 st.async(push) 직전에만 대기
cluster_barrier.wait(); 
// 3. Exit barrier 제거 (mbarrier를 통해 데이터 수신 여부만 확인)

왜 좋은가? 커널 시작 시 모든 CTA가 준비될 때까지 기다리는 대신, 각 CTA가 독립적으로 수행할 수 있는 초기화 및 데이터 로딩 작업을 먼저 진행합니다. 이를 통해 Peer CTA들의 초기화 지연 시간을 Overlap 시킬 수 있습니다. PR 설명에 따르면 이 최적화만으로도 런칭당 약 1.0~1.3us의 고정 비용을 절감했습니다.

3. FP32 Partials 및 결정론적(Deterministic) 연산

병합 경로(Merge path)에서 FP32 정밀도를 유지하도록 수정되었습니다. 이는 특히 분산된 워크스페이스나 DSM(Distributed Shared Memory)을 통한 병합 시 수치적 안정성을 보장합니다.

// csrc/cake_vsa_sm90/sm_90a/ 내의 커널 코드 변경 (개념적)
// Accumulator를 FP16/BF16으로 캐스팅하지 않고 최종 단계까지 FP32 유지
// 이를 통해 테스트 시 FP32 Reference와 비트 단위 일치(Bit-exact) 달성

성능 결과 및 교훈

H100 SXM 환경에서의 벤치마크 결과는 놀랍습니다.

  • 최대 성능 향상: h4-m1024-n1024-k8 케이스에서 기존 대비 1.196x (약 20%) 성능 향상.
  • 평균 성능 향상: 전체 벤치마크 지표에서 기하평균 1.03x 향상.
  • 고정 비용 절감: Cluster Rendezvous 최적화로 인해 작은 워크로드에서의 Latency Floor가 낮아짐.

일반적인 교훈

  1. Latency Model의 중요성: 이번 PR은 단순히 "빠르게" 만든 것이 아니라, 측정된 Latency 모델(measured latency model)을 기반으로 Planner의 비용 함수를 수정하여 최적의 경로를 선택하도록 유도했습니다.
  2. 하드웨어 특화 최적화: Hopper의 Cluster Barrier와 Warpgroup 구조를 깊게 이해하고, 동기화 지점을 최소화하거나 지연시키는 기법은 고성능 GPU 프로그래밍의 핵심입니다.
  3. 정밀도와 성능의 트레이드오프: FP32를 사용하면서도 성능을 올릴 수 있었던 이유는 연산량 자체보다 메모리 접근 패턴과 동기화 오버헤드가 병목이었기 때문입니다.

결론

이번 FlashInfer의 업데이트는 Hopper 아키텍처에서 Attention 연산을 최적화하려는 엔지니어들에게 훌륭한 사례를 제시합니다. 특히 st.async를 통한 데이터 푸시와 mbarrier를 이용한 동기화, 그리고 Warpgroup 간의 효율적인 역할 분담은 향후 LLM 추론 가속화의 표준적인 기법이 될 것입니다.

FlashInfer 팀의 이러한 세밀한 최적화 덕분에 우리는 더 빠르고 정확한 오픈소스 Attention 커널을 사용할 수 있게 되었습니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글