CuTe의 비밀 풀기: 처음부터 빠른 GEMM 작성하기

URL: https://schreibereric.substack.com/p/demystifying-cute-how-to-write-a

Title: Demystifying CuTe: How to Write a Fast GEMM from Scratch

2부 중 2부 — 복사 커널에서 빠른 행렬 곱셈까지

시리즈 개요: 이 글은 NVIDIA의 CuTe(CUDA Template) 라이브러리에 대한 2부작 시리즈의 두 번째 글이다. 1부에서는 레이아웃 대수, CuTe API, 그리고 효율적인 복사 연산의 기초를 다뤘다. 이번 글에서는 그 개념들을 적용해 cuBLAS를 능가하는 GEMM 커널을 만든다.

TL;DR: CuTe의 레이아웃 대수를 사용해 행렬 곱 커널을 만들고, (베이스라인) cuBLAS 성능의 102%에서 시작해 더블 버퍼링(111%), L2 스위즐링(112%), 3-스테이지 파이프라이닝(116%)으로 발전시키며, 이 모든 것을 깔끔하고 선언적인 코드로 구현한다. 그리고 더 빠른 반복을 위해 Python으로 포팅한다.

그림 1: CuTe에서 공유 메모리에서 레지스터로 복사하는 예시.

아직 1부를 읽지 않았다면, 다음 내용만은 알고 있으면 된다:

레이아웃은 함수다. CuTe 레이아웃은 (Shape, Stride) 튜플이며, 논리 좌표를 물리 메모리 오프셋으로 매핑한다. 예를 들어 (4,3):(3,1)은 4×3 행우선(row-major) 행렬을 나타낸다.

⊘로 타일링. 논리적 나눗셈 연산자는 텐서를 계층적으로 분할한다: (A,B,C)⊘(a,b,c)=((a,b,c),(A/a,B/b,C/c))

∘로 합성. 레이아웃 합성은 변환을 연결하며, 예를 들어 기존 복사 메커니즘에 스위즐링을 추가할 때 유용하다.

Copy/MMA Atom. 하드웨어 특화 프리미티브(TiledCopy, TiledMMA)는 스레드들이 데이터를 어떻게 협력적으로 이동/계산하는지를 기술하며 PTX 세부사항을 추상화한다.

이 도구들로 우리는 128×64 타일에서 1.48 TB/s를 달성하는 복사 커널을 만들었다. 이제 진짜 도전 과제인 행렬 곱셈을 해보자.

1부에서 레이아웃 대수와 복사 커널을 확립했으니, 자연스럽게 떠오르는 질문은 “CuTe의 추상화가 실제로 커널을 더 쉽게 작성하게 해주나?”이다. CUTLASS 저장소의 cutlass/examples/cute 아래에 GEMM 예제들이 있지만, 이들은 최적화된 구현으로 바로 뛰어들어 각 단계의 이유를 설명하지 않는다. 이번 글은 1부와 같은 단계별 접근으로 최소 예제를 점진적으로 확장해, 내가 학습한 과정을 따라갈 수 있게 한다.

모든 예제는 이 repo에서 전체 코드를 확인할 수 있다. 목표는 피크 성능에 도달하는 것이 아니라, CuTe의 레이아웃 대수가 전체 GEMM 커널로도 깔끔하게 확장되는 모습을 보여주는 것이다. 몇 가지 제약(블록 크기 최적화 없음; 타일 형태 128x128x64; 정렬이 좋고 충분히 큰 행렬)에 기반해 가능한 한 빠르게 커널을 만들고자 한다.

최신 서버급 GPU(A100+)에서 단순 GEMM은 개념적으로 (1) 텐서 데이터 로드, (2) 텐서 코어 연산, (3) 결과 기록의 세 단계로 생각할 수 있다. 따라서 이전 글의 코드를 GEMM 예제의 기반으로 사용할 수 있다. 다시 여러 GEMM 예제를 살펴볼 것이다. 보여주겠지만, CuTe 레이아웃을 사용한 베이스 구현이 있으면 더블 버퍼링이나 L2 그리드 스위즐링 같은 개념을 추가하는 데 큰 변경이 필요 없다. 비동기 복사를 사용하는 가장 단순한 GEMM 구현부터 시작하자.

먼저 동작 방식, 즉 문제 타일링을 정의하자. 완벽한 스위즐링을 위해서는 각 스레드가 새로운 뱅크를 맞춰야 한다. 각 뱅크는 4바이트를 보유한다. 따라서 K-major 형식에서 완벽한 스위즐링을 위해서는 K≥32∗4/2=64개의 BF16 값이 필요하다. M과 N은 재사용을 담당하므로 가능한 한 크게 잡아 I/O 복잡도를 낮추고 GEMM 속도를 올린다. 따라서 타일 형태를 128x128x64로 시작한다. 이 타일들은 출력 C 행렬에 해당한다.

