본문으로 건너뛰기

[onnxruntime] ONNX Runtime QMoE SwiGLU GEMV 최적화: Split-K2 커널로 LLM 추론 가속화

PR 링크: microsoft/onnxruntime#29167 상태: Merged | 변경: +559 / -15

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께, 모델 추론(inference) 성능은 사용자 경험과 서비스 비용에 지대한 영향을 미치고 있습니다. 특히 Mixture of Experts(MoE)와 같은 희소(sparse) 모델은 효율적인 연산이 더욱 중요합니다. ONNX Runtime은 다양한 하드웨어에서 모델을 최적화하여 실행하는 데 핵심적인 역할을 합니다. 이번 글에서는 microsoft/onnxruntime 레포지토리의 PR ([CUDA]: Split-K2 QMoE SwiGLU GEMV kernel)을 분석하여, LLM 추론 시 발생하는 주요 병목 중 하나인 QMoE(Quantized Mixture of Experts) FC1(First Fully Connected) 레이어의 SwiGLU GEMV(General Matrix-Vector) 연산을 어떻게 Split-K2 기법으로 최적화했는지 살펴보겠습니다.

이 PR은 GPT-OSS-20B와 같은 모델에서 단일 토큰 디코드 시 QMoE FC1 interleaved-SwiGLU GEMV 경로가 상당한 시간을 소모한다는 문제 인식에서 출발했습니다. 목표는 K 차원을 분할하여 FC1의 병렬성을 개선하고, 이를 통해 전반적인 디코드 처리량을 향상시키는 것입니다.

핵심 최적화: Split-K2 SwiGLU GEMV 커널

이 PR의 핵심은 Split-K2라는 기법을 QMoE FC1 interleaved-SwiGLU GEMV 연산에 적용한 것입니다. Split-K2는 행렬 곱셈의 K(내부) 차원을 두 부분으로 나누어 처리한 후, 그 결과인 부분 합(partials)을 나중에 결합하는 방식입니다. 이는 GPU의 병렬 처리 능력을 더욱 효과적으로 활용하여 연산 효율을 높입니다.

두 단계(Two-Pass) 접근 방식

기존의 단일 커널 방식 대신, 이 최적화는 두 단계의 커널 실행으로 FC1 연산을 분리합니다.

  1. 첫 번째 단계: 부분 합 계산 (moe_gemv_splitk_partials_kernel)

    • 이 커널은 K 차원을 두 부분으로 나누어 각각의 부분 행렬-벡터 곱셈을 수행합니다.
    • 계산된 FP32 partials는 QMoE workspace에 저장됩니다.

  2. 두 번째 단계: 부분 합 감소 및 SwiGLU 적용 (moe_gemv_splitk_reduce_swiglu_kernel)

    • 이 커널은 첫 번째 단계에서 계산된 두 개의 FP32 partials를 감소(reduce)시켜 최종 결과를 얻습니다.
    • 선택적으로 bias를 추가하고, interleaved SwiGLU 활성화 함수를 적용하여 FC1의 최종 출력을 생성합니다.

FC2(Second Fully Connected) 레이어는 기존의 moe_gemv_kernel 경로를 그대로 사용합니다.

코드 변경 분석

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu

이 파일에는 Split-K2 최적화를 위한 새로운 커널과 헬퍼 함수가 추가되었습니다. 특히 partial_epilogue와 moe_gemv_splitk_partials_kernel이 핵심적인 역할을 합니다.

Before (개념적):

// 기존 단일 커널에서 모든 GEMV 및 SwiGLU 연산 처리
__global__ void moe_gemv_interleaved_swiglu_kernel(...) {
  // ... GEMV 계산 ...
  // ... SwiGLU epilogue 적용 ...
}

After (주요 추가 코드):

template <typename Details, int CtaM, int CtaN, int Threads, typename AccT>
__device__ __forceinline__ void partial_epilogue(float* partial_out, void* tile_acc) {
  static constexpr int Interleave = Details::kInterleave;
  static constexpr int ThreadsPerInterleavedTile = Details::kThreadsPerInterleavedTile;
  static constexpr int WarpSize = Details::kWarpSize;
  static constexpr int WarpNum = Threads / WarpSize;
  static_assert(CtaM == 1);
  static_assert(Threads % WarpSize == 0);

  __shared__ float shmem[CtaM * CtaN * Interleave * WarpNum];
  int tid = threadIdx.x;
  int warp_id = tid / WarpSize;
  int lane_id = tid % WarpSize;
#pragma unroll
  for (int n = 0; n < CtaN; ++n) {
    float v = static_cast<float>(reinterpret_cast<AccT*>(tile_acc)[n]);
    v = warp_reduce_sum<Interleave, ThreadsPerInterleavedTile>(v);
    if (lane_id < Interleave * ThreadsPerInterleavedTile && lane_id % ThreadsPerInterleavedTile == 0) {
      shmem[warp_id * CtaN * Interleave + n * Interleave + lane_id / ThreadsPerInterleavedTile] = v;
    }
  }
  __syncthreads();

#pragma unroll
  for (int col = tid; col < CtaN * Interleave; col += Threads) {
    float val = 0.f;
#pragma unroll
    for (int warp = 0; warp < WarpNum; ++warp) {
      val += shmem[warp * CtaN * Interleave + col];
    }
    partial_out[col] = val;
  }
}

