CUDA 행렬 곱셈 최적화가 cuBLAS 성능의 93.7%에 도달했다

CUDA 행렬 곱셈 최적화 사례에서 단순 커널은 NVIDIA RTX A6000의 최고 성능 중 1.3%에 그쳤지만, 메모리 접근과 데이터 재사용을 단계적으로 개선한 커널은 cuBLAS 성능의 93.7%까지 도달했다. 이 설명은 GPU 커널을 설계하거나 딥러닝 연산의 성능 병목을 이해하려는 독자를 대상으로 한다.

측정은 CUDA 12.x를 사용하는 A6000(Ampere, GA102, 48GB GDDR6X)에서 이뤄졌다. 별도 조건이 없는 행렬은 4096×4096 단정밀도 부동소수점이며, 이 조건의 cuBLAS 기준 성능은 23,270 GFLOP/s다.

성능 차이를 만든 핵심은 연산량이 아니라 데이터 이동이다. 행렬은 행 우선으로 저장되므로 단순 커널에서 B의 열을 읽을 때 메모리 주소가 크게 떨어진다. 캐시 라인에 실린 값 가운데 필요한 값만 사용하게 되고, 여러 스레드가 같은 데이터를 반복해서 전역 메모리에서 읽는다.

타일링은 행렬을 작은 블록으로 나눠 빠른 메모리에 올린 뒤 여러 출력 계산에 재사용하는 방식이다. GPU에서는 공유 메모리에 타일을 올리고 스레드 블록이 함께 사용한다. 이 과정으로 전역 메모리 트래픽을 줄일 수 있지만, 공유 메모리만 사용하는 방식은 스레드당 계산량과 뱅크 충돌, 스레드 블록 점유율 때문에 성능이 제한된다.

이후 스레드마다 여러 출력 원소를 계산하도록 레지스터 타일링을 적용하고, 전역 메모리에서 float4 단위로 읽는다. 공유 메모리에서 A 타일을 전치해 계산 단계의 열 접근도 연속적인 읽기가 되도록 한다. 마지막으로 워프 안의 스레드와 데이터 배치를 조정해 레지스터 사용을 효율화한다. 이 순서로 성능이 309 GFLOP/s에서 21,779 GFLOP/s로 약 70배 높아졌다.

글은 남은 6.3% 차이가 cuBLAS의 행렬 크기·자료형·GPU 구조·행렬 비율별 특화 커널과 실행 시 자동 튜닝에서 비롯된다고 설명한다. 따라서 이 수치는 해당 A6000과 4096×4096 단정밀도 조건의 결과이며, 다른 하드웨어나 행렬 크기에 그대로 적용되는 보장은 아니다.

딥러닝 연산에 주는 의미

선형 계층은 행렬 곱셈을 직접 수행하고, 합성곱은 im2col 등의 방식으로 행렬 곱셈으로 바꿔 계산할 수 있다. 어텐션도 점수 계산과 출력 계산에 행렬 곱셈을 사용한다. 글에 따르면 PyTorch의 torch.mm()나 F.linear()는 cuBLAS로 연결되며, 합성곱은 cuDNN을 사용할 수 있다.

이 원리는 배치 크기와 시퀀스 길이, 혼합 정밀도, 여러 GPU에 행렬을 나누는 방식이 성능에 영향을 주는 이유를 이해하는 데 도움이 된다. 예를 들어 글은 작은 배치의 좁은 행렬이 GPU의 모든 SM을 채우지 못할 수 있고, FP16·BF16은 텐서 코어를 활용할 수 있다고 설명한다.

다만 이 글의 93.7% 결과는 A6000에서 표준 CUDA 코어를 이용한 커널에 관한 것이다. 텐서 코어와 TMA, 워프 그룹 MMA, 지속 커널을 사용하는 Hopper·Blackwell 세대의 기법은 별도로 다루며, H100의 FP16 행렬 곱셈 성능으로 764 TFLOP/s 이상을 제시한다. 이는 A6000의 측정값과 하드웨어·자료형 조건이 다르므로 직접적인 동급 비교로 읽어서는 안 된다.

관련 자료: CUDA 행렬 곱셈 최적화 설명과 시각화, GPU 메모리 계층, 스트리밍 멀티프로세서, 텐서 코어, HBM 메모리, NCCL 통신, 분산 병렬화.

추가 참고 자료: cuBLAS에 가까운 성능을 내는 CUDA 행렬 곱셈 커널 최적화, NVIDIA GPU의 고성능 행렬 곱셈 커널 분석, CUTLASS, Triton, CUDA C++ 모범 사례 가이드.