각 CTA(스레드 블록)는 이런 타일 하나를 계산하며 K 차원을 따라 반복한다. 스위즐링을 손쉽게 처리하기 위해 atom을 정의하고 이를 공유 메모리 레이아웃에 합성한다:

cpp
// 스위즐 atom 정의
auto swizzled_128B_atom = composition(
    Swizzle<3,3,3>{},
    make_layout(
        make_shape(Int<8>{}, make_shape(Int<8>{}, Int<8>{})),
        make_stride(Int<8>{}, make_stride(Int<1>{}, Int<64>{})))
);

// 공유 메모리 A와 B(128x64)에 적용
auto sA = tile_to_shape(swizzled_128B_atom, make_shape(Int<BM>{}, Int<BK>{}));
auto sB = tile_to_shape(swizzled_128B_atom, make_shape(Int<BN>{}, Int<BK>{}));

그림 2: A와 B 타일을 위한 스위즐된 공유 메모리 레이아웃.

커널 내부에서 비동기 복사 명령을 사용하더라도(이전 글에서 더 빠르다는 것을 확인) 실행 흐름은 동기적이다. 로드를 트리거하고, 완료를 기다린 뒤, 계산하고, 반복한다. 이는 로드-계산 “정지-출발(stop-and-go)” 패턴을 만든다.

cpp
// 메인 MMA 루프
for (int k = 0; k < (K + 64 - 1) / 64; k += 1) {
    // 1. 글로벌에서 공유 메모리로 비동기 복사 발행
    copy(copyA, tAgA(_,_,_,k), tAsA);
    copy(copyB, tBgB(_,_,_,k), tBsB);
    
    // 2. 데이터가 완전히 도착할 때까지 대기
    cp_async_fence();
    cp_async_wait<0>();
    __syncthreads();

    // 3. 공유 메모리에서 레지스터로 로드
    copy(copyS2R_A, tXsA, tXrA);
    copy(copyS2R_B, tXsB, tXrB);

    // 4. 계산
    gemm(mma, tCrA, tCrB, tCrC);

    __syncthreads();
}

두 하드웨어 프리미티브가 계산 단계와 데이터 이동 단계를 구동한다. SM80_16x8x16_F32BF16BF16F32_TN atom은 mma.sync.m16n8k16 텐서 코어 명령에 매핑된다: BF16 입력과 F32 누산을 갖는 16×8×16 행렬 곱이다. 이를 2×2로 타일링하면 32×16×16 워프 수준 연산이 된다. 공유 메모리 → 레지스터 전송에는 SM75_U32x4_LDSM_N이 ldmatrix.x4 명령을 내보낸다: 워프의 32개 스레드가 협력하며 각자 포인터 하나를 제공하고, 하드웨어가 8×8 타일 4개를 텐서 코어가 기대하는 레지스터 레이아웃으로 직접 로드한다. make_tiled_copy_A/B는 MMA 레이아웃으로부터 올바른 스레드별 포인터 매핑을 자동으로 유도하므로 수동 주소 연산이 필요 없다.

cpp
TiledMMA mma = make_tiled_mma(SM80_16x8x16_F32BF16BF16F32_TN{},
                                 Layout<Shape<_2,_2>>{},    // 2x2x1 MMA Atoms
                                 Tile<_32,_32,_16>{});      // LDSM용 32x32x16 Tiled MMA

Copy_Atom<SM75_U32x4_LDSM_N, bf16> s2r_atom;
auto copyS2R_A = make_tiled_copy_A(s2r_atom, mma);
auto copyS2R_B = make_tiled_copy_B(s2r_atom, mma);

그림 3: 우리가 만든 모든 GEMM 커널의 모든 계산 연산에 사용한 MMA atom.

주목할 만한 우회로 하나: 결과 기록(write-back)을 공유 메모리를 통해 라우팅해보았다. F32 누산값을 BF16으로 변환해 공유 메모리에 쓰고, 동기화한 뒤, 128비트 tiled copy로 더 대역폭 효율적인 글로벌 스토어를 하자는 아이디어였다. 비싼 스토어를 벡터화하면 공유 메모리 경유 비용을 상쇄할 수 있지 않을까? 결과는 더 느렸다. 문제는 공유 메모리를 거쳐야 하는 레이아웃 불일치다. MMA 누산기는 m16n8k16 명령이 만들어내는 패턴대로 스레드들에 흩어져 있고, 연속적인 행우선 덩어리로 있지 않다. 추가 __syncthreads()와 왕복 비용이 더 넓은 스토어로 절약되는 것보다 컸다. 각 스레드가 자신의 누산기 조각을 곧바로 글로벌에 쓰는 direct copy(tCrC, tCgC)가 더 단순하고 더 싸다.

