Fundamentals of CuTe Layout Algebra and Category-theoretic Interpretation

CUTLASS 3.x / CuTe 의 코드를 처음 만나면 — layout 이 곳곳에 퍼져 있다. mainloop 의 SMEM tile, epilogue 의 store, threadblock 의 partition, MMA 의 thread layout. 같은 추상이지만 매번 모양이 달라 보인다. 이 강의가 풀려는 문제 — 이 모든 자리에서의 layout 이 “같은 algebra” 로 통합된다는 사실을 정확히 보이는 것.

강의가 던지는 두 개의 질문.

1. layout 의 algebra 가 정확히 무엇인가 — 함수의 합성과 분해라는 frame 이 어디까지 정확하게 적용되는가.
2. 그 algebra 가 “수학적으로 옳다” 는 것의 의미는 무엇인가 — category-theoretic 시각이 어떻게 layout algebra 의 일관성을 보장하는가.

강의의 두 부분. 첫 번째 — algebra 자체 (§ 02–06). 두 번째 — category 의 시각 (§ 07). 첫 부분이 충분히 가치 있으므로 두 번째에 압도되지 않고 따라가는 게 좋은 학습 전략.

CuTe 의 한 줄 정의 — Layout = (Shape, Stride). Shape 가 좌표 공간, Stride 가 좌표를 자연수로 mapping 하는 가중치. Layout(c) 는 좌표 c 를 받아 메모리 index 를 돌려주는 함수.

(0,0) → 0   (0,1) → 1   ...  (0,7) → 7
(1,0) → 8   (1,1) → 9   ...  (1,7) → 15
...
(3,0) → 24  ...           (3,7) → 31

(0,0) → 0   (1,0) → 1   ...  (3,0) → 3
(0,1) → 4   (1,1) → 5   ...  (3,1) → 7
...
(0,7) → 28  ...           (3,7) → 31

Shape 와 Stride 자체가 tuple 이라는 점이 중요. 단순한 정수가 아니라 — nested tuple — 도 가능. 그 자리가 hierarchical layout 의 entry point (§ 04).

layout 의 algebra 는 — 5 개의 핵심 연산으로 거의 모두 표현된다. 각 연산이 함수의 합성/분해/제한 중 어느 것의 표현인지를 명시적으로 잡으면 직관이 단단해진다.

실전에서 자주 쓰는 wrapper — blocked_product, raked_product, zipped_divide, tiled_divide. 모두 logical_product / logical_divide 의 변형으로, mode 의 순서나 grouping 만 다르다.

CuTe 의 가장 큰 unique selling point — hierarchical layout. Shape 가 단순 tuple 이 아니라 nested tuple 일 수 있다. ((2,2), (2,2)) : ((4,1), (8,2)) 같은 모양. 좌표도 같은 모양으로 nested.

이걸 8 으로 나눠 보면:
8I  + 4J + 2i + j
↑outer↑      ↑inner↑

hierarchical layout 의 좌표 형태는 — natural coordinate (nested tuple, 모양이 Shape 를 따름) 또는 — 1-D coordinate (단일 정수, colexicographic 순서). 두 형태가 자유롭게 변환 가능. 실전 코드에서 한 자리에 한 형태로 통일.

SMEM 의 32-bank 구조는 — 같은 bank 의 32 thread access 가 동시에 일어나면 32× slowdown. 이걸 피하기 위한 swizzle 은 — 좌표를 “섞는” 함수. CuTe 에서는 swizzle 도 layout 의 합성으로 표현된다.

SMEM tile 의 layout 만 algebra 의 자리가 아니다. thread 가 tile 의 어느 element 를 담당하는가 도 같은 algebra. partition_S (source-partition), partition_D (destination-partition) 가 그 표현.

ThreadLayout = (32) : (1)
TiledShape   = (16, 16)
partition    = TiledShape ÷ ThreadLayout
             = ((4, 16), 32)
             = thread 안 좌표 4×16 = 8/thread × 8 (?) 

이 시각이 잡히면 — local_tile, partition_S, tiled_divide 같은 CuTe API 들이 모두 같은 algebra 의 다른 사용자 친화 wrapper 라는 사실이 명확해진다. 새 함수마다 외울 게 아니라 — “이 함수는 divide 의 무엇을 wrap 한 것인가” 의 시각으로.

강의의 두 번째 절반. layout 의 algebra 가 — category (범주) 의 morphism 으로 보일 때 — 모든 연산의 일관성이 자연스럽게 따라온다는 시각. 카테고리 이론 백그라운드 없는 사람도 핵심만 잡을 수 있도록 절제된 정리.

주의 — 이 시각이 implementation 을 더 빠르게 만들지는 않는다. CuTe 의 implementation 은 type-level template 으로 정해진 형태가 있고, category 시각은 그 형태가 왜 그러한지 의 답일 뿐. 강의의 메시지도 동일 — “code 를 짤 때 카테고리를 항상 의식할 필요는 없지만, 막혔을 때 카테고리 시각이 길을 보여준다”.

