Getting Started with CuTe DSL — Python 으로 만나는 CuTe

CUTLASS 는 NVIDIA 의 production GEMM/Conv 라이브러리. 그 안의 CuTe 가 layout 과 tile 의 algebra 를 표현하는 핵심 추상이다. 그러나 CuTe 는 C++ template으로 짜여 있고 — 한 번 코드를 펼친 사람이라면 안다 — 진입 장벽이 매우 높다. 이 강의는 CuTe 의 같은 추상을 파이썬으로 직접 만지는 새 DSL 을 깐다.

NVIDIA 가 CuTe Python DSL 을 만든 의도는 — 외부에 공개된 자료 흐름에서 — 세 갈래로 모인다.

1. CUTLASS 의 표현력은 살리고 진입 장벽은 낮춘다 — 같은 layout algebra 를 Python 에서 만지고, 빠른 iteration 가능.
2. Triton 이 못 다루는 자리 — Triton 은 사용자에게 layout 을 직접 노출하지 않는다. CuTe DSL 은 layout 의 모든 비트를 노출한다 — Tensor Core 의 mma 명령에 정확히 맞추는 자리에서 가치.
3. 새 hardware 의 빠른 채택 — Hopper 의 TMA, Blackwell 의 새 tensor memory accelerator 같은 새 명령은 CuTe 차원에서 가장 먼저 노출. DSL 이 같이 빨리 따라가야 산업이 받아쓸 수 있다.

CUTLASS C++ 코드를 한 번이라도 펼친 사람이라면 — 진입의 어려움이 어디서 오는지 안다. 그 어려움을 명시화한다.

CuTe Python DSL 의 답 — 같은 추상을 그대로 Python 에 노출. C++ template 의 deep nesting 없이, 같은 layout / tensor / copy / mma 를 Python 객체로 직접 만지고, 컴파일은 즉시 (JIT) 도는 형태.

CuTe 의 추상은 — C++ 이든 Python 이든 — 네 가지 객체에 모인다. 이 넷의 의미를 한 번 박아두면 나머지가 펼쳐진다.

이 네 객체로 한 GEMM 의 추상이 거의 다 표현된다 — “G/S/R memory 의 layout 들이 있고, 그 사이를 copy 가 stage 하고, 마지막에 MMA 가 누적” 의 그림.

CuTe 의 가장 큰 진입 장벽이자 가치 — layout algebra. 한 번 직관이 잡히면 같은 표기로 row-major, column-major, transposed, swizzled, blocked 가 모두 표현된다.

강의의 핵심 데모가 정확히 이 자리. 같은 GEMM 을 — Triton 도 아니고 CUTLASS C++ 도 아니고 — CuTe Python DSL 로 짜는 시퀀스. 외부 자료로 재구성한 개념적 형태.

// 깊은 template 의 한 예
using GemmKernel = cutlass::gemm::collective::CollectiveBuilder<
  cutlass::arch::Sm90,
  cutlass::arch::OpClassTensorOp,
  cutlass::half_t, LayoutA, AlignmentA,
  cutlass::half_t, LayoutB, AlignmentB,
  ElementAccumulator,
  TileShape, ClusterShape,
  cutlass::gemm::collective::StageCountAuto,
  cutlass::gemm::collective::KernelScheduleAuto
>::CollectiveOp;

// 이걸 device::Gemm 에 한 번 더 wrap 하고...
// 메인 함수에서 args, workspace, smem 사이즈 까지 직접 셋업

import cutlass.cute as cute