면책 조항: 다른 블로그 글들처럼 우리는 커널에 최적인 행렬 크기를 본다. 보다 미묘한 그림은 그림 5에 있다.

2048x2048 크기 행렬에서 이 커널을 프로파일링하면 이미 cuBLAS 속도의 최대 101%까지 도달한다. 다만 cuBLAS에 최대한 가깝게 보이도록 이 크기를 일부러 선택했음을 언급해야 한다. 그래도 첫 시도치고는 나쁘지 않다. 이 구현은 스위즐된 레이아웃으로 뱅크 충돌을 효과적으로 해결하지만, 메모리를 기다리는 동안 계산 유닛(텐서 코어)이 놀고, 텐서 코어가 계산하는 동안 메모리 버스가 노는 문제가 있다. 이를 고치려면 두 작업을 겹쳐야 한다.

알다시피, 로드-대기-계산 패러다임은 현대 GPU에 비최적이다. 대신 현재 타일을 계산하는 동안 다음 타일을 로드하고 싶다. 이를 위해 공유 메모리를 두 배로 할당해야 한다. 한 버퍼(Stage 1)에 쓰는 동안 다른 버퍼(Stage 0)에서 읽을 수 있게 말이다.

CuTe에서는 놀라울 정도로 우아하게 구현된다. 공유 메모리 레이아웃에 파이프라인 스테이지 차원을 추가하기만 하면 된다. (그리고 다른 차원의 스위즐링 패턴을 망치지 않도록 주의해야 한다.)

cpp
// 이전: make_shape(Int<BM>{}, Int<BK>{})
// 신규: 스테이지를 위한 3번째 차원 추가
auto sA = tile_to_shape(swizzled_128B_atom, make_shape(Int<BM>{}, Int<BK>{}, Int<2>{}));
auto sB = tile_to_shape(swizzled_128B_atom, make_shape(Int<BN>{}, Int<BK>{}, Int<2>{}));

하지만 레지스터 할당 시 작은 주의점이 생긴다. 이제 sA는 Rank-3(M, K, Stage)이므로, 순진하게 partition하면 Rank-3 레지스터 조각이 만들어진다. 레지스터에는 스테이지가 필요 없다. 수학에 필요한 값만 담으면 된다. 이를 위해 레지스터 할당용 “템플릿”을 만들도록 레이아웃을 슬라이싱한다:

cpp
// Stage 차원을 제거하기 위해 0번째 슬라이스 사용
Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); 
Tensor tCrB = thr_mma.partition_fragment_B(sB(_,_,0));

알고리즘은 단순 루프에서 Prologue + Loop 구조로 바뀐다. prologue는 루프 진입 전에 첫 타일을 로드해 “펌프를 프라임”하는 과정이다.

1. Prologue:

k=0에 대한 로드를 발행하고 async 그룹을 커밋한 뒤 기다린다.

cpp
// PROLOGUE: k=0을 smem_write 스테이지로 로드
copy(copyA, tAgA(_,_,_,0), tAsA(_,_,_,smem_write));
copy(copyB, tBgB(_,_,_,0), tBsB(_,_,_,smem_write));
cp_async_fence();

// write 스테이지 뒤집기
smem_write = (smem_write + 1) % 2;

// k=0이 도착할 때까지 대기(계산 시작을 위해)
cp_async_wait<0>();
__syncthreads();

1. 파이프라인된 루프:

메인 루프에서는 현재 데이터를 계산하면서 동시에 다음 타일 로드를 발행한다. 제공된 코드에서는 공유 메모리 지연을 더 숨기기 위해 레지스터 로드도 타일링(K_BLOCK_MAX에 대해 루프)하지만, 더블 버퍼링의 핵심 로직은 그대로다:

cpp
for (int k = 0; k < k_tile_max; k++) {
    // ... 레지스터 블록에 대한 내부 루프 ...
    
    // 타일의 시작이라면 NEXT 타일(k+1)에 대한 글로벌 로드 발행
    if (k_block == 0 && k_tile_next < k_tile_max) {
         copy(copyA, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,smem_write));
         copy(copyB, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,smem_write));
         cp_async_fence();
         
         // write 스테이지 인덱스 뒤집기
         smem_write = (smem_write + 1) % 2;
         k_tile_next++;
    }

    // 이전 단계/prologue에서 기다렸던 데이터로 연산
    gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);

    // 현재 타일이 끝났다면 read 스테이지를 뒤집고 위에서 발행한 로드를 기다림
    if (k_block == K_BLOCK_MAX - 1) {
        smem_read = (smem_read + 1) % 2;
        cp_async_wait<0>();
        __syncthreads();
    }
}

