[cutlass] NVIDIA Hopper에서 FP8 GEMM + GELU 퓨전 최적화: CuTe DSL 활용

PR 링크: NVIDIA/cutlass#3378 상태: Merged | 변경: +1937 / -0

들어가며

딥러닝 모델의 핵심 연산인 GEMM(General Matrix Multiply)은 대부분의 신경망 레이어에서 사용됩니다. 특히 트랜스포머(Transformer)와 같은 대규모 모델에서는 GEMM 연산 후 GELU(Gaussian Error Linear Unit)와 같은 활성화 함수가 뒤따르는 경우가 많습니다. 이러한 연산들을 개별적인 CUDA 커널로 실행하게 되면, GEMM의 결과가 글로벌 메모리에 쓰여지고, 다시 GELU 커널이 이 데이터를 읽어와 처리하는 과정에서 불필요한 메모리 접근 오버헤드(memory bound)가 발생하여 전체 성능을 저하시킵니다.

이러한 문제를 해결하기 위해 NVIDIA/cutlass 레포지토리의 이번 PR은 Hopper 아키텍처에 특화된 FP8 GEMM과 GELU 활성화 함수를 하나의 커널로 퓨전(fusion)하는 최적화 기법을 도입합니다. 이 PR의 핵심은 CuTe DSL(Domain Specific Language)을 활용하여 epilogue_op를 컴파일 타임에 특수화함으로써, 런타임 오버헤드 없이 다양한 GELU 구현을 유연하게 지원하고 메모리 접근을 최소화하는 데 있습니다.

코드 분석: 무엇이 어떻게 개선되었나

이번 PR은 examples/python/CuTeDSL/cute/hopper/kernel/dense_gemm/dense_gemm_fp8_gelu_persistent.py라는 새로운 파일을 추가하여 Hopper 아키텍처를 위한 고성능 FP8 GEMM + GELU 퓨전 커널을 구현합니다. 이 커널은 cutlass.cute 모듈을 기반으로 CuTe DSL을 사용하여 작성되었으며, 다음과 같은 주요 특징들을 가집니다.

1. 새로운 커널 파일 dense_gemm_fp8_gelu_persistent.py

이 파일은 Hopper 아키텍처의 최신 기능을 활용하여 FP8 GEMM과 GELU를 퓨전한 커널을 정의합니다. PR 설명에 따르면, 이 커널은 dense_gemm_persistent.py에서 파생되었으며, 메인 루프, 멀티캐스트, K-블록, 퍼시스턴트 스케줄링 동작을 유지합니다. 특히, epilogue_op를 cutlass.Constexpr로 전달하고 cute.compile을 통해 특수화하는 방식은 런타임 분기 없이 다양한 활성화 함수를 지원하는 핵심적인 개선점입니다.

Before (개념적):

# GEMM 커널 실행
output_gemm = gemm_kernel(A, B)
# GELU 커널 실행 (중간 결과를 메모리에서 로드)
output_gelu = gelu_kernel(output_gemm)

After (새로운 커널의 주요 특징):

# examples/python/CuTeDSL/cute/hopper/kernel/dense_gemm/dense_gemm_fp8_gelu_persistent.py 파일 상단 주석

"""
A high-performance batched FP8 dense GEMM + GELU fusion
(C = GELU(scale_a * scale_b * A * B)) example for the NVIDIA Hopper architecture
using CuTe DSL.
...
The fused epilogue accepts a compile-time elementwise operation and provides
two GELU implementations:
- ``erf`` evaluates ``0.5 * x * (1 + erf(x / sqrt(2)))``.
- ``poly11`` approximates the standard-normal CDF with a clipped degree-11
  polynomial and computes ``x * CDF_approx(x)``.
Both implementations evaluate the scale and activation in the accumulator
type before one conversion to the output type. The validated configuration
uses FP8 A/B, Float32 accumulation, and Float16 or BFloat16 output.
...
This GEMM kernel supports the following features:
    - Utilizes Tensor Memory Access (TMA) for efficient memory operations
    - Utilizes Hopper's WGMMA for matrix multiply-accumulate (MMA) operations
    - Implements TMA multicast with cluster to reduce L2 memory traffic
    - Support persistent tile scheduling to better overlap memory load/store with MMA between tiles
    - Support warp specialization to avoid explicit pipelining between mainloop load and MMA
    - Applies per-tensor scalar scale_a and scale_b factors in the epilogue
    - Accepts a compile-time elementwise epilogue operation
    - Provides erf-GELU and a clipped degree-11 CDF polynomial approximation
    - Views each epilogue fragment directly in the accumulator registers
"""

