행렬 곱 최적화: 입력 재사용과 타일링

지난 글에서는 산술 강도를 통해 계산량과 데이터 이동량의 관계를 살펴봤습니다. 원소별 연산을 하나의 커널로 합치면, 중간값을 저장했다가 다시 읽는 과정을 줄일 수 있었습니다. 이번에는 행렬 곱의 계산 시간을 줄이는 두 가지 방법을 살펴보겠습니다.

첫째는 한 번 가져온 입력을 여러 계산에 재사용해 데이터 이동량을 줄이는 것입니다. 공유 메모리와 레지스터에 입력을 보관하면 같은 값을 반복해서 가져오는 일을 줄일 수 있습니다. 둘째는 현재 입력으로 계산하는 동안 다음 입력을 가져와 기다리는 시간을 줄이는 것입니다. 필요한 데이터 이동을 계산과 동시에 진행하면, 입력이 준비될 때까지 계산이 멈춰 있는 시간을 줄일 수 있습니다.

행렬 곱에서 입력을 재사용할 수 있는 이유

행렬 A와 B를 곱해 C를 만든다고 해보겠습니다. C의 출력 하나는 A의 한 행과 B의 한 열에서 같은 위치의 값을 곱하고, 그 곱들을 모두 더해 구합니다. A의 행 길이가 8이라면 출력 하나를 만들기 위해 여덟 쌍의 값을 곱해 더합니다.

그림 1에서는 A가 4 × 8, B가 8 × 4이므로 C는 4 × 4입니다. c00은 C의 0행 0열, c01은 0행 1열의 요소입니다. 여기서 행과 열의 번호는 0부터 시작합니다.

먼저 c00과 c01을 보겠습니다. c00은 A의 0행과 B의 0열로, c01은 A의 0행과 B의 1열로 계산합니다. B에서 사용하는 열은 다르지만 A의 0행은 두 출력에 공통으로 필요합니다.

c00과 c10도 비슷합니다. 이번에는 A의 행이 달라지지만, B의 0열은 두 출력에 공통으로 필요합니다. 같은 출력 행의 요소들은 A의 한 행을, 같은 출력 열의 요소들은 B의 한 열을 함께 사용하는 것입니다.

출력마다 필요한 입력을 따로 가져온다면 같은 값을 반복해서 읽게 됩니다. 반대로 여러 출력을 함께 계산하면서 공통 입력을 보관해 두면, 이미 가져온 값을 다시 사용할 수 있습니다. 여러 출력을 함께 계산하면, 공통으로 필요한 입력을 한 번 가져와 재사용할 수 있습니다.

공유 메모리를 활용한 입력 재사용

입력을 어디에 보관하면 좋을까요? 앞서 살펴본 HBM에는 큰 모델의 가중치와 입력을 담을 수 있습니다. 공유 메모리와 레지스터는 연산 장치에 더 가까이 있지만, 보관할 수 있는 양이 훨씬 적습니다. 이 차이를 B300의 수치로 살펴보겠습니다.

아래는 B300을 기준으로 이번 글에서 활용하는 세 메모리 공간을 비교한 표입니다. HBM은 NVIDIA의 HGX 구성 사양, 온칩 용량은 CUDA 메모리 사양표를 기준으로 했습니다.

공간	용량	대역폭
HBM3e	GPU 한 개당 288 GB	최대 8 TB/s
공유 메모리	SM당 최대 228 KiB	SM당 약 192 GB/s, 전체 약 29 TB/s (아래 가정)
레지스터	SM당 256 KiB	비교할 B300 공개 수치 미확인

KiB는 1,024 byte입니다. 레지스터 용량은 SM당 65,536개의 32비트 레지스터를 바이트로 환산한 값이며, 여러 스레드에 나누어 할당됩니다. 공유 메모리는 L1 캐시 등과 함께 쓰는 SM당 256 KiB 공간의 일부로, 표에는 공유 메모리로 설정할 수 있는 최댓값을 적었습니다. 따라서 공유 메모리와 L1 캐시의 용량을 별도로 더하면 안 됩니다.