smem_write와 smem_read 인덱스를 수동으로 관리함으로써, Async Copy 엔진이 한 스테이지를 채우는 동안 텐서 코어가 다른 스테이지를 소비하도록 보장한다. 이는(계산이 로드보다 오래 걸린다는 가정하에) 글로벌 메모리 지연을 대부분 숨기며, 하드웨어 한계에 가까운 대역폭 상승을 만들어 cuBLAS의 111%까지 끌어올린다.

효율적인 더블 버퍼링을 하더라도 Partition Camping이라는 현상 때문에 성능 상한을 맞을 수 있다.

커널을 런치하면 CUDA는 보통 스레드 블록(CTA)을 래스터-스캔 순서(X 증가 후 Y 증가)로 배치한다. A, B 행렬이 큰 연속 메모리 블록(행우선/열우선)이기 때문에, 인접한 블록들이 같은 L2 캐시 파티션으로 매핑되는 주소를 자주 접근한다. 그 결과 특정 메모리 컨트롤러에 높은 경쟁이 생기고 다른 컨트롤러는 놀게 된다.

이를 고치려면 스레드 블록들이 타일을 소비하는 순서를 다시 스위즐해야 한다. 행 단위로 처리하는 대신 블록화된 패턴으로 처리한다. 이는 메모리 접근을 파티션 전체에 더 고르게 분산시키고, 새 블록에 필요한 데이터가 이미 L2에 상주할 가능성을 높인다.

커널에서는 이 로직이 맨 위에 위치한다. 물리적인 그리드 좌표(blockIdx)를 논리적인 문제 좌표(m_idx, n_idx)에서 분리한다.

cpp
// 1. GRID SWIZZLING (L2 캐시 접근 최적화)
    int m_idx = blockIdx.x;
    int n_idx = blockIdx.y;

    // 스위즐 팩터: A100에서 128x128에 대해 8이 잘 동작(L2 슬라이스 폭과 매칭)
    constexpr int swizzle_factor = 8; 
    
    if (gridDim.x >= swizzle_factor) {
        // 그리드를 선형 ID로 평탄화
        int tid = blockIdx.x + blockIdx.y * gridDim.x;
        
        int idx_outer = tid / swizzle_factor;
        int idx_inner = tid % swizzle_factor;
        
        int m_grid = gridDim.x;
        
        // “Z” 커브를 만들도록 좌표 재계산
        int n_swizzled = idx_outer / ((m_grid + swizzle_factor - 1) / swizzle_factor);
        int m_swizzled = (idx_outer % ((m_grid + swizzle_factor - 1) / swizzle_factor)) * swizzle_factor + idx_inner;

        // 경계 체크
        if (m_swizzled < m_grid) {
             m_idx = m_swizzled;
             n_idx = n_swizzled;
        }
    }

CuTe의 장점은 이 좌표 변환이 타일링 로직과 완전히 분리되어 있다는 점이다. 새로운 m_idx와 n_idx를 얻으면 좌표 객체를 만들기만 하면 된다:

cpp
// 스위즐된 좌표를 CTA 타일에 사용
    auto cta_coord = make_coord(m_idx, n_idx, _);

나머지 커널(gA_tile, gB_tile 로드 및 partition)은 완전히 동일하다. 수학이나 복사 로직을 한 줄도 다시 쓰지 않고도, 장치 전체의 실행 순서를 바꿔 하드웨어 활용을 최적화했다. 이는 실행 속도를 약간 개선해 cuBLAS의 112%까지 도달한다.

더블 버퍼링은 현재 계산과 다음 메모리 로드를 겹치게 해주지만, 빡빡한 의존성을 만든다: 현재 타일의 계산 시간은 다음 타일 로드 시간보다 반드시 길어야 한다. 글로벌 메모리가 순간적으로 혼잡하거나 계산이 너무 빠르면 파이프라인에 버블이 생긴다.

파이프라인을 더 견고하게 하기 위해 3-스테이지 버퍼링으로 이동한다. “다음 로드, 현재 계산” 대신 “다다음 로드, 다음 로드, 현재 계산”을 지향한다. 이는 준비된 데이터의 더 깊은 저장소를 만들어 메모리 도착 시간의 지터를 완화한다.

그림 4: 3-스테이지 파이프라이닝 시각화. 두 로드는 동시에 발행되지 않는다.

CuTe에서 2→3 스테이지 전환은 구조적으로는 단순하지만, 동기화 관점에서는 개념적으로 섬세하다. 먼저 공유 메모리 할당을 확장한다. 이전보다 공유 메모리를 50% 더 할당한다:

cpp
// 이전: make_shape(Int<BM>{}, Int<BK>{}, Int<2>{})
// 신규: 3 스테이지 사용
constexpr int STAGES = 3;
auto sA = tile_to_shape(swizzled_128B_atom, make_shape(Int<BM>{}, Int<BK>{}, Int<STAGES>{}));
auto sB = tile_to_shape(swizzled_128B_atom, make_shape(Int<BN>{}, Int<BK>{}, Int<STAGES>{}));

Prologue에서 알고리즘 복잡성이 크게 증가한다. k를 계산하는 동안 k+2를 로드해야 하므로, 루프가 돌기 전에 두 개 타일이 이미 비행 중이어야 한다.

cpp
// PROLOGUE: 첫 (STAGES - 1)개의 k-타일을 프리패치
// 3-스테이지의 경우 루프 전에 k=0과 k=1을 로드
for (int k_tile = 0; k_tile < STAGES - 1; k_tile++) {
    copy(copyA, tAgA(_,_,_,k_tile_index), tAsA(_,_,_,k_tile));
    copy(copyB, tBgB(_,_,_,k_tile_index), tBsB(_,_,_,k_tile));
    
    k_tile_index++;
    // 비동기 복사 그룹 커밋
    cp_async_fence(); 
}

더블 버퍼링에서는 cp_async_wait<0>을 사용했는데, 이는 “대기 중인 연산이 0개가 될 때까지 기다려라”(즉, 전부 끝날 때까지 기다려라)라는 뜻이다. 3-스테이지 버퍼링에서는 여러 복사 그룹을 동시에 비행시킨다. 전부가 끝날 필요는 없고, 가장 오래된 것(곧 계산할 대상)만 끝나면 된다.

cpp
// STAGES가 비행 중이다. 가장 오래된 것이 준비되어야 한다.
// 그러므로 (STAGES - 2) 그룹만 남도록 기다린다.
cp_async_wait<STAGES - 2>();
__syncthreads();

루프 내부에서는 두 단계 앞을 보는 로직으로 바뀐다:

cpp
// 1. k+2(“다다음” 타일)에 대한 글로벌 로드 발행
// smem_write에 쓰는데, 이는 순환적으로 smem_read보다 2스텝 앞
if (k_tile + STAGES - 1 < k_tile_max) {
    copy(copyA, tAgA(_,_,_,k_tile_index), tAsA(_,_,_,smem_write));
    copy(copyB, tBgB(_,_,_,k_tile_index), tBsB(_,_,_,smem_write));
    cp_async_fence();
    
    // 포인터 전진
    smem_write = (smem_write + 1) % STAGES;
    smem_read  = (smem_read  + 1) % STAGES;
}

// 2. 방금 기다렸던 가장 오래된 타일(k)로 계산
gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);

3-스테이지 버퍼링은 “공짜 점심”이 아니다. 블록당 공유 메모리 요구량이 증가하므로 맞추기가 더 어렵다.

NVIDIA GPU에서는 점유율(각 SM에서 활성 워프 수)이 종종 공유 메모리 용량에 의해 제한된다. 2→3 스테이지로 이동하며 블록 크기가 어떤 임계값(예: 48KB→72KB)을 넘으면, GPU는 한 SM에서 동시에 실행할 수 있는 스레드 블록 수를 줄여야 할 수 있다(우리 경우 3→2). 점유율이 너무 낮아지면 스레드 스위칭으로 메모리 로드 및 기타 지연을 숨기는 능력이 떨어져, 더 깊은 파이프라인의 이점이 상쇄될 수 있다.

하지만 큰 행렬 크기(예: 2048x2048)에서는 개선된 오버랩이 감소한 점유율보다 이득이 크다. 벤치마크에서 이 3-스테이지 구현은 성능을 더 끌어올려 cuBLAS의 116%에 도달했다.

솔직히 말하면 2048x2048을 타겟으로 삼은 이유는 그 지점에서 cuBLAS 성능이 가장 나빠 보였기 때문이다. 내 커널은 전반적으로 라이브러리와 비슷하게 추종하지만, 2048에서의 특정한 우위는 내 기발함이라기보다 그들이 상대적으로 덜 최적화된 탓일 가능성이 크다. 그럼에도 A100 출시 이후 여러 해가 지난 지금도 현실적인 크기에서 벤더 튜닝 성능을 넘길 수 있음을 보여주는 것은 만족스럽다.

그림 5: 다양한 행렬 크기에서 cuBLAS와 비교한 우리 커널 성능. 타일 형태 선택 방식 때문에, 우리 커널은 2048x2048 같은 정사각 행렬에서 가장 잘 동작한다.

