CuTe 레이아웃 대수 해설: GPU 커널의 데이터 배치를 범주론으로 다시 읽는 법

핵심 메시지

• CuTe에서 텐서는 반복자와 레이아웃으로 표현되며, 레이아웃은 모양(shape)과 스트라이드(stride)라는 두 개의 중첩 튜플만으로 데이터 배치를 기술한다.
• 레이아웃은 메모리 배치뿐 아니라 스레드·CTA 사이의 데이터 분배까지 담당해, MMA 계열 하드웨어 명령의 누산기 소유 패턴을 그대로 표현한다.
• 합성(composition)·여집합(complement)·병합(coalesce)이 논리적 분할과 논리적 곱의 밑바탕이며, 이 가운데 여집합이 가장 까다로운 연산이다.
• 연사들은 '다루기 쉬운(tractable) 레이아웃'만 추려 튜플 사상이라는 범주론적 대상과 일대일로 대응시켰고, 합성이 사상의 합성과 정확히 일치함을 보였다.
• 여집합은 사상이 '건드리지 않은 항목'을 그대로 채워 넣는 것으로 해석되며, 합성 알고리즘은 상호 세분화·당김·밂 구성으로 기술된다.

쉽게 이해하기

GPU에서 행렬 곱이나 어텐션 커널을 짤 때 가장 먼저 부딪히는 문제는 연산 자체가 아니라 데이터가 메모리 어디에 어떤 간격으로 놓여 있는가다. 엔비디아의 CuTe 라이브러리는 이 문제를 레이아웃이라는 하나의 개념으로 통일한다. 레이아웃은 모양과 스트라이드라는 두 튜플로 이루어지며, 좌표를 받아 반복자에서 얼마나 떨어진 위치를 읽을지 알려주는 함수 역할을 한다.

강연은 먼저 익숙한 예부터 훑는다. 행 우선과 열 우선 레이아웃은 어느 방향이 메모리에서 연속인지를 스트라이드 숫자로 드러낸다. 여기서 조금만 나아가면 같은 값을 여러 칸이 공유하는 브로드캐스트 레이아웃, 큰 텐서를 잘라낸 흔적이 남은 스트라이드 레이아웃, 작은 타일 단위로만 연속인 인터리브 레이아웃이 나온다. 블록 단위 스케일 값을 중복 저장하지 않고 브로드캐스트로 표현하는 사례처럼, 이 변형들은 실제 커널 최적화에서 곧바로 쓰인다.

레이아웃의 두 번째 쓰임은 데이터를 '누가 갖고 있는가'를 적는 일이다. 발표자들은 특정 MMA 명령의 누산기가 워프 32개 스레드에 어떻게 쪼개지는지, 그리고 워프 그룹 128개 스레드가 64×64 누산기 타일을 어떻게 나눠 갖는지를 스레드-값 레이아웃 하나로 적어 보인다. 하드웨어 명령을 쓰려면 이 소유 패턴을 정확히 알아야 하는데, 레이아웃이 그 고된 부기 작업을 대신해 준다.

이야기의 중심은 연산이다. 두 레이아웃을 이어 붙이는 합성, 어떤 레이아웃이 '덮지 않은 나머지'를 만들어 주는 여집합, 레이아웃 함수는 그대로 두고 표현만 간단히 줄이는 병합이 있다. 타일 하나를 지정하는 레이아웃과 타일 사이를 건너뛰는 격자의 레이아웃이 서로 여집합 관계라는 설명은 왜 이 연산이 논리적 분할의 핵심인지 잘 보여 준다. 다만 합성은 항상 되는 것이 아니어서, 두 레이아웃을 다시 '모양과 스트라이드' 형태로 인수분해할 수 있어야 한다.

후반부가 이 강연의 독자적인 기여다. 연구팀은 열 우선 레이아웃에서 모드를 치환·삭제·삽입·재괄호화해 얻을 수 있는 레이아웃만을 '다루기 쉬운 레이아웃'으로 정의하고, 이들을 튜플 사상이라는 범주론적 대상과 대응시켰다. 이 사전을 쓰면 레이아웃 합성은 사상의 합성이 되고, 여집합은 '사상이 도달하지 않은 항목을 채우기'라는 눈으로 보이는 조작이 된다. 합성 가능 여부까지 포함한 명시적 알고리즘도 제시되며, 오퍼라드나 배타적 논리합 기반 스트라이드로의 확장이 향후 과제로 남았다.

주요 인사이트

• 레이아웃을 '메모리 주소 계산 규칙'이 아니라 좌표 함수의 인수분해로 보면, 합성이 왜 항상 가능하지 않은지가 곧바로 설명된다. 두 함수를 이어 붙일 수는 있어도 그 결과를 다시 모양과 스트라이드로 쪼갤 수 있다는 보장이 없기 때문이다.
• 평평한 레이아웃 둘을 합성했는데 결과에 중첩 구조가 생기는 흔한 현상은, 사상의 정의역과 공역을 맞추려고 모드를 잘게 쪼개는 과정에서 필연적으로 나온다는 것이 이 이론의 설명이다.
• 이론을 다루기 쉬운 레이아웃으로 좁힌 것은 표현력을 포기한 타협이 아니다. 실무에서 행 우선으로 들어온 텐서를 나누고 타일링해 스레드에 분배하는 과정은 대체로 이 범위 안에서 닫히기 때문이다.
• 발표자들은 유리수 항목을 허용하면 임의의 레이아웃까지 같은 형식으로 적을 수 있다고 덧붙였다. 깔끔한 정수 이론을 먼저 세우고 필요할 때 확장하는 전형적인 수학적 전략이다.
• 여집합의 정의가 연구자마다 조금씩 다르다는 논의는 이 분야가 아직 정리 중임을 보여 준다. 압축성 조건이 필요한 곳은 분할 연산에서 원소를 잃지 않아야 할 때로 한정된다는 지적이 나왔다.

자주 묻는 질문

관련 AI 소식