template <typename Details, int CtaN, int Threads, int GroupSize, int SplitK,
          typename TypeA = typename Details::TypeDetailsA::Type, typename AccT = float>
__global__ void moe_gemv_splitk_partials_kernel(
    TypeA* act, uint8_t* weight, TypeA* scales, float* partials,
    int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts,
    int64_t weight_expert_stride, int64_t scale_expert_stride, int n, int k, int64_t expanded_num_rows) {
// ... (실제 부분 합 계산 로직) ...
}

partial_epilogue 함수는 각 워프(warp) 내에서 계산된 부분 합을 효율적으로 공유 메모리에 저장하고, 이를 다시 스레드 블록(thread block) 전체에서 감소시키는 역할을 합니다. moe_gemv_splitk_partials_kernel은 K 차원을 분할하여 각 부분의 GEMV 연산을 수행하고, 그 결과를 partials 버퍼에 저장합니다. 이후 moe_gemv_splitk_reduce_swiglu_kernel (diff에 직접 포함되지는 않았지만 PR 설명에 언급됨)이 이 partials를 최종적으로 처리합니다.

docs/contrib_ops/cuda/moe_qmoe.md 및 docs/contrib_ops/cuda/qmoe_gemv_experiments.md

이 문서들은 Split-K2 최적화가 어떻게 동작하는지, 그리고 어떤 환경 변수로 제어되는지 설명합니다. 특히 moe_qmoe.md에는 Split-K2 SwiGLU GEMV default path 섹션이 추가되어 fp16 INT4 interleaved-SwiGLU GEMV 경로에서 Split-K2가 기본으로 활성화됨을 명시합니다. 또한 ORT_DISABLE_MOE_GEMV_SPLITK2_SWIGLU=1 환경 변수를 통해 이 기능을 비활성화할 수 있음을 안내합니다.

--- a/docs/contrib_ops/cuda/moe_qmoe.md
+++ b/docs/contrib_ops/cuda/moe_qmoe.md
@@ -989,6 +989,27 @@ per-column INT4, block-wise INT4/INT8, and interleaved-SwiGLU GEMV kernels.
 | Kernel instantiation | `moe_gemv.cu` adds `__nv_bfloat16` details/instantiations (group sizes 0/32/64/128, INT4/INT8, bias on/off) under `ENABLE_BF16`. | The custom FC1/FC2 GEMV kernels run for BF16; no grouped-GEMM fallback when the FP16 gate would route. |
 | Profiling | GPT-OSS-20B, Qwen3.6-35B-A3B, and Gemma model shapes profiled with `block_size=64` for both dtypes. | BF16 matches FP16 routing and latency within noise (about 1.3x–1.5x faster than grouped GEMM); SwiGLU BF16 parity tests pass. |
 
+#### Split-K2 SwiGLU GEMV default path
+
+The fp16 INT4 interleaved-SwiGLU GEMV path uses a two-pass Split-K2 FC1 kernel by
+default for supported decode shapes. The first pass computes two K-split FP32
+partials into QMoE workspace, and the second pass reduces those partials, adds
+optional bias, and applies the interleaved SwiGLU epilogue. FC2 stays on the
+regular `moe_gemv_kernel` path.
+
+Set `ORT_DISABLE_MOE_GEMV_SPLITK2_SWIGLU=1` before process start to force the
+previous single-kernel FC1 SwiGLU GEMV path for debugging, A/B benchmarking, or
+bisecting numerical differences. On GPT-OSS-20B, Split-K2 reduced FC1 kernel
+work from about 21.42 us to 19.98 us and improved repeated CUDA-graph decode
+throughput by about 0.9% to 1.6% with valid focused-helper output. A 1000-sample
+MMLU smoke matched the opt-out fallback within noise. A future autotuner can
+replace this hand-selected default with per-shape route selection.
+
+```bash
+onnxruntime/test/python/transformers/profile_qmoe_gemv.py \
+  --case gpt_oss_20b_m1_top4_fp16_2880x2880_e32 \
+  --disable-splitk2-swiglu --warmup 5 --repeat 100 --nvtx
+```
 #### Accumulation policy
 
 The QMoE GEMV fast path accumulates fp16 activations in fp16 by default. Set

왜 이 최적화가 좋은가?

성능 향상