공유 메모리의 대역폭은 B300의 공식 제품 수치 대신, 규모를 비교하기 위한 계산 예시를 적었습니다. CUDA 메모리 활용 가이드의 구조 설명을 바탕으로 하면 SM당 한 사이클에 128 byte를 처리할 수 있습니다. 클럭 1.5 GHz, SM 150개를 가정하면 SM당 약 192 GB/s, 전체 합산 약 29 TB/s입니다. B300의 HBM 대역폭인 8 TB/s와 비교하면 약 3.6배로, 공유 메모리에서 입력을 재사용할 때 데이터 공급 능력이 얼마나 달라질 수 있는지 가늠할 수 있습니다.

이 합산치는 각 SM이 자기 공유 메모리를 동시에 사용할 때의 값이며, 한 SM이 전체 대역폭을 사용할 수 있다는 뜻은 아닙니다. 실제 처리량은 클럭과 접근 방식에 따라 달라집니다. 레지스터는 같은 기준으로 비교할 B300의 공개 대역폭 수치를 확인하지 못했지만, 여기에 유지한 입력을 재사용하면 공유 메모리에서 같은 값을 다시 읽는 과정까지 줄일 수 있습니다.

이 수치에서 먼저 볼 것은 용량의 차이입니다. 예를 들어 값 하나가 4 byte인 FP32로 4,096 × 4,096 행렬을 저장하면 64 MiB가 필요합니다. 이 행렬은 B300의 HBM에는 들어가지만, SM 하나의 공유 메모리에 통째로 담을 수는 없습니다. 큰 데이터는 HBM에 두고, 지금 계산할 일부를 연산 장치 가까이 가져와 사용하는 이유입니다.

대역폭도 무한하지 않습니다. HBM과 실제로 주고받는 양이 8 GB라면, 최대 8 TB/s를 온전히 사용하더라도 그 이동에만 최소 1 ms가 필요합니다. 같은 값을 반복해서 가져오느라 이동량이 늘면 데이터 공급에 더 많은 시간이 필요합니다. 이미 가져온 입력을 여러 계산에 쓰면, 같은 계산량을 더 적은 이동으로 처리할 수 있습니다.

이제 계산할 출력을 먼저 골라보겠습니다. 타일(tile)은 행렬에서 함께 다룰 작은 직사각형 영역을 뜻합니다. 그림 2는 전체 출력이 8 × 8인 행렬에서, 한 블록(block)이 왼쪽 위 4 × 4 영역을 담당하는 경우를 보여줍니다.

선택한 출력 16개에 필요한 입력은 A의 첫 네 행과 B의 첫 네 열입니다. 예를 들어 A의 0행에 있는 값 하나는 선택한 출력의 0행에 있는 네 요소를 계산할 때 사용됩니다. B의 0열에 있는 값 하나도 선택한 출력의 0열에 있는 네 요소에 사용됩니다.

블록의 스레드들이 협력해 함께 필요한 입력을 전역 메모리에서 공유 메모리로 가져옵니다. 그러면 블록 안의 여러 계산이 각자 같은 입력을 전역 메모리에서 다시 읽는 대신 공유 메모리에 보관된 값을 사용할 수 있습니다. 입력이 준비된 뒤 사용하도록 순서를 맞추는 일도 필요합니다.

지난 글에서 본 것처럼 전역 메모리 접근에는 캐시가 관여합니다. 따라서 전역 메모리를 네 번 읽는다고 반드시 HBM도 네 번 읽는 것은 아닙니다. 여기서의 핵심은 함께 쓸 입력을 블록이 보관하고 재사용하도록 계산을 배치하는 것입니다. 큰 행렬에서는 선택한 행과 열의 입력도 일부씩 가져오며, 이 과정은 마지막 그림에서 살펴보겠습니다.

레지스터를 활용한 입력 재사용

공유 메모리에 가져온 입력도 계산할 때마다 반복해서 읽으면 이동이 생깁니다. 이번에는 블록이 맡은 출력 영역 안에서, 한 스레드(thread)가 2 × 2 영역의 출력 네 개를 계산한다고 해보겠습니다.