@cute.kernel
def gemm(A, B, C,
         tA: cute.TiledCopy,
         tB: cute.TiledCopy,
         tC: cute.TiledMMA):
    bid_m = cute.block_idx(0)
    bid_n = cute.block_idx(1)

    # tile 추출 — layout algebra 로
    gA = cute.local_tile(A, (BM, BK), (bid_m, _))
    gB = cute.local_tile(B, (BN, BK), (bid_n, _))
    gC = cute.local_tile(C, (BM, BN), (bid_m, bid_n))

    sA = cute.make_smem_tensor(BM, BK)
    sB = cute.make_smem_tensor(BN, BK)
    rC = cute.make_reg_tensor(BM, BN)

    for k in range(K // BK):
        cute.copy(tA, gA[:,:,k], sA)
        cute.copy(tB, gB[:,:,k], sB)
        cute.cp_async_wait()
        cute.gemm(tC, sA, sB, rC)

    cute.copy(tC, rC, gC)

두 코드의 의도는 같다 — Hopper Tensor Core 위에서 fp16 GEMM 을 짜는 일. 차이는 표현의 두께. CuTe Python 은 같은 추상 (TiledCopy, TiledMMA, local_tile, copy) 을 명시적으로 다루지만 — template 의 깊이가 사라지고, JIT 로 빠르게 컴파일된다.

A: GMEM   M × K   ──tA──▶  sA: SMEM    BM × BK
B: GMEM   N × K   ──tB──▶  sB: SMEM    BN × BK
                              │
                              ▼
                         tC (TiledMMA)
                              │
                              ▼
                         rC: RMEM   BM × BN
                              │
                              └──tC──▶ C: GMEM   BM × BN

중요한 사실 한 가지 — CuTe Python DSL 이 만드는 PTX 는 CUTLASS C++ 가 만드는 PTX 와 같은 quality를 목표로 한다. 표면(C++ vs Python) 만 다르고, 깎은 결과는 같다는 게 디자인 의도. 그래서 production 에서 “Python DSL 로 짰는데 더 느리다” 같은 의외 없음.

• NVIDIA 내부 production GEMM — cuBLAS / CUTLASS 본체 안의 일부. 새 mma 변형의 첫 표현.
• FlashAttention-3 — Tri Dao 등의 FA-3 가 CUTLASS / CuTe 위에서. Hopper TMA 와 producer-consumer 분리 — Triton 이 표현 어려운 자리.
• BlackForestLab / Cohere 등의 production 학습 stack — 일부 hot path 가 CUTLASS / CuTe 기반 — 확인 필요.
• 학계의 새 알고리즘 prototyping — layout 자체가 핵심인 새 fusion 패턴 (예: low-bit GEMM 의 layout 변형).
• Triton 의 backend 로 — 확인 필요. Triton 의 일부 hopper 백엔드가 CuTe 의 추상을 차용하는 흐름.

강의에서 Vicki 가 강조했을 입장 — CuTe DSL 은 Triton 의 대체가 아니라 보완. 두 도구가 사용자 표면 차원에서 같이 산다는 디자인.

1. Blackwell 의 새 명령 노출 — tensor memory accelerator, new mma 변형, FP4 mma. CuTe DSL 이 가장 먼저 받는 자리.
2. autotune 통합 — layout / tile size 의 search 를 자동화. L082 Helion 의 흐름과 합쳐질 가능성.
3. PyTorch 통합 — torch.compile 과의 자연스러운 결합. 사용자가 PyTorch 안에서 CuTe 커널을 부르고, autograd 가 자동.
4. 비-NVIDIA backend — 확인 필요. 현재 NVIDIA only. AMD MI300 등 다른 hardware 로의 확장 가능성.
5. 학습 자료 — layout algebra 의 진입 장벽을 낮추기 위한 시각화 도구, 인터랙티브 튜토리얼. L103 의 본격 카테고리이론적 해석도 같은 흐름.

손에 새기기 — 실습 시퀀스

1. layout 의 그림 직접 그리기 — 종이 위에 (Shape=(4,4), Stride=(4,1)) 와 (Shape=(4,4), Stride=(1,4)) 의 메모리 offset 을 모두 그려본다. transpose 가 데이터 이동 없이 가능한 이유 손에.
2. nested layout — Shape=((2,2),(2,2)) 의 4×4 의 offset 을 그려본다. tile 단위 layout 의 직관.
3. CUTLASS Python DSL 설치 + first kernel — 공식 README 의 hello world. JIT 컴파일 시간 직접 측정.
4. simple GEMM — § 05 의 형태로 fp16 GEMM 한 번 짜기. cuBLAS baseline 과 시간 비교.
5. 같은 GEMM 의 lowering dump — CuTe IR / MLIR / PTX 를 각각 dump 해본다. layout 이 어떻게 thread mapping 으로 lowering 되는지 추적.

• CuTe Python DSL 의 정확한 API — 노트의 코드 예시는 개념적 형태. cute.local_tile, cute.copy, cute.gemm 같은 정확한 함수 이름은 NVIDIA 공식 문서 확인 필요.
• release 상태 — public release 인지, NVIDIA 내부 도구인지, beta 인지. 발표 시점에 따라 다름.
• JIT 컴파일 시간 실제 수치 — § 02 의 “초 단위” 가 정확한가. 첫 호출의 cold start 와 두 번째의 warm 시간 비교.
• 비-NVIDIA backend 가능성 — AMD / Intel 으로의 확장 — NVIDIA 의 입장과 산업의 호환성 요구.
• Triton 과의 정확한 경계 — 같은 fused attention 을 두 도구로 짠 결과의 성능 비교가 어디 있는가.
• FlashAttention-3 의 채택 — FA-3 가 CUTLASS / CuTe 위에 어떻게 짜였는가. 그 코드를 CuTe Python 으로 재구현하면 같은 성능이 나오는가.