Split-K2 최적화는 LLM 추론 성능에 가시적인 개선을 가져왔습니다.

  • FC1 커널 작업 시간 감소: Nsight 프로파일링 결과, FC1 커널 작업 시간이 약 21.42 us에서 17.59 us + 2.39 us = 19.98 us로 감소했습니다. 이는 QMoE 호출당 약 1.44 us의 FC1 연산 시간 단축을 의미합니다.
  • 모델 수준 처리량 증가: GPT-OSS-20B 모델의 CUDA-graph 디코드 벤치마크에서 +0.9%에서 +1.6% (평균 +1.29%)의 처리량(throughput) 향상과 -1.28%의 지연 시간(latency) 감소를 보였습니다. 이는 대규모 모델의 실제 서비스 환경에서 상당한 이점으로 작용할 수 있습니다.

병렬성 극대화

Split-K2는 K 차원을 분할함으로써 GPU의 병렬 처리 장치(예: Streaming Multiprocessors)가 더 많은 작업을 동시에 처리할 수 있도록 합니다. 이는 특히 K 차원이 큰 행렬 곱셈에서 GPU 리소스 활용률을 높여 전체 연산 시간을 단축시키는 효과를 가져옵니다.

견고한 구현을 위한 리뷰 과정

PR 리뷰 과정에서 여러 중요한 개선 사항이 반영되어 최적화의 견고성과 정확성이 크게 향상되었습니다.

  • 환경 변수 일관성 유지: 초기 구현에서는 Split-K2 경로가 ORT_MOE_GEMV_FP16_ACCUM=1과 같은 기존 환경 변수의 의미를 무시할 수 있다는 피드백이 있었습니다. tianleiwu는 use_fp32_accum 변수를 먼저 계산하고, Split-K2 두 단계 경로(항상 FP32 partials를 사용)를 use_fp32_accum이 true일 때만 활성화하도록 수정했습니다. 이로써 ORT_MOE_GEMV_FP16_ACCUM=1이 요청되면 Split-K2는 비활성화되고 단일 커널 경로로 폴백되어 환경 변수의 동작이 문서화된 대로 유지됩니다.

  • 타입 안전성 확보: moe_gemv_splitk_partials 버퍼가 float*로 선언되었지만, 실제로는 FP16 또는 FP32 partials를 저장할 수 있어 혼란을 줄 수 있다는 지적이 있었습니다. tianleiwu는 Split-K2 partials를 항상 FP32로 처리하도록 강제하여 float* 타입이 버퍼 내용과 일치하도록 수정했습니다. 이는 잠재적인 타입 불일치 오류를 방지하고 코드의 명확성을 높입니다.

  • 테스트 안정성 개선: 테스트 코드에서 환경 변수를 런타임에 변경하는 방식이 함수 로컬 스태틱(function-local statics)에 캐싱된 값 때문에 불안정할 수 있다는 문제가 제기되었습니다. 이 문제는 Split-K2를 기본 경로로 설정하고 ORT_DISABLE_MOE_GEMV_SPLITK2_SWIGLU를 opt-out 방식으로 변경함으로써 해결되었습니다. 이제 테스트는 런타임에 환경 변수를 조작하지 않고, 프로세스 시작 전에 설정된 값을 기반으로 Split-K2 기본 경로를 안정적으로 검증합니다.

일반적인 교훈

이 PR은 다음과 같은 중요한 개발 및 최적화 교훈을 제공합니다.

  1. 정확한 프로파일링의 중요성: Nsight와 같은 도구를 사용한 상세한 커널 프로파일링이 실제 병목을 식별하고 최적화 효과를 정량화하는 데 필수적입니다.
  2. 환경 변수 및 설정의 일관성: 사용자에게 제공되는 환경 변수나 설정 옵션은 그 의미와 동작이 일관되고 예측 가능해야 합니다. 기존 시스템과의 호환성을 고려한 설계가 중요합니다.
  3. 테스트 코드의 견고성: 런타임 환경에 의존하는 테스트는 예상치 못한 동작을 일으킬 수 있으므로, 환경 변수 캐싱과 같은 미묘한 문제를 고려하여 견고하게 작성되어야 합니다.
  4. 성능과 정확도 간의 균형: 성능 향상과 함께 MMLU와 같은 벤치마크를 통해 모델의 정확도 회귀가 없는지 확인하는 과정은 필수적입니다.

마치며

이번 Split-K2 QMoE SwiGLU GEMV kernel 최적화는 ONNX Runtime이 LLM 추론 성능을 지속적으로 향상시키기 위한 노력을 보여주는 좋은 예시입니다. K 차원 분할을 통한 병렬성 극대화와 두 단계 커널 실행 전략은 QMoE FC1 레이어의 병목을 효과적으로 해소하여, GPT-OSS-20B와 같은 모델에서 의미 있는 처리량 향상을 달성했습니다.

이러한 최적화는 단순히 성능 수치를 개선하는 것을 넘어, 코드의 견고성과 유지보수성을 높이는 방향으로 리뷰 과정을 거쳐 완성되었습니다. 앞으로는 per-shape autotuning을 통해 각 모델 형태에 맞는 최적의 경로를 자동으로 선택하거나, launch-fused reduction과 같은 기법을 통해 추가적인 성능 개선을 모색할 수 있을 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글