[onnxruntime] ONNX Runtime: GatedDeltaNet 및 GatedRMSNorm의 CUDA 커널 최적화 분석

PR 링크: microsoft/onnxruntime#33601 상태: Merged | 변경: +546 / -188

들어가며

최근 LLM 추론 성능을 극대화하기 위해 GatedDeltaNet과 같은 선형 어텐션(Linear Attention) 구조가 주목받고 있습니다. 하지만 기존 커널 구현은 범용적인 접근 방식을 사용하여 하드웨어 자원을 충분히 활용하지 못했습니다. 이번 PR은 NVIDIA SM80+(Ampere 이상) 아키텍처를 타겟으로 GatedDeltaNetChunkedKernel과 GatedRMSNorm 커널을 대대적으로 재설계하여, 메모리 대역폭 활용도를 높이고 불필요한 동기화 오버헤드를 제거했습니다.

코드 분석

1. GatedDeltaNetChunkedKernel: 메모리 I/O 및 동기화 최적화

가장 큰 변화는 uint4를 이용한 벡터화된 글로벌 메모리 I/O 도입입니다. 기존에는 스칼라 단위로 데이터를 로드/저장했으나, 16바이트 정렬(Alignment)이 보장될 경우 uint4를 사용하여 트랜잭션 효율을 극대화했습니다.

// Before: 스칼라 단위의 개별 요소 저장
pack.output[tok * v_stride + hv * DV + v0 + c] = FloatToElem<T>(p.scale * t_f[t * L::kVf + c]);

// After: 16바이트 정렬 확인 후 uint4 벡터화 저장
if (vec_out) {
  uint4 out4;
  out4.x = PackHalf2(...); // 8개 요소를 4개의 PackHalf2로 처리
  *reinterpret_cast<uint4*>(pack.output + ...) = out4;
}

또한, __syncthreads() 호출을 줄이기 위해 SyncBeforeStore 플래그를 도입하고, 공유 메모리 라운드 트립을 제거하여 커널 내 동기화 지점을 6개 이상 줄였습니다.

2. GatedRMSNormWarp128Kernel: 워프 셔플 기반의 병렬화

기존 GatedRMSNorm은 cub::BlockReduce를 사용하여 1개 행(row)당 CTA를 할당했으나, 이는 오버헤드가 컸습니다. 이를 4-warp/CTA 구조의 워프 셔플(Warp Shuffle) 방식으로 변경했습니다.

// Before: cub::BlockReduce 사용 (1-row/CTA)
// After: Warp Shuffle을 통한 레지스터 기반 연산
// REG:26 STACK:0 SHARED:0, 0 barriers

왜 이게 좋은가

이번 최적화는 하드웨어의 특성을 깊이 있게 활용했습니다.

1. 성능 수치: H200 GPU 기준 GatedDeltaNetChunkedKernel은 1.57배(34.74ms -> 22.18ms), GatedRMSNormWarp128Kernel은 2.94배(3.08ms -> 1.05ms)의 Latency 감소를 기록했습니다.
2. 정렬 가드(Alignment Guard): 리뷰 과정에서 지적된 IOBinding 시의 정렬 문제(2바이트 오프셋 등)를 해결하기 위해 IsAligned16 함수를 도입했습니다. 이는 벡터화된 I/O가 정렬되지 않은 메모리 주소에서 발생시킬 수 있는 Fault를 방지하는 필수적인 안전장치입니다.
3. 교훈: GPU 커널 최적화의 핵심은 '메모리 트랜잭션의 벡터화'와 '불필요한 공유 메모리 접근 및 동기화 제거'에 있습니다. 특히, 범용적인 커널을 작성할 때는 입력 데이터의 정렬 상태에 따른 스칼라/벡터 경로 분기 처리가 필수적임을 보여줍니다.

리뷰어 피드백 반영

리뷰어들은 특히 '정렬되지 않은 메모리 접근'에 대한 위험성을 강조했습니다. tianleiwu는 IOBinding을 통해 2바이트 오프셋으로 할당된 텐서가 기존 uint4 로직에서 Fault를 일으킬 수 있음을 지적했고, 이를 위해 런타임에 정렬을 확인하고 스칼라 경로로 폴백(fallback)하는 로직을 추가하여 안정성을 확보했습니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-shuffle-functions
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory-alignment

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