CuTe — CUTLASS Tensor의 layout algebra

SOTA GPU GEMM 또는 attention kernel 안 — 같은 tensor 가 동시에 여러 모양으로 보여야 한다. global memory 에서는 (M, K), shared memory 에서는 swizzled (M_block, K_block), register 에서는 thread 별로 분할된 tile, MMA instruction 에서는 architecture 별로 정의된 shape — 같은 데이터에 대한 다섯 layout. 이 다섯 사이의 변환을 손으로 짜면 인덱스 계산만으로 코드의 절반.

강의가 풀려는 두 질문.

1. layout 자체를 1급 객체로 다루는 algebra 가 있다면 — 같은 tensor 의 여러 view 사이 변환을 어떻게 generic 하게 표현할까?
2. 그 algebra 가 GPU kernel 작성을 어떻게 단순화하는가 — 특히 Hopper 의 새 hardware (TMA, async, MMA) 와 결합될 때.

본 노트는 transcript 실패로 — CUTLASS GitHub repo (github.com/NVIDIA/cutlass), CuTe 의 media/docs/cute/ 문서, Cris Cecka 의 GTC talk, Pradeep Ramani 의 CuTe blog 시리즈로 재구성.

CuTe 의 layout 은 (shape, stride) pair — shape 은 multi-dim 의 크기, stride 는 한 차원으로 한 칸 움직일 때 linear index 가 얼마나 변하는지. 같은 데이터를 다른 shape/stride 로 보면 다른 layout — 같은 메모리, 다른 view.

그리고 — layout 은 compile-time 또는 runtime. shape/stride 의 일부가 compile-time const 면 — 컴파일러가 그 부분을 unroll/inline. 일부가 runtime 이면 그 부분만 일반 코드. CuTe 의 type system 이 두 모드를 혼합 가능.

CuTe 의 가장 큰 디자인 — hierarchical mode. shape 가 그냥 (M, N) 이 아니라 — 같은 차원이 ((tile_M, num_tiles_M), N) 같이 중첩. 이게 GPU 의 hierarchy (thread → warp → block → grid) 와 정확히 맞물린다.

// hierarchical shape — CuTe
// 4096 = 4 * 16 * 2 * 32
auto shape  = make_shape(_4{}, _16{}, _2{}, _32{});
auto stride = make_stride(_1{}, _4{}, _64{}, _128{});
auto layout = make_layout(shape, stride);

// 한 thread 의 입장 — (t, w, w_idx, b)
// t = thread idx within warp (0..3)
// w_idx = thread group within warp (0..15)
// w = warp within block (0..1)
// b = block idx (0..31)
auto coord = make_coord(t, w_idx, w, b);
auto linear_idx = layout(coord);
// linear_idx 가 그 thread 가 다룰 element

hierarchical layout 의 또 다른 사용 — data 의 swizzling. shared memory 의 bank conflict 를 피하기 위해 data 를 permuted layout 으로 두기. 같은 algebra 안에서 swizzled view 를 표현 가능.

CuTe 의 핵심 — layout 위에 정의된 연산들. compose, divide, coalesce. 이 셋이 복잡한 layout 변환을 generic 하게 표현한다.

// CUTLASS 3.x GEMM 의 hot loop (의사코드)
auto A_tile = local_tile(A, tile_shape, block_coord);
auto B_tile = local_tile(B, tile_shape, block_coord);

auto A_thr = local_partition(A_tile, thr_layout, thr_idx);
auto B_thr = local_partition(B_tile, thr_layout, thr_idx);

// thr 단위 fragment register
auto rA = make_fragment_like(A_thr);
auto rB = make_fragment_like(B_thr);
auto rC = make_fragment_like(C_part);

// load + MMA
copy(A_thr, rA);
copy(B_thr, rB);
gemm(rA, rB, rC);     // MMA atom

CuTe 의 “layout 은 함수” 이라는 통찰의 가장 큰 응용 — thread layout 도 같은 객체. (thread idx) → (어떤 element 를 다룰지) 의 매핑이 layout 함수.