출력 하나는 여러 곱의 합이므로, 계산하는 동안에는 아직 완성되지 않은 값이 있습니다. 부분합은 지금까지 계산한 곱을 더한 값입니다. 이 예시에서는 네 출력의 부분합을 각각 0으로 시작하고, 스레드의 레지스터에 유지합니다.

그림 3은 각 출력의 첫 번째 곱을 계산하는 순간입니다. A[0, 0]과 A[1, 0]을 a0와 a1으로, B[0, 0]과 B[0, 1]을 b0와 b1으로 표시했습니다. A[0, 0]처럼 대괄호 안의 두 숫자는 행과 열의 위치를 뜻합니다.

a0는 b0와 곱해 c00에, b1과 곱해 c01에 더합니다. a1도 b0, b1과 각각 곱해 c10과 c11에 더합니다. 공유 메모리에서 읽어온 네 값으로 네 곱을 계산하고, 각 입력을 두 번씩 사용한 것입니다. 네 출력을 완전히 따로 계산하면서 매번 두 입력을 읽는 방식과 비교하면, 같은 입력을 다시 읽는 과정을 줄일 수 있습니다.

이어서 두 번째 곱에 필요한 A[0, 1], A[1, 1], B[1, 0], B[1, 1]을 읽고, 네 부분합에 새로운 곱을 더합니다. 이 과정을 여덟 번째 곱까지 이어가면 출력 네 개가 완성됩니다. 계산 도중의 부분합을 매번 밖에 저장하지 않고, 레지스터에서 이어서 갱신하는 것입니다.

공유 메모리에서는 블록 안의 여러 계산이 공통 입력을 사용했습니다. 레지스터에서는 한 스레드가 읽어온 입력을 여러 곱셈에 사용하고, 자신의 부분합을 유지합니다. 재사용은 한 단계에서 끝나지 않고, 계산을 더 작은 단위로 나누는 과정에서도 이어집니다. NVIDIA의 계층적 행렬 곱 설명

타일 크기와 자원 사용량

여러 출력을 함께 계산할수록 같은 입력을 더 많이 사용할 기회가 생깁니다. 그렇다면 블록이나 스레드가 맡는 출력 타일을 가능한 한 크게 잡으면 될까요?

타일을 크게 잡으면 계산 도중 보관해야 할 데이터도 많아집니다. 블록이 맡은 출력 영역이 커지면 함께 가져올 입력이 늘 수 있고, 스레드가 맡는 출력이 많아지면 유지할 부분합도 늘어납니다. 앞의 B300 표에서 봤듯이 공유 메모리와 레지스터는 모두 SM마다 한정된 공간입니다.

한 블록이 공유 메모리나 레지스터를 많이 사용하면, 같은 SM에 다른 블록을 함께 배치할 여지가 줄어들 수 있습니다. 앞서 배운 워프(warp) 스케줄링과 연결하면, 어떤 워프가 데이터를 기다릴 때 대신 실행할 수 있는 다른 워프가 줄어들 가능성도 있습니다.

반대로 작은 타일은 블록당 자원 사용량을 줄일 수 있지만, 각 블록이 공통 입력을 가져와 활용하는 범위도 작아집니다. 같은 전체 출력을 계산하기 위해 더 많은 블록으로 나누면, 블록 사이에서 같은 입력을 다시 가져오는 일이 늘 수도 있습니다.

따라서 입력 재사용을 늘리는 이점과, 동시에 실행할 작업을 충분히 확보하는 이점을 함께 봐야 합니다. 상주하는 워프가 많다는 사실만으로 빠른 것도, 타일이 크다는 사실만으로 빠른 것도 아닙니다. 실제 선택은 행렬의 모양과 자원 사용량, 실행 결과를 함께 보며 결정합니다. CUTLASS의 타일 크기와 자원 사용 설명

데이터 이동과 계산의 동시 진행

