본문으로 건너뛰기

[onnxruntime] [CUDA] HyperConnectionLowRankMix 최적화: Speculative Decoding과 Hopper 아키텍처를 위한 튜닝

PR 링크: microsoft/onnxruntime#33548 상태: Merged | 변경: +503 / -166

들어가며

LLM 추론 속도를 높이기 위한 기법 중 하나인 Speculative Decoding(추측 디코딩)은 한 번의 스텝에서 여러 개의 토큰($S$)을 동시에 검증합니다. 이 과정에서 $S \le 8$ 수준의 작은 시퀀스 길이에 대한 연산 효율성이 전체 처리량(Throughput)에 큰 영향을 미칩니다.

최근 Microsoft의 ONNX Runtime 레포지토리에 반영된 PR은 HyperConnectionLowRankMix 커널을 최적화하여, 특히 NVIDIA Hopper(SM90) 아키텍처에서의 성능을 극대화하고 Speculative Decoding 시나리오($S=8$)에서의 효율을 개선했습니다. 이번 글에서는 실제 코드 변경 사항을 통해 어떤 최적화 기법들이 적용되었는지 살펴보겠습니다.

코드 분석: 핵심 변경 사항

1. Fused Row-Group Dispatch (4x2 전략)

기존에는 행(row) 수에 관계없이 고정된 방식으로 커널을 런칭했으나, 이번 PR에서는 $S=8$인 경우를 위해 4x2 Fused Row-Group 전략을 도입했습니다. 이는 HcDownKernel에서는 CTA당 4개 행을 처리하고, HcUpPreKernel에서는 CTA당 2개 행을 처리하도록 분배하여 하드웨어 리소스를 최적으로 활용합니다.

Before:

// 기존에는 행 수에 따른 세밀한 분기가 부족했음
LaunchFused<T>(...);

After:

// docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md 발췌
// rows에 따른 최적의 row group 설정
| rows | row groups |
| ---- | ---------- |
| 1    | 1 x 1      |
| 2..7 | 2 x 2      |
| 8    | 4 x 2      |

이러한 정교한 디스패치는 특히 $S=8$인 Speculative Decoding 환경에서 H200 기준 약 0.04~0.08ms의 성능 이득을 가져왔습니다.

2. Hopper PDL (Programmatic Dependent Launch) 도입

NVIDIA Hopper 아키텍처(SM90)와 CUDA 12.3+ 환경에서는 커널 간의 의존성을 하드웨어 수준에서 더 효율적으로 관리할 수 있는 PDL(Programmatic Dependent Launch) 기능을 지원합니다. 이를 통해 HcDownKernel과 HcUpPreKernel 사이의 스트림 직렬화를 최적화했습니다.

Code Snippet (Environment Variable):

// onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu
#include "core/platform/env_var_utils.h"

// ORT_HYPER_CONNECTION_PDL 환경 변수를 통해 제어
// SM90+에서 cudaLaunchAttributeProgrammaticStreamSerialization 활성화

또한, 가중치(weight)가 상수(constant initializer)인 경우 __ldcg(load cache global)를 사용하여 캐시 효율을 높였습니다.

Before/After (Memory Access):

// Before
sv[r][i] = qi < nvec ? sp[qi] : make_uint4(0u, 0u, 0u, 0u);

// After (__ldcg 적용으로 캐시 히트율 향상)
sv[i] = qi < nvec ? __ldcg(sp + qi) : make_uint4(0u, 0u, 0u, 0u);

3. Register Spill 방지 및 Rank-320 특화

HcUpPreKernel에서 Rank가 320인 경우를 특화(Specialization)하여, 로컬 레지스터 배열(float4 pq[8][ROWS])을 사용하는 대신 부분합을 직접 누적하도록 수정했습니다. 이는 GPU의 한정된 자원인 레지스터 사용량을 줄여 Register Spill을 방지하고 점유율(Occupancy)을 높이는 효과가 있습니다.

Before (Register Array 사용):

float4 pq[8][ROWS]; // 레지스터 압박 유발 가능성
// ... 연산 후 누적

After (Direct Accumulation):

// HcUpPreKernel<T, ROWS, 4, 320> 특화 버전
// va[r]에 직접 K-split partial sums를 누적하여 레지스터 절약

4. Deterministic Unfused Path 및 Branch Tiling

기존의 Unfused Path는 cuBLAS GEMM을 사용했으나, 이번 업데이트에서는 결정론적(Deterministic) 결과를 보장하기 위해 직접 구현한 FusedInjectPostGateKernel로 교체되었습니다. 특히 리뷰어들의 피드백을 반영하여 8개 이상의 브랜치(num_branches > 8)도 처리할 수 있도록 grid.y 방향으로 타일링(Tiling)을 적용했습니다.

Reviewer Feedback 반영:

"num_branches가 8보다 큰 경우에도 대응할 수 있도록 타일링을 적용하고 테스트 케이스를 추가해달라."

After (Tiling Logic):

// FusedInjectPostGateKernel에서 8개 브랜치 단위로 타일링 수행
// grid.y를 통해 arbitrary num_branches 지원

또한, 불필요한 스크래치 버퍼 할당을 제거하여 메모리 오버헤드를 줄였습니다.

Before/After (Memory Management):

// onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc
// Before
auto inject = GetScratchBuffer<T>(rows * static_cast<size_t>(branches), GetComputeStream(context));

// After
// inject 스크래치 버퍼 할당 로직 삭제 (커널 퓨전으로 불필요해짐)

왜 이게 좋은 최적화인가?

  1. 아키텍처 특화 활용: 단순히 범용 코드를 작성하는 것에 그치지 않고, Hopper 아키텍처의 PDL이나 __ldcg 같은 최신 기능을 적극적으로 활용하여 하드웨어 잠재력을 끌어올렸습니다.
  2. 워크로드 맞춤형 설계: Speculative Decoding에서 자주 발생하는 $S=8$ 케이스를 타겟팅하여 4x2 Row-Group 분배를 설계한 점은 실전 성능에 큰 영향을 미칩니다. 벤치마크 결과, 기본 설정 대비 약 1% 이상의 성능 향상을 확인했습니다.
  3. 안정성과 확장성: 리뷰어의 지적을 수용하여 브랜치 수 제한을 풀고 타일링을 도입함으로써, 최적화와 범용성 사이의 균형을 맞췄습니다. 또한 Deterministic 연산을 보장하여 모델 재현성 문제를 해결했습니다.

일반적인 교훈

  • Small Batch/Sequence 최적화: LLM 서비스 환경에서는 큰 배치보다 작은 배치나 짧은 시퀀스의 지연 시간(Latency)이 중요할 때가 많습니다. 이때는 커널 런칭 오버헤드와 리소스 분배가 핵심입니다.
  • Register Pressure 관리: CUDA 커널에서 배열을 로컬 변수로 선언할 때는 항상 Register Spill 가능성을 염두에 두어야 합니다. 특수 케이스에 대한 Specialization은 이를 해결하는 좋은 방법입니다.

이 PR은 최신 GPU 아키텍처의 기능을 이해하고, 실제 서비스 시나리오(Speculative Decoding)의 특성을 반영한 수준 높은 최적화 사례라고 할 수 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글