[onnxruntime] ONNX Runtime WebGPU: FlashAttentionDecodeQKV 성능 최적화 분석

PR 링크: microsoft/onnxruntime#29586 상태: Merged | 변경: +10 / -2

들어가며

최근 LLM 추론 성능 최적화는 AI 엔지니어링의 핵심 과제입니다. 특히 브라우저 환경에서 WebGPU를 활용한 추론 시, GPU의 하드웨어 자원을 얼마나 효율적으로 활용하느냐가 전체 토큰 생성 속도(TPS)를 결정합니다. 이번에 분석할 microsoft/onnxruntime의 PR은 FlashAttentionDecodeQKV 연산의 Workgroup 크기와 타일링 파라미터를 조정하여, GPU 점유율(Occupancy)을 높이고 연산 비용을 획기적으로 줄인 사례입니다.

코드 분석

이번 최적화는 onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc 파일 내에서 수행되었습니다. 핵심은 m_tile == 1인 경우(즉, 디코딩 단계)에만 더 공격적인 하드웨어 설정을 적용하는 것입니다.

1. 타일링 파라미터 최적화

기존에는 tile_size_k_vec이 8로 고정되어 있었으나, 이를 디코딩 단계에서 32로 확장했습니다.

// Before
const uint32_t tile_size_k_vec = 8;

// After
const uint32_t tile_size_k_vec = (m_tile_ == 1u) ? 32u : 8u;

이 변경은 MatMulNBitsProgram과 동일한 하드웨어 활용 패턴을 따르도록 설계되었습니다. m_tile == 1인 디코딩 시에는 메모리 제약이 상대적으로 덜하여 더 큰 타일링이 가능하지만, 프리필(prefill) 단계에서는 공유 메모리 제한(32KB)을 고려하여 기존 설정을 유지하는 세심함이 돋보입니다.

2. Workgroup Size 조정

Workgroup 크기를 64에서 128로 상향 조정하여 병렬 처리 능력을 극대화했습니다.

// Before
program.SetWorkgroupSize(64)

// After
const uint32_t workgroup_size = (m_tile == 1u) ? 128u : 64u;
program.SetWorkgroupSize(workgroup_size)

sub_tile_count는 WorkgroupSizeX() / tile_size_k_vec로 계산되므로, 결과적으로 8에서 4로 조정되어 연산 효율성이 최적화되었습니다.

왜 이게 좋은가

이 최적화는 GPU의 연산 유닛을 더 촘촘하게 채우는 'Occupancy' 개선에 초점을 맞추고 있습니다. Qwen3-1.7B 모델을 기준으로 측정한 결과는 매우 인상적입니다.

• GPU 실행 시간: 0.83 ms → 0.47 ms (약 43% 감소)
• 토큰 생성 처리량: 257 tps → 282 tps (약 10% 향상)

교훈

1. 하드웨어 특성 반영: 범용적인 설정보다는 연산 단계(Decode vs Prefill)에 따른 하드웨어 제약 조건을 분리하여 최적화하는 것이 중요합니다.
2. 성공적인 벤치마킹: MatMulNBits와 같은 고성능 커널의 디스패치 형태를 벤치마킹하여 유사한 성능을 끌어내는 전략은 매우 효과적인 최적화 방법론입니다.
3. 제약 조건 준수: 성능 향상을 위해 무작정 파라미터를 키우기보다, 공유 메모리 제한(32KB)과 같은 하드웨어적 한계를 고려하여 안전한 분기 처리를 구현해야 합니다.

참고 자료

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

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