실제 사용 — autotuning 의 base. 같은 kernel 의 여러 thread 분배 후보 (4×8, 8×4, 16×2, 2×16) 를 각각 다른 thread layout 으로 두고 — 같은 source 가 컴파일됨. CUTLASS 의 GEMM 이 새 architecture 에 빠르게 적응하는 이유의 큰 부분.

CuTe 의 두 번째 핵심 — hardware instruction 자체를 “atom” 으로 추상. NVIDIA 의 MMA (Matrix Multiply Accumulate) 명령어, TMA (Tensor Memory Accelerator) copy, async copy 등이 모두 atom. 각 atom 은 자기만의 layout 을 강제.

CUTLASS 2.x 는 — C++ template 의 deep nesting + 거의 manual 한 indexing. 새 architecture 마다 코드 거의 새로. CUTLASS 3.x 는 CuTe 위에서 — 같은 kernel 이 layout 변경만으로 새 hardware 적응. 실용적 차이가 큼.

실용적 진입 시퀀스 — (1) CUTLASS examples 의 가장 작은 GEMM 부터 정독. (2) media/docs/cute/ 의 layout tutorial 따라가기. (3) 자기 손으로 layout 객체를 만들고 print 해보기 — print(layout) 가 동작. (4) 작은 GEMM 을 CuTe 로 짜보기. (5) FA3 코드를 한 줄씩 읽기.

손에 새기기 — 실습 시퀀스

1. CuTe layout 만들고 print — auto L = make_layout(make_shape(_4{}, _3{}), make_stride(_3{}, _1{})). print(L) 로 layout 의 모양 확인.
2. 같은 데이터의 두 layout — row-major 와 column-major 를 둘 다 만들고 — 같은 logical (i,j) 에 대해 다른 linear idx 가 나오는지 확인.
3. hierarchical shape — (4, 8) 의 single-mode 와 ((2, 2), 8) 의 hierarchical 두 가지로 — coalesce 후 같은 layout 인지 확인.
4. compose · divide — 작은 layout 두 개를 compose 했을 때의 결과를 print 로 확인. divide 의 결과도 마찬가지로.
5. local_tile · local_partition — 큰 tensor 에서 한 thread block 의 tile 을 잡고, 그 안의 한 thread 의 fragment 를 잡는 패턴.
6. CUTLASS 의 작은 GEMM — examples/00_basic_gemm 을 빌드하고 — 코드를 한 줄씩 주석.
7. thread layout 변경 실험 — 같은 GEMM 을 다른 thread layout (4×8 vs 8×4 vs 16×2) 으로 — 결과 동일한지, 성능 차이는 어떤지.
8. FA3 코드의 한 forward block 정독 — Tri Dao 의 FA3 코드 (CUTLASS 위) 의 한 함수를 한 줄씩 — CuTe 의 어느 도구가 어디 쓰이는지.

• 강의의 정확한 deep-dive — layout algebra 의 수학적 디테일 (compose 의 정확한 정의, partition 의 edge case) 를 강의에서 어디까지 풀었는지. 원본 영상 확인 필요.
• 실전 데모 코드 — 강의 안에서 화이트보드/슬라이드의 구체적 예제. 본 노트는 일반적 패턴으로 재구성.
• swizzling 의 디테일 — bank conflict 회피의 swizzle layout — CuTe 가 어떻게 표현하는지. 강의에서 다뤘는지.
• TMA 의 새 자리 — Hopper TMA 의 descriptor 디자인 — CuTe 가 어떻게 자동 생성하는지의 디테일.
• multi-cluster (Hopper TMA cluster) — 4 thread block 의 cooperative 처리 — CuTe 가 어떻게 표현하는지.
• CUTLASS 2 → 3 migration — 사용자가 자기 코드를 옮기는 실용적 가이드. 강의에서 언급되었는지.
• CuTe 의 future direction — Blackwell (SM_100) 의 새 instruction 들과의 결합. 강의 시점에 어디까지.
• 다른 hardware 와의 비교 — AMD CDNA, Apple Silicon 위에서 비슷한 abstraction 시도 — 강의에서 비교 했는지.