[onnxruntime] ONNX Runtime WebGPU QMoE: 단일 토큰 라우팅 최적화로 성능 대폭 향상

2026년 10월 9일수정: 2026년 10월 9일

PR 링크: microsoft/onnxruntime#33058 상태: Merged | 변경: +413 / -49

들어가며

최근 Microsoft의 ONNX Runtime 레포지토리에서는 [WebGPU] Reduce QMoE single-token gate sorting to stable TopK라는 제목의 Pull Request(PR)가 머지되었습니다. 이 PR은 양자화된 Mixture-of-Experts (QMoE) 모델에서 단일 토큰 라우팅 시 발생하는 성능 병목 현상을 해결하는 데 중점을 둡니다. 특히, 기존의 복잡하고 비효율적인 정렬 방식을 더 빠르고 안정적인 TopK 기반 방식으로 개선하여, WebGPU 환경에서의 QMoE 연산 속도를 크게 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 이 최적화가 가지는 일반적인 기술적 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 onnxruntime/contrib_ops/webgpu/moe/ 디렉토리 내의 게이트(gate) 로직 구현에 집중되어 있습니다. 특히 단일 토큰 처리 시 라우팅 전문가(expert)를 선택하는 방식이 개선되었습니다.

1. `gate_1token.wgsl.template` 변경

가장 중요한 변경은 WebGPU 셰이더 템플릿 파일인 gate_1token.wgsl.template에서 이루어졌습니다. 기존에는 단일 토큰 라우팅 시 모든 전문가를 대상으로 버블 정렬(bubble sort)과 유사한 방식으로 Top-K 전문가를 선별했습니다. 이는 k 값이 작을 때는 괜찮을 수 있으나, 전문가 수가 많아질수록 비효율적이었습니다.

Before (기존 로직 추정 - 버블 정렬 방식):

    // K is small, use a simple bubble sort
    for (var i = 0u; i < workgroup_size_x - 1u; i++) {
        for (var j = 0u; j < workgroup_size_x - 1u - i; j++) {
            // ... (bubble sort logic) ...
        }
        workgroupBarrier();
    }

After (개선된 TopK 방식):

#if normalize_routing_weights
    var candidate_val = max_val;
    var candidate_idx = local_idx;
    for (var i = 0u; i < K; i++) {
        shared_vals[local_idx] = candidate_val;
        shared_idxs[local_idx] = candidate_idx;
        workgroupBarrier();
        for (var remaining = workgroup_size_x; remaining > 1u; remaining = (remaining + 1u) / 2u) {
            let stride = (remaining + 1u) / 2u;
            if (local_idx < remaining - stride) {
                let other = local_idx + stride;
                if (shared_idxs[local_idx] >= cols ||
                    (shared_idxs[other] < cols &&
                     (shared_vals[other] > shared_vals[local_idx] ||
                      (shared_vals[other] == shared_vals[local_idx] && shared_idxs[other] < shared_idxs[local_idx])))) {
                    shared_vals[local_idx] = shared_vals[other];
                    shared_idxs[local_idx] = shared_idxs[other];
                }
            }
            workgroupBarrier();
        }
        if (local_idx == 0u) {
            selected_vals[i] = shared_vals[0];
            selected_idxs[i] = shared_idxs[0];
        }
        workgroupBarrier();
        if (local_idx == selected_idxs[i]) {
            candidate_val = -MAX_FLOAT; // Invalidate selected expert
            candidate_idx = cols;
        }
        workgroupBarrier();
    }
    if (local_idx < K) {
        shared_vals[local_idx] = selected_vals[local_idx];
        shared_idxs[local_idx] = selected_idxs[local_idx];
    }
    workgroupBarrier();
#else
    // ... (original bubble sort logic retained for non-normalized case) ...
#endif

개선된 로직은 normalize_routing_weights 플래그에 따라 분기됩니다. 이 플래그가 활성화된 경우 (즉, 라우팅 가중치가 정규화될 때), K개의 상위 전문가를 찾기 위해 병렬 Top-K 리덕션 알고리즘을 사용합니다. 이 알고리즘은 각 워크그룹 내에서 K개의 최대 로짓(logit) 값과 해당 인덱스를 효율적으로 찾아냅니다. 특히, 동일한 로짓 값을 가진 경우 낮은 전문가 인덱스를 우선하는 안정적인(stable) 정렬 방식을 유지합니다. 또한, 이미 선택된 전문가는 다음 라운드에서 제외되도록 처리하여(-MAX_FLOAT로 설정) 중복 선택을 방지합니다.

