High-Performance FP16 General Matrix Multiplication on NVIDIA Ada Lovelace via CUTLASS and WMMA: A Benchmark and Pipeline-Depth Analysis

기술 요약: CUTLASS 및 WMMA를 이용한 NVIDIA Ada Lovelace 기반 고성 성능 FP16 GEMM

문제 정의
일반 행렬 곱셈(GEMM)은 딥러닝과 과학 계산의 계산적 기초입니다. NVIDIA의 Ada Lovelace 아키텍처(특히 AD107 칩을 탑재한 소비자용 RTX 4060)는 FP16 텐서 코어 이론적 피크 처리량인 60.55 TFLOPS를 제공하지만, 피크 성능에 근접한 성능을 달성하려면 GPU 메모리 계층 구조의 정교한 오케스트레이션이 필요합니다. 과제는 텐서 코어를 포화시키기 위해 타일링 전략, 소프트웨어 파이프라인 깊이, 그리고 스레드 블록 점유율(occupancy) 사이의 균형을 맞추는 것입니다. 이때 공유 메모리 압박이나 레지스터 제약에 의해 병목 현상이 발생하지 않도록 해야 합니다. 기존 문헌들은 라이브러리 프로파일러 수준에서 GEMM을 광범위하게 특성화해 왔으나, 소비자용 하드웨어에서 CUTLASS 디바이스 레벨 API를 사용하여 임의의 파이프라인 깊이와 타일 기하학(tile geometries)의 영향을 체계적으로 평가하는 데에는 공백이 존재합니다.

방법론
저자들은 CUDA 13.1 및 Windows 11 환경에서 NVIDIA RTX 4060 (Ada Lovelace)을 대상으로 수행한 FP16→FP32 GEMM 연산의 체계적인 벤치마크를 수행했습니다. 연구에서는 세 가지 구현 계층을 비교했습니다:

1. cuBLAS: 참조 모델으로 사용되는 벤더 최적화 폐쇄형 소스 라이브러리.
2. CUTLASS 디바이스 레벨 API: 타일 크기, 파이프라인 깊이($S$), 에필로그 벡터화의 컴파일 타임 커스터마이징을 허용하는 오픈 소스 C++ 템플릿 라이브러리.
3. 커스텀 WMMA 커널: 교육적 검증을 위해 CUDA Warp Matrix Multiply Accumulate (WMMA) API를 사용하여 직접 튜닝한 커널.

실험은 고정된 문제 크기 $N = 8192$ ($M=N=K=8192$)에 집중되었습니다. 방법론은 다음과 같습니다:

• 파이프라인 깊이 변화: $256 \times 128$ 및 $128 \times 128$ 스레드 블록 타일을 사용하여 파이프라인 깊이 $S=3$ 및 $S=4$를 갖는 CUTLASS 구성을 테스트했습니다.
• 분석적 모델링: 파이프라인 효율성($S/(S+\kappa)$)을 기반으로 성능 모델을 개발하고, 이를 확장하여 처리량을 예측하기 위한 점유율 제약($\gamma(S)$)을 포함했습니다.
• 검증: 상대 오차 임계값 $10^{-2}$ 내에서 cuBLAS에 대한 커스텀 WMMA 커널의 수치적 정확성을 보장했습니다.
• 측정 프로토콜: 클럭 스케일링 및 초기화 오버헤드를 완화하기 위해 워밍업 단계를 포함한 호스트 측 타이밍을 사용했으며, cuBLAS와 CUTLASS 실행 스트림을 분리해야 했던 특정 Windows 드라이버 안정성 문제(TDR)를 기록했습니다.

주요 결과

