[onnxruntime] ONNX Runtime WebGPU: NVIDIA GPU를 위한 FlashAttention 루프 언롤링 최적화

PR 링크: microsoft/onnxruntime#33395 상태: Merged | 변경: +44 / -4

들어가며

최근 microsoft/onnxruntime 레포지토리에 흥미로운 성능 최적화 PR이 병합되었습니다. 이 PR은 WebGPU 백엔드에서 NVIDIA GPU를 사용할 때, FlashAttention의 핵심 연산인 QK(Query-Key) 및 PV(Projected-Value) 리덕션 과정을 루프 언롤링(Loop Unrolling)하여 성능을 개선하는 내용을 담고 있습니다. 특히 대규모 언어 모델(LLM)의 프리필(Prefill) 단계에서 발생하는 병목을 효과적으로 해결했습니다.

코드 분석

1. C++ 코드: flash_attention.cc

핵심 변경 사항은 셰이더 코드를 생성하는 GenerateShaderCode 함수 내부에 있습니다. 특정 조건(is_nvidia_ && use_shm_path_ && qkv_head_size_ / 4 <= 32)을 만족할 때, 루프 대신 직접 연산을 나열하는 코드를 생성합니다.

// Before: 루프를 통한 일반적인 리덕션
// After: 런타임에 루프를 언롤링하여 셰이더 생성
const bool unroll_shm_inner = is_nvidia_ && use_shm_path_ && qkv_head_size_ / 4 <= 32;
if (unroll_shm_inner) {
  auto& impl = shader.AdditionalImplementation();
  impl << "fn flash_qk_unrolled(k : u32) -> qk_precision {\n"
       << "  var score = qk_precision(0);\n";
  for (int i = 0; i < qkv_head_size_ / 4; ++i) {
    impl << "  score += dot(vec4<qk_precision>(q_tile[" << i
         << "u]), vec4<qk_precision>(k_tile[k][" << i << "u]));\n";
  }
  // ... (이하 생략)
}

2. WGSL 템플릿: flash_attention.wgsl.template

생성된 언롤링 함수를 실제 셰이더 로직에 적용합니다. #if unroll_shm_inner 전처리기를 사용하여 기존의 루프 방식과 최적화된 방식을 분기 처리합니다.

// Before
for (var k = 0u; k < max_k_step; k++) {
  var score = qk_precision(0);
  for (var i = 0u; i < head_size_vec; i++) {
    score += dot(vec4<qk_precision>(q_tile[i]), vec4<qk_precision>(k_tile[k][i]));
  }
}

// After
#if unroll_shm_inner
  var score = flash_qk_unrolled(k);
#else
  // ... 기존 루프 로직
#endif

왜 이게 좋은가

이 최적화는 NVIDIA GPU의 공유 메모리(Shared Memory) 아키텍처를 최대한 활용합니다. 루프를 언롤링함으로써 컴파일러가 인덱스 계산을 상수화(Constant Folding)하고, 레지스터 할당 및 파이프라인 효율을 극대화할 수 있게 됩니다.

성능 지표를 보면, 4096 토큰 입력 시 Subgroup Matrix 모드에서 최대 80.08%의 성능 향상을 보였습니다. 다만, 모든 헤드 사이즈에서 이득을 보는 것은 아닙니다. PR 작성자는 헤드 사이즈가 512인 경우 오히려 성능이 크게 퇴보(Regression)하는 현상을 확인하고, qkv_head_size_ / 4 <= 32라는 엄격한 가드레일을 설정하여 최적화의 범위를 제한했습니다. 이는 무조건적인 최적화가 아닌, 아키텍처 특성을 고려한 정교한 튜닝의 중요성을 보여줍니다.

교훈

1. 조건부 최적화: 모든 케이스에 최적화를 적용하는 것이 항상 정답은 아닙니다. 특정 임계값(Threshold)을 넘어서면 오히려 오버헤드가 발생할 수 있음을 인지해야 합니다.
2. 데이터 기반 의사결정: 리뷰 과정에서 헤드 사이즈별 성능 데이터를 상세히 비교하여 최적의 지점을 찾아낸 점이 인상적입니다.
3. 테스트의 중요성: 89개의 템플릿 테스트와 고정밀도 검증을 통해, 성능 향상과 동시에 결과값의 정확성(Byte-identical)을 완벽히 보존했습니다.

참고 자료

• https://www.w3.org/TR/WGSL/
• https://onnxruntime.ai/docs/execution-providers/WebGPU-ExecutionProvider.html

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