L103 자체가 자료의 합성이고 — 그 자료들을 어떤 순서로 읽으면 가장 효율적인지가 강의의 부수적 가치.

1. CUTLASS 의 cute/00_quickstart.md — Layout = (Shape, Stride) 의 30분 자료. 가장 작은 entry.
2. cute/01_layout.md — coordinate mapping, hierarchical, integer types. 1 시간.
3. cute/02_layout_algebra.md — composition / divide / product / complement 의 정의. 1–2 시간.
4. L103 슬라이드 (layout_algebra.pdf) — 같은 자료의 정리, category 시각 추가.
5. Colfax — Tutorial: Matrix Transpose in CUTLASS — algebra 가 swizzle 위에서 어떻게 작동하는가의 구체.
6. Jay Shah — “A note on the algebra of CuTe Layouts” — paper-level 정리. category 시각의 깊이 보충.
7. L057 — CuTe — 다른 발표자 (Jay Shah) 의 더 가벼운 introduction. 보충 시각.
8. L086 — FlashAttention 3 — algebra 가 production 코드에 어떻게 들어가는가. 응용.
9. L101 — Learning CUTLASS — algebra 위에 build 된 mainloop 의 hierarchy.
10. NVIDIA cutlass example/cute — 실제 코드. transpose / sgemm / FA3.

• data dependent control flow — predication 이 compile-time 으로 결정되지 않으면 layout algebra 의 표현 영역 밖. attention 의 causal mask 같은 자리가 부분적으로 해당.
• dynamic shape — Shape 가 compile-time constant 가 아니면 — 같은 algebra 가 그대로 적용되지만 type-level 의 모든 결정이 runtime 으로 옮겨가 효율이 떨어진다.
• async pipeline 의 sequencing — “이 stage 가 끝나면 다음 stage 가 시작” 같은 dependency 는 layout 의 영역이 아니라 scheduling 의 영역.
• register allocation — 같은 layout 이라도 — register 에 어떻게 매핑되는가는 — compiler / SASS 의 결정. layout algebra 가 직접 통제하지 않는다.
• numerical correctness — algebra 는 “좌표 → index” 의 일관성만 보장. 결과의 numerics 는 algebra 밖.
• cross-CTA collective — TMA multicast 같은 cross-CTA 통신은 — layout 의 stride 만으로 표현하기 어렵다. 추가 추상이 필요.

이 한계들은 “algebra 가 잘못됐다” 는 의미가 아니라 — “이 algebra 는 layout 의 자리만 다룬다” 는 사실. mainloop 의 dependency, async barrier, scheduling 은 다른 추상의 영역이고, 그 영역들과 layout algebra 가 합쳐져 한 GEMM kernel 을 이룬다.

손에 새기기 — 실습 시퀀스

1. 좌표 → index 직접 계산 — 5 개 layout (flat / hier / swizzled) 각각의 좌표를 손으로 메모리 grid 에 그려본다. 함수의 직관 형성.
2. composition 손계산 — A=(4,4):(4,1), B=(2,2):(2,1) 의 A ∘ B 를 손으로. crd2idx 로 검증.
3. logical_divide 의 결과 파악 — A=(8,8):(8,1), B=(4,4):(8,1) 의 A ÷ B 가 어떤 모양인지. CuTe 의 print_layout 으로 검증.
4. swizzle 의 효과 시각화 — 8×8 SMEM 위 naive vs Swizzle<3,0,3> 의 bank id 분포를 직접 그린다.
5. partition_S 의 결과 따라가기 — 16×16 tile 위 32 thread layout 의 partition 결과가 (thread, value) 쌍으로 어떻게 나오는지 손으로.
6. TiledMMA 코드 읽기 — CUTLASS example 의 한 GEMM 의 TiledMMA 를 — atom + atom_layout 의 product 로 분해해 읽기.
7. category 의 한 ax. 검증 — divide 후 product 가 원본을 복원하는 한 케이스를 손으로 확인. adjunction 의 직관.

• category 시각의 정확한 statement — 강의에서 어떤 universal property / functoriality 를 정확하게 statement 했는지 — slide 본문 확인 필요. (확인 필요)
• divide 와 product 의 adjunction 의 정식 형태 — 강의가 categorical adjunction 을 정식으로 statement 했는지, 또는 metaphor 로만 사용했는지. (확인 필요)
• complement 의 universal property — “unique up to iso” 류의 정식 statement 가 있었는지. (확인 필요)
• 구체적 worked example — Hopper WGMMA 의 한 case 를 슬라이드에서 직접 풀었는지. (확인 필요)
• L057 vs L103 의 분담 — Jay Shah 가 두 강의에 모두 참여했다는 가정 하에, 두 강의의 깊이 차이의 정확한 형태. (확인 필요)
• Layout 위 operation 의 complete list — 강의에서 정리한 operation 의 정확한 갯수 / 분류. (확인 필요)