normalize_routing_weights가 비활성화된 경우, 기존의 전체 전문가를 대상으로 하는 소프트맥스(softmax) 계산 방식이 유지됩니다. 이는 연산 결과의 일관성을 보장하기 위함입니다.

2. `gate_1token.h` 파일 추가

gate_1token.h 파일은 새로운 Gate1TokenProgram 클래스를 정의합니다. 이 클래스는 k (선택할 전문가 수), is_fp16 (FP16 사용 여부), has_router_weights (라우터 가중치 존재 여부), normalize_routing_weights (정규화 여부) 등의 파라미터를 받아 셰이더 코드를 생성하는 역할을 합니다. 이는 기존의 GateProgram 클래스와 유사하지만, 단일 토큰 처리에 특화된 로직을 캡슐화합니다.

3. 테스트 케이스 강화 (`moe_test.cc`)

PR은 단순히 로직을 변경하는 것을 넘어, 다양한 엣지 케이스를 다루는 테스트 케이스를 추가하거나 강화했습니다.

안정적인 Top-K 및 Tie-breaking: [0, -inf, -inf]와 같이 특정 로짓이 -inf인 경우, 또는 모든 로짓이 동일한 경우(k == num_experts)에도 올바른 전문가 인덱스가 선택되는지 검증합니다. 리뷰어들의 피드백에 따라, 기존 Comparator로는 -inf 로짓 처리 시 잘못된 인덱스가 선택되는 문제가 발견되었고, 이를 수정하고 직접 인덱스를 검증하는 테스트가 추가되었습니다.
```
// FP16/FP32 negative-infinity logits ([0, -inf, -inf] selects [0, 1])
// and all-equal logits with k == num_experts ([0, 1, 2, 3])
```
다양한 전문가 풀 크기: 3개, 5개, 6개, 512개 등 다양한 수의 전문가 풀에 대한 테스트를 추가하여, 특히 작은 수의 전문가 풀에서의 경계값(boundary tie) 처리를 검증합니다. 512개 전문가의 경우, 어댑터의 워크그룹 제한(workgroup limit)이 512 미만이면 해당 테스트를 건너뛰도록 수정되었습니다. 이는 실제 하드웨어 제약 조건을 반영한 것입니다.
```
// Cover boundary ties in 3-, 5-, 6-, and 512-expert pools
// Skip the 512-expert case when either adapter workgroup limit is below 512.
```
정규화되지 않은 라우팅: normalize_routing_weights=0인 경우, 모든 전문가의 로짓 합을 사용하여 정규화하는 기존 동작을 유지하고, 이에 대한 회귀(regression) 테스트를 추가했습니다. 이는 최적화 과정에서 기존 연산의 정확성이 변경되지 않음을 보장합니다.
```
// Exercise FP16/FP32 nonnormalized routing with an excluded expert contributing appreciably to the all-expert softmax denominator.
```

4. 문서 업데이트 (`docs/contrib_ops/webgpu/qmoe_gate.md`)

새로운 WebGPU QMoE 단일 토큰 게이트 구현에 대한 상세 문서가 추가되었습니다. 이 문서는 정규화된 Top-K, 안정적인 Tie-breaking 및 -inf 처리, 전체 전문가 기반의 정규화 폴백(fallback), 어댑터 제한, 그리고 벤치마크 범위 등을 설명합니다. 이는 기능 변경 사항에 대한 명확한 이해를 돕습니다.

왜 이게 좋은가?

이번 PR의 핵심 가치는 성능 향상과 안정성 확보에 있습니다.

성능 향상

PR 설명에 따르면, Qwen3.8-Flash-Next 모델 (512 전문가, k=10, FP16 I/O)을 사용한 H200 환경에서 다음과 같은 성능 개선이 관찰되었습니다:

16-토큰 사전 생성(Prefill): 11.64초에서 0.90초로 대폭 단축
안정적인 디코드(Steady Decode): 토큰당 약 0.674초에서 ~0.045초/토큰으로 감소

이는 단일 토큰 라우팅 로직이 전체 추론 시간에서 상당한 부분을 차지하며, 이 부분의 최적화가 전체 성능에 미치는 영향이 매우 크다는 것을 보여줍니다. 특히, 기존의 버블 정렬 방식 대신 병렬 Top-K 리덕션을 사용함으로써 GPU 워크로드를 더 효율적으로 활용하게 된 결과입니다.

또한, H200/Vulkan 환경에서 256개 전문가와 Top-4 라우팅을 사용하는 합성 단일 토큰 블록-FP8 QMoE 그래프 벤치마크에서는 다음과 같은 결과가 나왔습니다:

기존: 7.09 ms/run (평균)
개선 후: 5.56 ms/run (평균)

이는 약 1.28배의 성능 향상에 해당합니다. 비록 이 수치는 전체 추론 시간의 일부인 게이트 커널만의 속도 향상이며, 전체 모델 속도 향상을 의미하는 것은 아니지만, 핵심 연산의 효율성이 크게 개선되었음을 시사합니다.

안정성 및 정확성 확보

단순히 속도만 높이는 것이 아니라, 다양한 엣지 케이스에 대한 테스트 강화와 코드 수정을 통해 연산의 정확성과 안정성을 높였습니다.

정확한 Tie-breaking: 동일한 로짓 값에 대한 전문가 인덱스 선택 시, 기존 Comparator의 버그(예: [0, -inf, -inf]에서 잘못된 인덱스 선택)를 수정하고, 이를 직접 검증하는 테스트를 추가했습니다. 이는 모델의 예측 결과가 일관되고 정확하게 유지되도록 보장합니다.
하드웨어 제약 조건 고려: 512개 전문가를 사용하는 경우, 실제 GPU 어댑터의 maxComputeWorkgroupSizeX 및 maxComputeInvocationsPerWorkgroup 제한을 확인하고, 이를 초과할 경우 테스트를 건너뛰도록 하여 하드웨어 호환성을 높였습니다.
기존 동작 유지: normalize_routing_weights=0인 경우, 기존의 전체 전문가 기반 정규화 방식을 그대로 유지하는 테스트를 추가하여, 최적화가 연산의 의미를 변경하지 않음을 확인했습니다.

일반적인 교훈

이 PR은 다음과 같은 일반적인 기술적 교훈을 제공합니다:

알고리즘 선택의 중요성: 복잡한 정렬 알고리즘(예: 버블 정렬)은 특정 조건에서는 간단하지만, 대규모 데이터셋이나 병렬 처리 환경에서는 비효율적일 수 있습니다. 문제의 특성(예: Top-K 선택)에 맞는 효율적인 알고리즘(예: 병렬 Top-K 리덕션)을 선택하는 것이 성능 향상의 핵심입니다.
엣지 케이스 테스트의 중요성: 딥러닝 모델, 특히 MoE와 같이 복잡한 구조를 가진 경우, 다양한 입력 값(예: -inf, 동일 값, 경계값)과 하드웨어 설정에 대한 철저한 테스트가 필수적입니다. 리뷰어들의 날카로운 지적과 이를 반영한 테스트 강화는 코드의 안정성과 신뢰성을 크게 높입니다.
문서화의 가치: 새로운 기능이나 최적화가 적용될 때는, 해당 변경 사항의 목적, 동작 방식, 성능 영향, 그리고 테스트 범위 등을 명확하게 문서화하는 것이 중요합니다. 이는 다른 개발자들의 이해를 돕고, 향후 유지보수를 용이하게 합니다.
성능 측정의 신중함: PR 설명에서 언급된 것처럼, 특정 커널의 성능 향상이 반드시 전체 모델의 성능 향상으로 직결되는 것은 아닙니다. 전체 추론 파이프라인을 고려한 종합적인 성능 측정과 분석이 필요합니다.

결론

Microsoft ONNX Runtime의 이 PR은 QMoE 모델의 단일 토큰 라우팅 로직을 최적화하여 WebGPU 환경에서의 성능을 크게 향상시킨 훌륭한 사례입니다. 기존의 비효율적인 정렬 방식을 병렬 Top-K 리덕션으로 대체하고, 다양한 엣지 케이스에 대한 철저한 테스트와 문서화를 통해 정확성과 안정성까지 확보했습니다. 이러한 개선은 ONNX Runtime을 사용하는 개발자들이 더 빠르고 효율적인 딥러닝 모델 추론을 경험할 수 있도록 기여할 것입니다.

참고 자료

https://github.com/microsoft/onnxruntime/blob/main/docs/contrib_ops/webgpu/qmoe_gate.md

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

PR Analysis 의 다른글

이전글 [onnxruntime] ONNX Runtime WebGPU: Windows D3D12를 활용한 모델 로딩 가속화
현재글 : [onnxruntime] ONNX Runtime WebGPU QMoE: 단일 토큰 라우팅 최적화로 성능 대폭 향상
다음글 [sglang] MiniMax-H3 모델의 RMSNorm 및 AdaLN 연산 최적화: Triton을 활용한 커널 융합