Hierarchical Tiling

개요

Hierarchical Tiling은 GEMM(General Matrix Multiply) 연산에서 메모리 계층 구조에 맞춰 행렬을 여러 단계로 분할하는 기법이다. 단순 Square Tiling이 단일 캐시 레벨을 기준으로 타일 크기를 결정하는 반면, Hierarchical Tiling은 레지스터, Shared Memory, L1/L2 캐시, DRAM 등 각 메모리 계층의 대역폭·레이턴시·용량 특성에 최적화된 타일 크기를 설계하여 데이터 재사용을 극대화하고 병렬성을 완전히 활용한다.

GPU의 메모리 계층은 수백 GB/s 대역폭의 Global Memory(HBM)에서 수 TB/s의 Shared Memory, 그리고 수 KB의 레지스터 파일까지 수직적 구조를 갖는다. Hierarchical Tiling은 이 각 계층에서 데이터가 한 번 로드되면 해당 계층 내에서 최대한 재사용되도록 타일 크기를 조절함으로써, 계층 간 데이터 이동 비용을 최소화한다. NVIDIA CUTLASS, cuBLAS 등 고성능 GEMM 라이브러리의 핵심 최적화 기법으로, H100/A100 GPU에서 이론적 피크 성능의 90% 이상을 달성하는 데 기여한다.

핵심 개념

메모리 계층과 타일 분할

Hierarchical Tiling의 핵심은 각 메모리 계층에 맞는 타일 크기를 설계하는 것이다:

각 계층의 타일은 상위 계층의 타일을 서브 디바이드(subdivide)하여 정의된다. 예를 들어, 하나의 Threadblock(CtaTile)은 여러 개의 Warp Tile로 분할되고, 각 Warp Tile은 다시 여러 개의 MMA Atom으로 분할된다.

Threadblock Tiling (Global → Shared Memory)

Threadblock Tiling은 Global Memory에서 Shared Memory로의 효율적 데이터 로드를 담당한다:

C = A × B  (M×K) × (K×N) → (M×N)

for cta_n = 0 to N step CtaTileN:
  for cta_m = 0 to M step CtaTileM:
    for cta_k = 0 to K step CtaTileK:
      // CTA(Threadblock)가 담당하는 M×N 타일
      // CtaTileK 만큼 K축을 순회하며 데이터 로드
      Load A[cta_m:cta_m+CtaTileM, cta_k:cta_k+CtaTileK] → SMEM
      Load B[cta_k:cta_k+CtaTileK, cta_n:cta_n+CtaTileN] → SMEM
      // Shared Memory에서 MMA 연산

타일 파라미터 (CUTLASS 기준):

성능 트레이드오프:
- 큰 타일: Global Memory 로드 횟수 감소 (좋음), Register/Shared Memory 소비 증가 (나쁨)
- 작은 타일: 로드 횟수 증가 (나쁨), 메모리 사용량 감소 (좋음)

Warp Tiling (Shared Memory → Register)

Warp Tiling은 Shared Memory에서 Register로의 효율적 데이터 로드와 Tensor Core 활용을 담당한다:

// Threadblock 내부에서 각 워프가 담당하는 서브타일
for warp_n = 0 to CtaTileN step WarpTileN:
  for warp_m = 0 to CtaTileM step WarpTileM:
    for warp_k = 0 to CtaTileK step WarpTileK:
      // K축 전체를 unroll하여 연산
      Load SMEM → Register (WarpTileM × WarpTileK + WarpTileN × WarpTileK)
      MMA(Register A, Register B, Register C)

타일 파라미터:

중요 고려사항:
- Shared Memory bank conflict 없어야 함
- 워프 내 데이터 중복 로드를 최소화
- Tensor Core 사용 시 MMA instruction 크기에 맞춰야 함

Thread/Atom Tiling (Register 내 MMA)

Thread/Atom Tiling은 실제 Tensor Core/CUDA Core MMA 연산의 입출력을 관리한다:

// Warp 내부에서 MMA instruction 수준 연산
for mma_k = 0 to WarpTileK step MmaK:
  for mma_n = 0 to WarpTileN step MmaN:
    for mma_m = 0 to WarpTileM step MmaM:
      mma_instruction(d, a, b, c)  // Tensor Core 연산

타일 파라미터 (Tensor Core 기준):

레지스터 사용량 계산:

레지스터 수 ≈ (WarpTileM × WarpTileK + WarpTileN × WarpTileK + WarpTileM × WarpTileN) × bytes_per_element / 4

→ 너무 큰 WarpTile은 레지스터 부족으로 occupancy 감소

CUTLASS의 5계층 분해 구조

CUTLASS 3.x는 5계층 API로 GEMM을 분해한다:

비교/분석

타일 크기 트레이드오프

타일 크기 결정 공식

Shared Memory 기준:

Stage 크기 = (CtaTileM × CtaTileK + CtaTileN × CtaTileK) × element_size
총 Shared Memory = Stage 크기 × num_stages
→ SM당 Shared Memory 한도 (96KB~228KB) 이내어야 함

