CUDA 행렬곱 최적화 단계별 정리: 타일링부터 텐서코어와 CUTLASS 추상화까지

핵심 메시지

• 발표자는 소비자용 GPU 한 대로 나이브 행렬곱 커널부터 시작해 단계마다 최적화를 더하며 cuBLAS 대비 성능을 끌어올렸다.
• 가장 큰 초반 이득은 알고리즘이 아니라 메모리 접근 순서를 바꿔 병합 읽기를 만든 데서 나왔다.
• 스레드 하나가 더 많은 계산을 맡는 타일링을 1차원에서 2차원으로 확장하며 산술 강도가 올라가고 성능도 함께 올랐다.
• 데이터 타입만 반정밀도로 바꾸면 오히려 성능이 크게 떨어지고, 텐서코어 명령을 실제로 쓰도록 코드를 바꿔야 이득이 생긴다.
• 직접 짠 커널은 인덱싱 수학이 감당하기 어려울 만큼 복잡해지며, CUTLASS는 바로 그 부분을 추상화해 준다.

쉽게 이해하기

발표자는 파이토치 팀에서 일하며 분산 학습을 담당하고 있고, 행렬곱 커널 최적화는 개인 시간에 공부한 주제라고 밝혔다. 최신 데이터센터 GPU를 쓸 수 없어 소비자용 카드 한 대로 시작했고, 한 달 반 남짓 커널을 고쳐 쓰며 왜 CUTLASS 같은 라이브러리가 필요한지 몸으로 이해하게 됐다고 말했다. 강연 제목이 '어렵게 배우기'인 이유도 그 과정을 그대로 밟았기 때문이다.

출발점은 값을 하나씩 곱해 더하는 가장 비효율적인 커널이다. 여기서 첫 개선은 계산 방식이 아니라 반복 순서를 바꿔 메모리를 연속으로 읽게 만든 것이었고, 그것만으로 성능이 크게 뛰었다. 이어 전역 메모리 대신 공유 메모리에 블록을 올려 캐싱했고, 스레드 블록 크기를 줄여 스트리밍 프로세서가 놀지 않도록 점유율을 손봤다.

다음 단계는 스레드 하나가 맡는 일을 늘리는 타일링이다. B의 값 하나를 레지스터에 캐싱해 한 번에 네 개를 처리하는 1차원 타일링만으로도 cuBLAS 대비 성능이 몇 퍼센트대에서 20% 수준으로 올랐고, 두 방향 모두 캐싱하는 2차원 타일링으로 확장하자 30% 근처에 도달했다. 이 시점의 연산은 결국 슬라이스끼리의 외적을 누적하는 형태가 된다.

그다음은 데이터를 얼마나 빨리 밀어 넣느냐의 문제였다. 한 번에 128비트를 읽는 벡터화 로드를 쓰면 같은 데이터를 옮기는 데 필요한 명령 수가 4분의 1로 줄고, 프로파일러에서도 로드 명령의 종류가 바뀌는 것이 확인된다. 여기에 워프 단위 타일을 한 층 더 끼워 넣자 성능은 더 올랐지만, 대가로 중첩 반복문과 인덱싱 계산이 사람이 감당하기 힘든 수준으로 불어났다.

단정밀도로 절반을 조금 넘긴 지점에서 발표자는 텐서코어로 넘어간다. 흥미로운 대목은 데이터 타입만 반정밀도로 바꿨을 때다. 이론상 최대 성능은 두 배가 되지만 실제로는 원래의 4분의 1 수준으로 떨어졌다. 텐서코어 명령을 쓰도록 코드를 고쳐야 비로소 이득이 나오며, 반정밀도 입력을 곱하더라도 누적은 단정밀도로 하는 편이 수치적으로 안전하다는 점도 짚었다. 실시간 채팅에서는 원시 커널의 진짜 병목이 텐서코어 미사용이 아니라 가장 안쪽 반복문의 동적 인덱싱이라는 지적도 나왔다.

마지막 열쇠는 더블 버퍼링이다. 데이터를 싣고 계산하고 저장하는 일을 순서대로 하는 대신, 타일 메모리를 두 벌 두고 한쪽에서 계산하는 동안 다른 쪽으로 다음 데이터를 실어 나른다. 발표자는 이를 분산 시스템의 파이프라인 개념이 하드웨어 수준까지 내려온 것으로 설명했고, 버퍼를 여럿으로 늘리면 곧 파이프라인 단계 수 개념이 된다고 덧붙였다. CUTLASS에서는 이 단계 수는 물론 공유 메모리 뱅크 충돌을 피하는 스위즐링까지 인자 하나로 지정할 수 있고, 타일 모양은 정적 템플릿으로 미리 정해 두면 인덱싱 계산을 라이브러리가 대신 처리한다.

주요 인사이트

• 성능 최적화의 초반 이득이 알고리즘 변경이 아니라 메모리를 어떤 순서로 읽느냐에서 나온다는 점은, GPU 커널에서 데이터 공급이 얼마나 지배적인지 보여 준다.
• 반정밀도로 바꾸기만 하면 빨라질 것이라는 기대는 사실과 반대였다. 하드웨어 기능은 명시적으로 그 명령을 쓰도록 코드를 바꿔야 실제 성능이 된다.
• 직접 만든 커널의 진짜 한계는 산술 성능이 아니라 사람이 관리할 수 없는 인덱싱 복잡도였고, CUTLASS의 존재 이유가 바로 거기에 있다.
• 더블 버퍼링과 파이프라인 단계 수는 새로운 개념이 아니라, 큐를 정적으로 잡아 두는 저지연 시스템 설계 관행이 GPU 안으로 들어온 것에 가깝다.
• 프로파일러가 내놓는 요약 페이지의 최적화 제안을 매번 확인하는 반복 습관이, 다음에 무엇을 고쳐야 할지 알려 주는 가장 실용적인 나침반으로 소개됐다.

자주 묻는 질문

관련 AI 소식