내 구현을 cuBLAS의 SASS 디스어셈블리와 비교하면 두 가지 핵심 차이가 드러난다. 첫째, 타일링 전략이 약간 다르다. 둘째, 서로 다른 비동기 복사 명령을 사용한다: 나는 SM80_CP_ASYNC_CACHEALWAYS(L2와 L1에 기록)를 쓰지만, cuBLAS는 캐시 오염을 줄이기 위해 L1을 우회하는 SM80_CP_ASYNC_CACHEGLOBAL을 사용한다. 현재 내 구현은 성능을 위해 L1 히트에 의존하고 있으므로, 레이아웃 조정 없이 우회 명령으로 바꾸면 처리량이 떨어진다.

앞으로의 추가 최적화에는 다음이 포함된다:

• 타일링 정교화: A100의 특정 캐시 계층에 더 잘 맞도록 타일 차원을 실험.

• Hopper(H100): 비동기 복사와 타일링 로직을 하드웨어로 가속하는 TMA(Tensor Memory Accelerator)를 활용하도록 파이프라인을 적응. 다만 모든 입력 크기에 TMA를 적용하는 것이 최적은 아니다. 핸들을 계산하는 데 레지스터와 시간이 더 들어, I/O 인덱스 계산보다 비용이 커질 수 있다.

• Blackwell(B200): 새로운 Tensor Memory 기능 활용.

지금까지는 CUDA C++로 모든 것을 구현했다. 하지만 NVIDIA는 이제 대안을 제공한다: 동일한 효율적인 코드로 컴파일되는 Python DSL이다. 우리의 3-스테이지 커널이 어떻게 번역되는지 보자.

NVIDIA는 강력한 컴파일러가 뒷받침하는 더 높은 수준의 추상화로 점점 이동하고 있다. C++는 절대적 제어(직접 PTX 어셈블리 통합, 세밀한 컴파일 관리) 측면에서 여전히 황금 표준이지만, CuTe Python DSL도 매력적인 대안이다.

Python은 개발 속도에서 이점을 갖는다: JIT로 작은 프로젝트에서 더 빠른 컴파일, PyTorch 통합을 통한 더 단순한 스캐폴딩, 표준 Python 디버거/프로파일러와의 매끄러운 통합이 가능하다. 반면 C++는 인라인 어셈블리를 쓰거나 복잡한 빌드 시스템 통합이 필요할 때 강점을 유지한다. 중요한 점은, 보겠지만 Python의 편의성이 더 이상 속도의 대가를 요구하지 않는다는 것이다.

이 섹션에서는 이전 섹션에서 개발한 완전한 3-스테이지 파이프라인 GEMM 커널을 CUDA에서 Python으로 포팅하고, 구현 접근을 비교하며, 성능 차이를 측정한다.

타입 선언과 커널 정의 CUDA에서는 타입 시스템 전반에 걸쳐 템플릿 파라미터가 필요하다:

cpp
template <class ProblemShape, class CtaTiler, 
          class StrideA, class SmemLayoutA, ...>
__global__ void kernelCuteSwizzledPipeline3StageOptimized(
    ProblemShape problem_shape, CtaTiler cta_tiler,
    bf16 *A, StrideA a_stride, SmemLayoutA a_smem_layout,
    ...)

Python DSL은 데코레이터와 런타임 객체를 사용한다:

python
@cute.kernel
def gemm_kernel(
    mA: cute.Tensor,
    mB: cute.Tensor,
    mC: cute.Tensor,
    sA_layout: cute.ComposedLayout,
    sB_layout: cute.ComposedLayout,
    ...
):

@cute.kernel 데코레이터는 CUDA C++로의 JIT 컴파일을 처리하고, 템플릿 인스턴스화를 자동으로 생성한다. 이는 템플릿 메타프로그래밍을 제거하면서도 컴파일 타임 최적화를 유지한다.

레이아웃 구성 스위즐된 공유 메모리 레이아웃은 문법적 변환을 잘 보여준다. CUDA 버전은 템플릿 타입을 합성한다:

cpp
auto swizzled_128B_atom = composition(
    Swizzle<3,3,3>{},
    make_layout(
        make_shape(Int<8>{}, make_shape(Int<8>{}, Int<8>{})),
        make_stride(Int<8>{}, make_stride(Int<1>{}, Int<64>{})))
);

Python은 중첩 튜플로 같은 레이아웃을 표현한다:

python
layout_atom_outer = cute.make_layout(
    (8, (8, 8)),
    stride=(8, (1, 64)),
)
swizzle = cute.make_swizzle(3, 3, 3)
swizzle_atom = cute.make_composed_layout(swizzle, 0, layout_atom_outer)

두 표현 모두 계층 구조를 가진 8×64 레이아웃 atom을 인코딩한다. 하지만 중첩 튜플은 타입 대수보다 수학적 표기처럼 읽힌다. 의미는 동일하고, 문법은(내 의견으로는) 더 접근하기 쉬워진다.