이 주석은 새로운 커널이 Hopper의 TMA, WGMMA, Persistent Scheduling, Warp Specialization 등 최신 기능을 활용하여 FP8 GEMM과 GELU를 퓨전하며, 특히 컴파일 타임에 epilogue_op를 특수화하여 erf-GELU와 poly11-GELU를 지원함을 명확히 보여줍니다. 스케일링과 활성화 연산은 어큐뮬레이터 레지스터에서 직접 수행되어 중간 결과를 메모리에 쓸 필요가 없습니다.

2. GELU 구현 (erf_gelu 및 poly11_gelu)

이 PR은 두 가지 GELU 구현을 제공합니다. 이 함수들은 epilogue_op로 전달되어 컴파일 타임에 커널 내에서 직접 연산됩니다.

erf_gelu (정확한 구현):

# examples/python/CuTeDSL/cute/hopper/kernel/dense_gemm/dense_gemm_fp8_gelu_persistent.py

def erf_gelu(acc_vec):
    """Evaluate erf-form GELU in the accumulator type."""
    half = cute.full_like(acc_vec, 0.5)
    one = cute.full_like(acc_vec, 1.0)
    inv_sqrt2 = cute.full_like(acc_vec, 0.7071067811865476)
    return half * acc_vec * (one + cute.erf(acc_vec * inv_sqrt2))

이 함수는 표준 erf 함수를 사용하여 GELU를 계산합니다. cute.erf는 CuTe DSL 환경에서 erf 연산을 수행하며, 어큐뮬레이터 타입(Float32)으로 직접 연산하여 정밀도를 유지합니다.

poly11_gelu (다항식 근사 구현):

# examples/python/CuTeDSL/cute/hopper/kernel/dense_gemm/dense_gemm_fp8_gelu_persistent.py

def poly11_gelu(acc_vec):
    """Evaluate the clipped degree-11 CDF approximation to erf-GELU."""
    lower = cute.full_like(acc_vec, -3.5)
    upper = cute.full_like(acc_vec, 3.5)
    clipped = cute.where(acc_vec < lower, lower, acc_vec)
    clipped = cute.where(clipped > upper, upper, clipped)

    half = cute.full_like(acc_vec, 0.5)
    zero = cute.full_like(acc_vec, 0.0)
    one = cute.full_like(acc_vec, 1.0)
    # Phi(z) ~= 0.5 + z * (c1 + z^2 * (c3 + ... + z^8*c11)).
    c1 = cute.full_like(acc_vec, 0.39639100184010506)
    c3 = cute.full_like(acc_vec, -0.06247543670746915)
    c5 = cute.full_like(acc_vec, 0.0077960139440769235)
    c7 = cute.full_like(acc_vec, -0.000606554913963472)
    c9 = cute.full_like(acc_vec, 2.5871031157700182e-05)
    c11 = cute.full_like(acc_vec, -4.5557832301351553e-07)
    clipped_sq = clipped * clipped
    cdf = half + clipped * (
        c1
        + clipped_sq
        * (
            c3
            + clipped_sq
            * (c5 + clipped_sq * (c7 + clipped_sq * (c9 + clipped_sq * c11)))
        )
    )
    # Guard against small polynomial overshoot before forming GELU.
    cdf = cute.where(cdf < zero, zero, cdf)
    cdf = cute.where(cdf > one, one, cdf)
    return acc_vec * cdf

이 함수는 erf 함수 대신 11차 다항식 근사를 사용하여 GELU를 계산합니다. 다항식 근사는 일반적으로 erf 함수보다 계산 비용이 낮아 더 빠른 성능을 제공할 수 있습니다. 입력 값을 클리핑하여 수치적 안정성을 확보하고, cdf 값도 0과 1 사이로 제한하여 정확성을 높였습니다.

3. 컴파일 타임 에필로그 특수화

리뷰어 jackyangNJ의

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://www.nvidia.com/en-us/data-center/hopper-architecture/
• https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/#FP8

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.