Speaking Composable Kernel (CK)

2018년 AMD MI100 출시. 첫 matrix-core (MFMA) 가 들어가는 generation. NVIDIA 에는 이미 cuDNN/cuBLAS/CUTLASS 가 있었다. AMD 에는 — 같은 위치의 라이브러리 스택이 없었다. CK 가 그 자리를 메우러 출발한 프로젝트.

강의가 깐 큰 질문 세 개.

1. 왜 fused / composable 한 단위로 짜는가 — 한 모델 안에 수백 개의 layer 가 있다. 각각 hand-tuned 라이브러리 호출이면 polynomial 한 코드 폭발. tile-level building block 으로 합성 의 idea(§ 04).
2. AMD 의 MFMA 가 NVIDIA 의 mma.sync 와 어떻게 다른가 — instruction 레벨에서 미묘한 차이. 코드 portable 이 어려운 자리(§ 03).
3. 이 stack 위에 어떻게 올라가는 ML 프레임워크 — Meta 의 AI Template 와의 협업이 turning point. PyTorch 모델이 CK 커널로 lowering 되는 path(§ 07).

가장 미묘한 차이 — wavefront 가 64 thread. CUDA 코드의 “warp 32” 가정이 거의 모든 곳에서 깨진다. shared memory bank 수도 다르고 (32 vs 32 — 같지만 wavefront 가 64 라서 access pattern 이 다름), MFMA 의 tile shape 도 다르다(예: v_mfma_f32_16x16x16, v_mfma_f32_32x32x8).

강의의 가장 중요한 abstraction. “CK 의 코드는 일반 CUDA 보다 high-level 하다” — tile 단위의 op 들이 사용자가 보는 primitive 다.

// CK tile programming — 약식
// 1. DRAM 위에 tile_window 를 만든다 — “이 region 을 본다”
auto a_dram_window = make_tile_window(
    a_dram_tensor,
    Sequence<128, 32>{},                 // tile shape
    {iM, iK});                              // tile origin

// 2. LDS (shared) 위에 distribution 만들기
auto a_lds_block = make_static_distributed_tensor<...>();

// 3. load — DRAM tile → LDS
load_tile(a_lds_block, a_dram_window);

// 4. sweep — tile 의 각 element 위 연산
sweep_tile(a_lds_block, [&](auto idx) {
  // idx 가 tile 안 좌표
  acc(idx) += a_lds_block(idx) * b_lds_block(idx);
});

// 5. MFMA — block-level matrix multiply
block_tile_gemm_xdl_cshuffle(c_acc, a_lds_block, b_lds_block);

강의의 specific 한 예시. CNN 의 convolution 을 어떻게 CK 의 GEMM building block 위에서 짜는가. 답 — implicit GEMM.

• traditional approach (im2col) — input NHWC 를 (N·H·W, K·K·C) 의 큰 matrix 로 펼침. memory 폭발.
• implicit GEMM — input 을 그대로 두고, GEMM 의 index 계산 안에 conv 의 sliding window 를 직접 표현. “GEMM_M = N·H·W, GEMM_N = output channel, GEMM_K = K·K·input_channel” 의 mapping.

// implicit GEMM — input tensor index 가 conv 의 spatial index 로 변환
// gemm_m = n * H * W + h * W + w
// gemm_k = ky * Kx * C + kx * C + c
// (n, h+ky*dilation, w+kx*dilation, c) → input

// CK 가 이 mapping 을 tile_window 의 stride 로 표현
auto a_window_strided = make_implicit_gemm_input_window(
    input_tensor, kernel_size, strides, padding);
// 그 다음은 일반 GEMM 코드 그대로

strict 하게 — CUTLASS 의 layer 와 거의 1:1 대응. CUTLASS 의 “collective” 가 CK 의 “block tile” 에 대응, “kernel” 이 “grid” 에 대응. 사용자는 보통 L4 entry 만 쓰고, customization 이 필요할 때만 한 layer 씩 내려간다.

CK 가 라이브러리라면, AI Template (AIT) 는 그 위에 올라가는 model compiler. PyTorch 모델을 AIT 의 graph 표현으로 받아 — 각 op 를 CK 의 fused kernel 로 lowering, 한 .so 파일로 codegen. inference 시점에 PyTorch 가 그 .so 를 호출.

FlashAttention (Tri Dao) 의 origin 은 CUTLASS 위. 같은 algorithm 을 AMD 에 옮기는 작업이 CK 의 큰 첫 stress test. 강의에서 Haocong 이 직접 설명한 자리.

• online softmax + tiled attention — algorithm 자체는 hardware-agnostic. tile shape, MFMA instruction 만 hardware-specific.
• CK 위에서 거의 그대로 mapping — block tile GEMM 두 번 (Q·K^T, attn·V) + epilogue 안에 online softmax. tile programming 의 sweep 이 softmax statistics 누적.
• upstream 통합 — Tri Dao 의 official FlashAttention repo 에 ROCm path 가 합류. 같은 PyTorch SDPA API 가 H100/MI300 모두에서 동작.