[onnxruntime] WebGPU 성능의 한계를 넘어서: GatedDeltaNet 커널 최적화 분석

PR 링크: microsoft/onnxruntime#32645 상태: Merged | 변경: +316 / -84

들어가며\n\n최근 브라우저 기반의 AI 추론이 가속화되면서 WebGPU는 ONNX Runtime의 핵심 실행 프로텍터(Execution Provider)로 자리 잡고 있습니다. 특히 GatedDeltaNet과 같은 최신 재귀적 신경망 구조는 시퀀스 모델링에서 뛰어난 성능을 보이지만, WebGPU 환경에서는 하드웨어 리소스 제약과 연산 오버헤드라는 두 가지 큰 벽에 부딪히곤 합니다.\n\n이번에 분석할 PR은 GatedDeltaNet의 WebGPU 커널을 최적화하여 불필요한 중복 연산을 제거하고, 메모리 대역폭을 극대화하는 벡터화 I/O를 도입한 사례입니다. 시니어 엔지니어의 관점에서 이 변경사항들이 왜 중요한지, 그리고 실제 코드에서 어떻게 구현되었는지 심층 분석해 보겠습니다.\n\n## 코드 분석: 핵심 변경 사항\n\n### 1. Dispatch Sizing 및 오버플로우 방지 (gated_delta_net.cc)\n\n가장 먼저 눈에 띄는 변화는 커널 실행을 위한 워크그룹 할당(Dispatch) 로직의 수정입니다. 기존 코드에서는 32비트 정수 연산 과정에서 발생할 수 있는 오버플로우 위험이 있었고, 이는 잘못된 워크그룹 수 할당으로 이어질 수 있었습니다.\n\ncpp\n// Before\n.SetDispatchGroupSize((onnxruntime::narrow<uint32_t>(total_tokens * hv) + 63u) / 64u)\n\n// After\n.SetDispatchGroupSize(onnxruntime::narrow<uint32_t>(\n (static_cast<uint64_t>(total_tokens) * static_cast<uint64_t>(hv) + 63u) / 64u))\n\n\ntotal_tokens와 hv(head_size_v)의 곱이 uint32_t 범위를 넘어설 경우, 잘못된 워크그룹 크기가 계산되어 연산이 누락되거나 불필요한 오버헤드가 발생할 수 있습니다. 이를 uint64_t로 캐스팅하여 안전하게 계산한 뒤 narrow하게 변환하는 방식은 대규모 텐서를 다루는 커널에서 필수적인 방어적 프로그래밍입니다. 또한, PR 설명에 따르면 이 수정을 통해 기존에 발생하던 64배의 과도한 워크그룹 할당 문제를 해결했습니다.\n\n### 2. 워크그룹 수준의 연산 공유 (gated_delta_net.wgsl.template)\n\nGPU 연산에서 가장 비싼 작업 중 하나는 exp()나 sigmoid() 같은 초월 함수(Transcendental functions)입니다. 기존에는 모든 스레드가 루프 내부에서 독립적으로 gate_scale을 계산했습니다.\n\nwgsl\n// Before (Inside loop)\nfor (var token = first_token; token < last_token; token++) {\n var gate_scale = 1.0;\n // ... 매 토큰마다 모든 스레드가 gate 계산 수행 ...\n}\n\n// After (Workgroup Shared Memory)\nvar<workgroup> gate_scale: f32;\n\n// ... Inside loop ...\nif (key_index == 0u) {\n // 워크그룹 내의 단 한 스레드만 gate_scale을 계산\n // ... gate_scale 계산 로직 ...\n}\nworkgroupBarrier(); // 모든 스레드가 계산 완료를 기다림\n\n\n변경 후에는 var<workgroup>을 사용하여 gate_scale을 공유 메모리에 배치했습니다. 워크그룹 내의 첫 번째 스레드(key_index == 0u)만 이 비싼 연산을 수행하고, workgroupBarrier()를 통해 결과를 공유합니다. 이는 수천 개의 스레드가 동일한 값을 중복 계산하던 낭비를 획기적으로 줄여줍니다.\n\n### 3. 벡터화 I/O (vec4) 도입\n\n메모리 대역폭 최적화의 꽃은 벡터화입니다. WebGPU에서 f32 하나씩 읽는 것보다 vec4<f32>로 128비트를 한 번에 읽는 것이 훨씬 효율적입니다.\n\ncpp\n// gated_delta_net.cc\nconst bool vectorized_value_io = !use_packed_qkv && dv % kValueChannelsPerWorkgroup == 0;\nconst int value_io_components = vectorized_value_io ? onnxruntime::narrow<int>(kValueChannelsPerWorkgroup) : 1;\n\nprogram.AddOutput({output, ProgramTensorMetadataDependency::Type, value_io_components});\n\n\n이 코드는 조건이 충족될 경우(데이터 크기가 워크그룹 크기의 배수이고 버퍼가 패킹되지 않았을 때) 출력 타입을 벡터 타입으로 설정합니다. 이를 통해 셰이더 내부에서는 f32 대신 vec4를 사용하여 메모리 접근 횟수를 1/4로 줄일 수 있게 됩니다.\n\n## 왜 이게 좋은가: 성능과 설계의 교훈\n\n### 1. 메모리 계층 구조의 영리한 활용\nGPU는 연산 능력에 비해 메모리 대역폭이 병목이 되는 경우가 많습니다. 이번 PR은 vec4를 통한 메모리 병합(Memory Coalescing) 최적화와 워크그룹 공유 메모리를 통한 연산 재사용을 동시에 달성했습니다. 이는 전형적이지만 가장 강력한 GPU 최적화 패턴입니다.\n\n### 2. WebGPU 제약 사항에 대한 유연한 대응\n리뷰 댓글에서 언급되었듯, WebGPU는 기본적으로 8개의 Storage Buffer 바인딩 제한을 가집니다. 입력 텐서가 많아지면 ONNX Runtime은 버퍼를 하나로 합치는 packed_qkv 모드로 전환됩니다. 이번 최적화는 이러한 제약 조건(!use_packed_qkv)을 체크하여, 최적화가 가능한 상황에서만 벡터화를 활성화하도록 설계되었습니다. 이는 모든 환경에서 동작하면서도 특정 환경에서는 최대 성능을 내는 'Graceful Enhancement'의 좋은 예시입니다.\n\n### 3. 정확한 Dispatch의 중요성\nGPU 커널에서 Dispatch 사이즈를 잘못 계산하는 것은 성능 저하뿐만 아니라 잘못된 결과(Silent Data Corruption)를 초래할 수 있습니다. 64배나 많은 워크그룹이 생성되었다는 것은 GPU 스케줄러에 엄청난 부하를 주었음을 의미합니다. 이를 바로잡은 것만으로도 커널 런칭 오버헤드가 크게 줄어듭니다.\n\n## 결론\n\n이번 PR은 단순히 코드를 깔끔하게 만드는 것을 넘어, WebGPU라는 플랫폼의 특성을 깊이 이해하고 하드웨어 성능을 끝까지 쥐어짜는 최적화의 정석을 보여줍니다. 특히 초월 함수 연산을 워크그룹 단위로 캐싱하고, 조건부 벡터화를 통해 I/O 효율을 높인 점은 다른 WebGPU 커널 개발 시에도 반드시 참고해야 할 기법입니다.

참고 자료

• https://gpuweb.github.io/gpuweb/ — WebGPU 공식 사양 문서
• https://onnxruntime.ai/docs/execution-providers/WebGPU-ExecutionProvider.html — ONNX Runtime WebGPU 실행 프로바이더 가이드
• https://developer.mozilla.org/en-US/docs/Web/API/GPUDevice/limits — WebGPU 장치 제한 사항(Storage Buffer limit 등) 설명

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