공유 메모리에 전체 입력을 담을 수 없다면, 필요한 입력을 조금씩 가져와 계산해야 합니다. 여기서 나누어 가져오는 것은 같은 출력을 완성하는 데 필요한 입력입니다. 새로운 출력으로 넘어가는 것이 아니라, 기존 부분합에 아직 계산하지 않은 곱을 더하는 과정입니다.

그림 5처럼 출력 하나에 여덟 곱을 더해야 한다고 해보겠습니다. 첫 두 곱을 더하는 일을 계산 1, 다음 두 곱을 더하는 일을 계산 2로 묶고, 같은 방식으로 계산 4까지 나눕니다. 입력 1은 계산 1에 필요한 A와 B의 값들이고 입력 2, 3, 4도 각각 대응하는 계산에 필요한 값들입니다. 블록은 자신이 맡은 모든 출력에 대해 이 네 단계를 진행합니다.

입력을 잠시 보관하려고 확보한 메모리 공간을 버퍼(buffer)라고 부릅니다. 별도의 하드웨어 장치를 뜻하는 것은 아닙니다. 이번 예시의 버퍼는 공유 메모리에 마련한 입력 보관 공간이며, 출력의 부분합은 앞서 설명한 레지스터에 계속 유지합니다.

먼저 버퍼가 하나라면 입력 1을 채운 뒤 계산 1을 수행합니다. 해당 입력의 사용이 끝나면 같은 버퍼를 입력 2로 채우고 계산 2를 수행합니다. 그림의 순차 진행에서는 입력을 가져오는 동안 계산이 기다리고, 계산하는 동안에는 다음 입력을 가져오지 않습니다.

버퍼가 두 개라면 현재 사용하는 입력을 유지하면서 다른 버퍼에 다음 입력을 가져올 수 있습니다. 버퍼 0의 입력 1로 계산 1을 수행하는 동안, 버퍼 1에는 입력 2를 준비합니다. 계산 1이 끝나고 입력 2가 준비되어 있으면 계산 2를 시작할 수 있습니다.

이제 입력 1을 다 사용했으므로 버퍼 0을 입력 3으로 채울 수 있습니다. 계산 2는 버퍼 1을 사용하고 있으니, 버퍼 0을 다시 채우는 일과 동시에 진행할 수 있습니다. 이후에는 두 버퍼의 역할을 번갈아 바꿉니다. 이처럼 현재 단계의 계산과 다음 단계의 데이터 준비를 동시에 진행하도록 구성하는 것을 파이프라이닝(pipelining)이라고 합니다.

이 과정에서도 지켜야 할 순서는 있습니다. 계산은 자신에게 필요한 입력이 준비된 뒤 시작해야 하고, 아직 계산에 사용하는 버퍼를 새로운 입력으로 덮어쓰면 안 됩니다. 필요한 동기화로 이 순서를 지키면서, 다른 버퍼에서 진행할 수 있는 데이터 이동을 먼저 시작하는 것입니다.

두 경우의 계산량과 데이터 이동량은 같습니다. 달라지는 것은 이동과 계산이 시간상 함께 진행되는지입니다. 버퍼를 두 개 만들었다는 사실만으로 동시 진행이 보장되는 것은 아니며, 실제 실행도 데이터 이동과 계산을 함께 진행하도록 구성해야 합니다. 입력 준비가 계산보다 오래 걸리면 기다림이 남을 수도 있습니다. 또한 버퍼를 추가하면 공유 메모리 사용량이 늘어나므로, 앞서 본 자원 사용의 절충도 함께 적용됩니다.

행렬 곱에서는 함께 계산할 출력을 고르고 공통 입력을 공유 메모리에 보관한 뒤, 레지스터로 읽어온 값도 여러 계산에 사용했습니다. 그리고 남은 데이터 이동은 현재 계산과 동시에 진행하도록 구성했습니다. 다음 글에서는 Attention으로 넘어가, 이러한 방식으로 입력을 나누고 계산을 이어갈 때 softmax의 합산 과정이 어떤 어려움을 만드는지 살펴보겠습니다.