[논문 리뷰] Stream-K: Work-centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU

Muhammad Osama, Duane Merrill외 3명|arXiv (Cornell University)|2023. 01. 09.

Parallel Computing and Optimization Techniques인용 수 5

한 줄 요약

Stream-K는 GPU에서 조밀한 행렬-행렬 곱셈(GEMM)을 위한 작업 중심의 병렬 분해 기법을 도입하여, 출력 행렬을 토막내는 대신 MAC 루프 반복을 스트리밍 다중처리기(SMs) 간에 균일하게 분배한다. 이는 다양한 문제 형태에서 거의 100%의 프로세서 활용도를 달성하며, cuBLAS 및 CUTLASS 대비 최대 14배의 피크 성능 향상과 평균 6.7배의 성능 향상을 이룬다. 이는 복잡한 커널 집합 대신 정밀도당 단일 커널을 사용함으로써 달성된다.

ABSTRACT

We introduce Stream-K, a work-centric parallelization of matrix multiplication (GEMM) and related computations in dense linear algebra. Whereas contemporary decompositions are primarily tile-based, our method operates by partitioning an even share of the aggregate inner loop iterations among physical processing elements. This provides a near-perfect utilization of computing resources, regardless of how efficiently the output tiling for any given problem quantizes across the underlying processing elements. On GPU processors, our Stream-K parallelization of GEMM produces a peak speedup of up to 14$ imes$ and 6.7$ imes$, and an average performance response that is both higher and more consistent across 32,824 GEMM problem geometries than state-of-the-art math libraries such as CUTLASS and cuBLAS. Furthermore, we achieve this performance from a single tile size configuration per floating-point precision, whereas today's math libraries employ complex kernel-selection heuristics to select from a large ensemble of kernel variants.

연구 동기 및 목표

• 넓은 GPU에서 타일 기반 GEMM 작업 부하의 정량화 비효율성으로 인한 성능 저하 문제를 해결하기 위해.
• 수학 라이브러리에서 복잡한 커널 선택 히وري스틱이 필요 없도록 하여 다양한 문제 형태에서 일관되고 높은 성능을 확보하기 위해.
• MAC 루프 수준에서 작업을 정량화함으로써 문제 크기나 형태에 관계없이 거의 100%의 프로세서 활용도를 달성하기 위해.
• 대규모 커널 집합을 정밀도당 단일이고 높은 이식성 있는 커널로 대체함으로써 라이브러리의 부피와 유지보수 복잡성을 줄이기 위해.

제안 방법

• Stream-K는 총 MAC 루프 반복 수를 GPU 스트리밍 다중처리기(SMs) 간에 균일하게 나누어 균형 잡힌 워크로드 분포를 보장한다.
• 작업 정량화의 원자 단위로 MAC 루프 반복을 사용하며, 이는 전체 출력 타일 대비 일정하고 최소한의 비용을 갖는다.
• 동적으로 p개의 분할 실선(O(p)개)를 생성하며, 여기서 p는 SM의 수이다. 이는 문제 크기 대신 프로세서 폭에 따라 워크로드 분포를 스케일링한다.
• 정밀도당 고정된 차단 요소를 사용하여, 복수의 커널 버전이나 복잡한 선택 히وري스틱이 필요 없도록 한다.
• 통신 및 동기화 오버헤드를 최소화하여, 문제 크기 대신 SM 수에 따라 증가하는 오버헤드를 줄인다.
• CUTLASS 2.11 내부에 구현되었으며, FP64, FP16→FP32 및 기타 정밀도 형식에서 일관된 성능을 제공한다.

(a) Data parallel decomposition with grid size $g$ =9 CTAs, large $128\times 128\times 128$ CTA work volumes, and 75% processor utilization ceiling

실험 결과

연구 질문

• RQ1 MAC 루프 수준에서 정량화하는 작업 중심의 분해 기법이 현대 GPU에서 타일 기반 GEMM보다 더 높고 일관된 성능을 달성할 수 있는가?
• RQ2 정밀도당 단일 커널이 cuBLAS와 같은 수학 라이브러리의 대규모 커널 집합을 성능 저하 없이 대체할 수 있는가?
• RQ3 다양한 문제 기하학적 형태에서 Stream-K의 워크로드 균형 조절은 타일 기반 방법 대비 정량화 효율성 측면에서 어떻게 비교되는가?
• RQ4 단일 통합 작업 중심 분해 기반의 통합 설계로 커널 선택 히وري스틱을 제거할 경우 성능에 어떤 영향을 미치는가?

주요 결과

• Stream-K는 32,824개의 GEMM 문제 기하학적 형태에서 cuBLAS 대비 최대 14배의 피크 성능 향상과 평균 6.7배의 성능 향상을 달성했다.
• FP64에서는 cuBLAS 대비 6% 높은 처리량을, FP16→FP32에서는 13% 높은 처리량을 달성했다.
• Stream-K의 성능 분포 범위는 cuBLAS와 이상화된 CUTLASS 오라클 모두보다 좁아, 더 뛰어난 일관성과 활용도를 보였다.
• 계산이 제한되는 영역(FP64에서 150ops/바이트 이상, FP16→FP32에서 400ops/바이트 이상)에서는 Stream-K가 항상 cuBLAS 커널 집합을 능가했다.
• Stream-K는 cuBLAS 대비 배포 크기를 최대 20배까지 줄였으며, 정밀도당 단일 커널을 사용해 수백 개의 전용 커널이 아닌 최고 성능을 달성했다.
• 문제 형태에 관계없이 거의 100%의 프로세서 활용도를 달성하여, 타일 기반 방법에서 발생하는 정량화 비효율성으로 인한 이론적 한계(75–90%)를 제거했다.

(b) Data parallel decomposition with grid size $g$ =18 CTAs, smaller $128\times 64\times 128$ CTA work volumes, and 90% processor utilization ceiling