[triton] Triton Hopper 커널 최적화: Persistent Matmul에서 Epilogue 오버랩 제거하기

PR 링크: triton-lang/triton#9279 상태: Merged | 변경: +15 / -3

들어가며

GPU 커널 최적화의 세계에서는 때때로 직관에 반하는 결과가 나타나곤 합니다. 일반적으로 연산의 파이프라이닝(Pipelining)을 통해 Epilogue(결과 저장 단계)와 Prologue(다음 타일의 로드 단계)를 오버랩하면 성능이 향상될 것이라 기대합니다. 하지만 NVIDIA H200과 같은 최신 Hopper 아키텍처 기반의 Mixed Precision Matmul에서는 오히려 이러한 오버랩이 성능 저하를 유발할 수 있습니다. 본 글에서는 Triton 레포지토리의 최근 PR을 통해, 왜 특정 상황에서 루프를 Flatten하지 않는 것이 더 나은 성능을 내는지 분석합니다.

코드 분석

이번 최적화의 핵심은 FLATTEN_LOOPS라는 tl.constexpr 플래그를 도입하여, Hopper 아키텍처의 특정 조건에서 루프 Flattening을 제어하는 것입니다.

1. _p_matmul.py의 루프 제어 변경

기존에는 tl.range 호출 시 flatten=True가 하드코딩되어 있었습니다. 이를 파라미터화하여 외부에서 제어할 수 있도록 변경했습니다.

# Before
for block_id in tl.range(tl.program_id(0), num_blocks, NUM_SMS, flatten=True, ...):

# After
for block_id in tl.range(
    tl.program_id(0), num_blocks, NUM_SMS,
    flatten=FLATTEN_LOOPS,
    disallow_acc_multi_buffer=DISALLOW_ACC_MULTI_BUFFER,
    warp_specialize=True,
):

2. opt_flags.py에서의 조건부 적용

HopperMXScaleLayout을 사용하는 경우, 즉 Hopper 아키텍처의 특수 연산을 수행할 때는 FLATTEN_LOOPS를 False로 설정하여 오버랩을 비활성화합니다.

# After
target_kernel_kwargs=dict(
    maxnreg=maxnreg,
    # For some reason, overlapping the epilogue is slower for hopper bf16 x mxfp4
    FLATTEN_LOOPS=not is_hopper_scale,
),

왜 이게 좋은가

성능 수치

이 변경을 통해 H200 GPU에서 대역폭이 2650 GBps에서 2800 GBps로 약 150 GBps 향상되었습니다.

기술적 통찰

저자의 가설에 따르면, Epilogue와 Prologue를 오버랩하지 않음으로써 SM(Streaming Multiprocessor) 내의 다른 블록이 Tensor Core를 더 효율적으로 활용할 수 있게 됩니다. 이는 마치 하드웨어 수준의 'Ping-pong' 스케줄링과 유사한 효과를 Warp Scheduler가 스스로 수행하게 만드는 결과로 이어집니다.

일반적인 최적화 교훈은 다음과 같습니다:

1. 하드웨어 아키텍처 특수성: 범용적인 최적화 기법(오버랩)이 최신 아키텍처(Hopper)의 스케줄링 메커니즘과 충돌할 수 있습니다.
2. Occupancy의 중요성: 때로는 파이프라이닝을 포기하고 Occupancy를 높이는 것이 전체적인 처리량(Throughput) 측면에서 더 유리할 수 있습니다.
3. 실측 기반의 튜닝: 이론적인 성능 향상 모델보다 실제 하드웨어에서의 프로파일링 결과가 우선되어야 함을 보여줍니다.

결론

이번 PR은 무조건적인 파이프라이닝이 항상 정답은 아님을 시사합니다. 특히 Hopper와 같은 복잡한 스케줄링 구조를 가진 GPU에서는 커널의 실행 흐름을 유연하게 제어할 수 있는 구조를 갖추는 것이 중요합니다.

참고 자료

• https://triton-lang.org/main/index.html
• https://docs.nvidia.com/cuda/hopper-tuning-guide/index.html

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