• 최적의 구성: 파이프라인 깊이 $S=3$, 스레드 블록 타일 $256 \times 128$, 에필로그 벡터화 4를 갖춘 CUTLASS 구성이 31.9 TFLOPS(이론적 피크의 52.7%)를 달축했습니다. 이는 CUTLASS 프로파일러의 내장 s116816gemm_f16 커널 성능과 일치하며, 벤치마크 오버헤드로 인한 차이를 감안할 때 cuBLAS 참조값(31.5 TFLOPS)을 1.3% 상회했습니다.
• 파이프라인 깊이 이상 현상: 직관과는 반대로, 파이프라인 깊이를 $S=4$로 증가시켰을 때 성능 저하가 발생하여 31.1 TFLOPS(2.5% 감소)를 기록했습니다.
• 병목 분석: $S=4$에서의 성능 저하는 점유율 붕괴(occupancy collapse) 때문인 것으로 분석되었습니다. $S=3$(72 KB 공유 메모리 필요)은 SM당 2개의 스레드 블록을 허용했지만, $S=4$(96 KB 필요)는 점유율을 SM당 1개 블록으로 강제했습니다. 이는 동시 실행되는 워프 수를 절반으로 줄여 메모리 레이턴시를 숨기는 스케줄러의 능력을 저하시켰습니다. 또한, 스레드당 레지스터 압박이 약 64개에서 72개로 증가했습니다.
• 모델 정확도: 단순 분석 모델은 $S=2$와 $S=3$에 대해서는 성능을 정확히 예측했으나, $S=4$에서는 실제 값보다 13.8% 높게 예측했습니다. 점유율 상한을 포함한 확장 모델은 $S=4$ 예측치를 측정값의 1% 이내로 보정했습니다.
• 커스텀 커널 성능: 커스텀 WMMA 커널은 $N=4096$에서 25.1 TFLOPS(피크의 41.5%)를 달성했습니다. 이 낮은 성능은 소프트웨어 파이프라이닝(단일 버퍼링)의 부재와 최적화된 CUTLASS 템플릿에 비해 높은 동기화 오버헤드 때문인 것으로 나타났습니다.

주요 기여

1. 재현 가능한 벤치마킹 프레임워크: Ada Lovelace GPU에서 임의의 컴파일 타임 파라미터를 사용하여 GEMM 커널을 인스턴스화할 수 있는 CUTLASS 디바이스 레벨 API 기반의 오픈 소스 프레임워크.
2. 파이프라인 깊이에 대한 경험적 증거: 더 깊은 파이프라인($S=4$)이 공유 메모리 및 점유율 제약으로 인해 항상 더 높은 처리량을 보장하지는 않는다는 실험적 증거를 제시하며, 이는 단순한 처리량 모델과 상충됩니다.
3. 보정된 분석 모델: $S=2$ 및 $S=3$의 처리량을 성공적으로 예측하지만, $S=4$ 영역을 정확히 모델링하기 위해서는 점유율 캡(occupancy-capped) 확장이 필요한 성능 모델.
4. 교육적 참조 자료: 저수준 텐서 코어 프로그래밍을 위한 참조 역할을 하는 검증된 커스텀 WMMA 커널.
5. 드라이버 동작 관찰: CUTLASS와 cuBLAS 커널을 교차 실행할 때 발생하는 특정 Windows 드라이버 불안정성(TDR 관련)을 기록하여, Windows 환경에서의 교차 라이브러리 벤치마킹 시의 실질적인 제약 사항을 명시했습니다.

의의 및 주장
본 논문은 현대 소비자용 GPU에서 고성능 GEMM을 달성하는 것은 파이프라인 깊이가 유일한 결정 요인이 아닌 다중 병목 현상 문제임을 주장합니다. 본 연구는 공유 메모리 압박과 점유율 제약이 더 깊은 소프트웨어 파이프라인의 기대 이익을 역전시킬 수 있음을 강조합니다. 소비자급 RTX 4060이 이론적 텐서 코어 피크의 약 53%(31.9 TFLOPS)를 달성할 수 있음을 입증함으로써, 본 연구는 특정 워크로드에 대해 저렴한 하드웨어에서도 데이터센터급 성능 수준에 접근 가능하다는 것을 시사합니다. 저자들은 자신들의 오픈 소스 도구와 연구 결과를 GPU 최적화 기법을 연구하기 위한 자원으로 제시하며, 접근 가능한 고성능 컴퓨팅 및 에너지 효율적인 머신러닝 추론이라는 더 넓은 목표에 기여하고자 합니다. 본 연구는 모든 GEMM 최적화 문제를 해결한다고 주장하는 것이 아니라, Ada Lovelace 아키텍처에서 파이프라인 깊이와 점유율 간의 트레이드오프에 대한 구체적이고 재현 가능한 분석을 제공합니다.