Register 기준:

누적기(accumulator) 레지스터 = CtaTileM × CtaTileN / (warp 수) / 32
→ 총 레지스터 사용량 ≤ 255 (SM당)
→ occupancy = min(255/레지스터수, SM당 최대 워프수)

일반적인 타일 크기 예시 (H100/A100)

동작 원리

Pipelining (소프트웨어 파이프라이닝)

메모리 접근 지연시간(memory latency)을 숨기기 위해 더블 버퍼링을 사용한다:

Stage 1: [Load SMEM_0 from Global] [MMA on SMEM_1] [Store results]
Stage 2: [Load SMEM_1 from Global] [MMA on SMEM_0] [Store results]
...

• Threadblock-scoped Shared Memory tiles: 2-Stage 이상의 더블 버퍼링
• 한 타일은 현재 MMA 연산 사용
• 다른 타일은 다음 반복을 위한 Global Memory 로드
• Warp-scoped Matrix Fragments: 2-Stage 레지스터 더블 버퍼링

Threadblock Rasterization

L2 Cache 활용도를 극대화하기 위해 연속적으로 런칭되는 Threadblock이 인접한 GEMM 영역을 담당하도록 매핑한다:

• 같은 SMEM 타일을 동시에 접근할 확률 증가 → L2 cache hit율 향상
• threadblock_swizzle.h에서 다양한 스케줄링 함수 제공

Parallelized Reductions

Split-K (Threadblock 간 병렬성):
- M, N이 작고 K가 클 때 유효
- K 차원을 여러 Threadblock으로 분할 → 병렬 연산 후 Reducer
- 2-커널 필요: Partitioned-K GEMM + Batched Reduction

Sliced-K (Warp 간 병렬성):
- Threadblock 내에서 K 차원을 여러 warp로 분할
- Split-K보다 적은 오버헤드
- warp 간 부분 합(partial sum)을 마지막에 reduction

Hopper Warp Specialization

Hopper 아키텍처(SM90)에서 도입된 Warp Specialization은 Producer-Consumer 분리로 비동기 파이프라인을 구현한다:

• Producer Warps: TMA(Tensor Memory Accelerator)로 Global → Shared Memory 비동기 로드
• Consumer Warps: Shared Memory에서 Tensor Core MMA 연산
• Persistent Kernel: Grid 크기를 SM 수만큼만 런칭하여 커널 런칭 오버헤드를 amortize
• Tile Scheduler: 런칭된 Threadblock에 동적으로 타일 할당

장단점

장점

• 메모리 대역폭 최적화: 각 계층의 특성에 맞는 타일 크기로 데이터 이동 비용 최소화
• 병렬성 극대화: Grid/Threadblock/Warp 수준의 계층적 병렬성 활용
• ** occupancy 조절:** 타일 크기 조정으로 레지스터/Shared Memory 사용량 제어
• 하드웨어 활용: Tensor Core, TMA, Warp Specialization 등 최신 하드웨어 기능 활용
• 범용성: FP16, BF16, FP8, INT8 등 다양한 정밀도 지원

단점

• 설계 복잡성: 다수의 타일 파라미터를 튜닝해야 함
• 하드웨어 의존: 아키텍처별 최적 타일 크기가 다름
• 커널 런칭 오버헤드: Persistent Kernel이 아닌 경우 비효율적
• 경계 타일 처리: 문제 크기가 타일 크기로 나누어 떨어지지 않으면 불필요한 연산 발생

관련 기술

• CUTLASS: NVIDIA 고성능 선형대수 CUDA C++ 템플릿 라이브러리 (Volta~Blackwell 지원)
• cuBLAS: NVIDIA 기본 선형대수 라이브러리 (내부적으로 CUTLASS 기반 커널 사용)
• cuBLASXt: Multi-GPU + CPU 하이브리드 타일링
• cuBLASDx: On-Chip GEMM (Device-Side GEMM)
• Tensor Core: NVIDIA GPU의 매트릭스 연산 가속기 (MMA instruction)
• TMA (Tensor Memory Accelerator): Hopper 아키텍처의 비동기 메모리 로드 엔진
• Software Pipelining: 메모리 레이턴시 숨기기 위한 더블 버퍼링 기법

핵심 정리

Hierarchical Tiling은 메모리 계층 구조에 맞춰 행렬을 다단계로 분할하는 GEMM 최적화 기법으로, Threadblock(Warp)/Thread/Atom 수준에서 각각 최적의 타일 크기를 설계하여 데이터 재사용을 극대화한다. NVIDIA CUTLASS의 5계층 분해 구조(Device/Kernel/Collective/TiledMMA/Atom)는 이 기법의 표준적인 구현 프레임워크로, Hopper 아키텍처의 Warp Specialization과 결합되어 이론적 피크 성능의 90% 이상을 달성한다. 타일 크기는 Shared Memory 용량, Register 수, occupancy 간의 트레이드오프를 통해 결정되며, Pipelining과 Threadblock Rasterization과 같은 기법이 결합되어 최적의 성능을 제공한다.