텐서 파티셔닝과 인덱싱 핵심 partition API는 거의 변하지 않는다. 두 버전 모두 동일한 메서드 이름을 사용한다:

cpp
// CUDA
auto thr_copy_a = copyA.get_thread_slice(threadIdx.x);
auto tAgA = thr_copy_a.partition_S(gA_tile);  // Source
auto tAsA = thr_copy_a.partition_D(sA);       // Destination

python
# Python
thr_copy_A = tiled_copy_A.get_slice(tidx)
tAgA = thr_copy_A.partition_S(gA)  # (CPY, CPY_M, CPY_K, k_tiles)
tAsA = thr_copy_A.partition_D(sA)  # (CPY, CPY_M, CPY_K, STAGES)

partition 연산은 스레드별 텐서 뷰를 만든다. 하지만 텐서 인덱싱에서 가장 눈에 띄는 문법 차이가 드러난다:

copy(copyA, tAgA(_,_,_,k_tile_index), tAsA(_,_,_,k_tile));

python
cute.copy(tiled_copy_A, tAgA[None, None, None, k_tile_index], 
          tAsA[None, None, None, k_tile])

CUDA는 _로 전체 차원을 선택하고, Python은 None을 사용한다. 이는 NumPy의 브로드캐스팅 관례와 맞아 과학 Python 사용자에게 익숙하다. 둘 다 같은 연산(첫 세 모드의 모든 원소를 선택하고 네 번째를 인덱싱)을 의미한다.

메모리 관리 추상화는 상당히 다르다. CUDA는 수동 포인터 산술이 필요하다:

cpp
extern __shared__ char smem_[];
bf16* smem_ptr = reinterpret_cast<bf16*>(smem_);
Tensor sA = make_tensor(make_smem_ptr(smem_ptr), a_smem_layout);
Tensor sB = make_tensor(make_smem_ptr(smem_ptr + cosize_v<SmemLayoutA>), 
                        b_smem_layout);

Python은 오프셋을 자동으로 처리하는 할당기를 제공한다:

python
smem = cutlass.utils.SmemAllocator()
sA = smem.allocate_tensor(io_dtype, sA_layout, 128)
sB = smem.allocate_tensor(io_dtype, sB_layout, 128)

할당기는 바이트 오프셋을 계산하고 정렬을 관리하며 타입이 붙은 텐서 뷰를 반환한다.

메인 파이프라인 3-스테이지 파이프라인 로직은 언어 간에 직접적으로 번역된다. 두 구현 모두 prologue에서 첫 두 타일을 프리패치한다:

cpp
// CUDA Prologue
for (int k_tile = 0; k_tile < num_smem_stages - 1; k_tile++) {
    if (k_tile < k_tile_max) {
        copy(copyA, tAgA(_,_,_,k_tile_index), tAsA(_,_,_,k_tile));
        copy(copyB, tBgB(_,_,_,k_tile_index), tBsB(_,_,_,k_tile));
    }
    k_tile_index++;
    cp_async_fence();
}

python
# Python Prologue
for k_tile in range(num_smem_stages - 1):
    if k_tile < k_tile_max:
        cute.copy(tiled_copy_A, tAgA[None, None, None, k_tile_index], 
                  tAsA[None, None, None, k_tile])
        cute.copy(tiled_copy_B, tBgB[None, None, None, k_tile_index], 
                  tBsB[None, None, None, k_tile])
        k_tile_index = k_tile_index + 1
    arch.cp_async_commit_group()

메인 루프 구조도 일관된다. 각 타일의 첫 k-block에서 다음 타일에 대한 글로벌 로드를 발행한다. 각 계산 전에 다음 k-block용 레지스터 데이터를 프리패치한다:

python
for k_tile in range(k_tile_max):
    for k_block in cutlass.range_constexpr(k_block_max):
        if k_block == k_block_max - 1:
            smem_read_current = smem_read
            arch.cp_async_wait_group(num_smem_stages - 2)
            arch.sync_threads()
        
        # 다음 k-block을 레지스터로 프리패치
        next_k_block = (k_block + 1) % k_block_max
        cute.copy(tiled_copy_s2r_A, 
                  tCsA_copy_view[None, None, next_k_block, smem_read_current],
                  tCrA_copy_view[None, None, next_k_block])
        
        if k_block == 0:
            # 다음 타일을 위한 글로벌 로드 발행
            if k_tile + num_smem_stages - 1 < k_tile_max:
                cute.copy(tiled_copy_A, tAgA[None, None, None, k_tile_index],
                          tAsA[None, None, None, smem_write])
        
        # 현재 데이터로 계산
        cute.gemm(tiled_mma, tCrC, tCrA[None, None, k_block], 
                  tCrB[None, None, k_block], tCrC)

Python 버전은 루프 언롤링을 위해 cutlass.range_constexpr()를 사용하며, #pragma unroll을 명시적 API 호출로 대체한다. 파이프라인 상태 관리(smem_read, smem_write, k_tile_index 추적)는 두 언어에서 동일한 로직을 따른다.

2048×2048 BFloat16 행렬에 대해 NCU로 측정한 벤치마크 결과는, 3-스테이지 파이프라인 커널이 149 TFLOPS를 달성하며 PyTorch의 128 TFLOPS(상대 116%)와 NVIDIA 자체 Python CuTe DSL 예제의 142 TFLOPS(상대 105%)를 상회함을 보여준다. Python CuTe DSL 커널은 0.1155ms, 동등한 CUDA 코드는 약 0.1158ms로, 사실상 CUDA 버전과 동일한 성능을 달성한다.

이 성능 동등성은 Python DSL이 JIT를 통해 동일한 CUDA 코드로 컴파일되어야 하기 때문에 발생한다.

그림 6: Python CuTe DSL로 포팅한 3-스테이지 파이프라인 커널을 표준 PyTorch(cuBLAS 사용) 및 NVIDIA의 Python CuTe DSL 예제와 비교. 역시 우리 커널은 2048x2048 정사각 행렬에서 가장 잘 동작한다.

이 2부작 시리즈는 CuTe의 레이아웃 대수와 이를 CUDA 커널 및 Python DSL 대응물에서 사용하는 방법을 보여줬다. CuTe 레이아웃 대수를 살펴보며 (Shape, Stride) 튜플이라는 기초에서 시작해 타일링(⊘)과 합성(∘)으로 확장했고, 이론적으로 GPU 커널에서 인덱싱 오류를 피하는 강력한 추상화임을 보여줬다. 기본 구현에서 스위즐된 비동기 복사까지 메모리 복사를 점진적으로 최적화했다. 이 복사 커널을 바탕으로, 2048×2048 행렬에서 101%의 cuBLAS(베이스라인)에서 시작해 더블 버퍼링(111%), L2 스위즐링(112%), 3-스테이지 파이프라이닝(116%)으로 발전하는 GEMM 커널을 만들었다. CuTe 레이아웃 추상화는 코드를 깔끔하게 만들었고, 성능 손실 없이 선언적인 방식으로 GEMM 커널을 구축하는 데 도움을 줬다. 또한 전체 커널을 Python으로 거의 동일한 성능으로 포팅할 수 있었고, 대부분의 행렬 크기에서 NVIDIA가 제공한 CuTe Python 예제보다 더 빠른 실행 시간을 보여줬다.

CuTe를 언제 사용할까. 표준 연산에서 튜닝이 뛰어난 cuBLAS/cuDNN을 사용하라. 반면 커스텀 퓨전 커널, 새로운 아키텍처(어텐션 변형, 상태공간 모델) 구축, 혹은 타일링 전략을 빠르게 반복해야 한다면 CuTe를 선택하라. 레이아웃 대수는 partition을 명시적이고 조합 가능하게 만든다. 매우 특수한 커널은 여전히 손으로 튜닝한 CUDA가 이득일 수 있지만, 아직 그런 요구가 필요한 사례는 보지 못했다.

Python DSL. Python CuTe DSL은 850줄짜리 CUDA 커널을 ~650줄로 줄이면서도 JIT 컴파일을 통해 거의 동일한 성능을 달성했다. Python은 더 빠른 반복(컴파일 오버헤드가 훨씬 짧음), 매끄러운 PyTorch 통합, 표준 도구(디버거, 프로파일러, Jupyter)를 제공한다. C++는 직접적인 PTX 제어와 기존 CUDA 코드베이스 측면에서 장점이 있다. 두 구현이 cuBLAS와 동급 성능을 보인다는 것은 기반 추상화가 타당함을 입증한다. 언어 선택은 이제 성능 역량이 아니라 개발 맥락의 문제로 바뀐다.

이 블로그 글을 위한 컴퓨트를 제공해준 Verda와 Paul에게 큰 감사를 전하고, 벤치마킹을 제대로 하는 법에 대한 교정과 가이드를 해준 Szymon에게도 감사하며, 코드의 멍청한 버그들을 고치는 데 도움을 준 Lukas에게도 감사한다.

• 문서: Nvidia | Github

• GPUMode 강의 15 (CUTLASS, CuTe 대수 개요)

• GPUMode 강의 57 (Cris Cecka와 함께 하는 CuTe 심층 분석)

• 1부: CuTe 레이아웃 이해하기 — 이번 글의 모든 내용의 기초