[triton] Triton AMD 커널 최적화: TDM 로드 파이프라이닝 개선을 통한 성능 향상

PR 링크: triton-lang/triton#10057 상태: Merged | 변경: +None / -None

들어가며

고성능 컴퓨팅(HPC) 환경에서 행렬 곱셈(GEMM) 연산은 전체 워크로드의 성능을 결정짓는 핵심 요소입니다. 특히 AMD의 최신 아키텍처인 gfx1250을 타겟팅하는 Triton 커널에서, 데이터 로드와 연산 간의 파이프라인 효율은 성능의 병목 지점이 되곤 합니다. 이번에 살펴볼 PR은 gemm_tdm_pipelined_single_warp_per_simd_schedule_kernel 내에서 TDM(Triton Data Mover) 로드 명령의 위치를 최적화하여, 메모리 레이턴시를 더 효과적으로 숨김으로써 성능을 개선한 사례입니다.

코드 분석

로드 시점의 재배치 (Pipeline Re-scheduling)

기존 코드에서는 루프 내부의 SubIteration1 단계에서 TDM 로드를 수행하고 있었습니다. 이는 루프 전체의 관점에서 볼 때, 데이터가 준비되기까지의 대기 시간을 충분히 활용하지 못하는 구조였습니다.

Before

# SubIteration1
# TDM load for next tile
producer = issue_loads(producer, a_desc, b_desc, 0, 0, a_buffer, b_buffer, BLOCK_K, NUM_BUFFERS, TRANSPOSE_B, pred=pred)

After

# SubIteration3 하단으로 이동
ttgl.amd.gfx1250.tdm.async_wait((NUM_BUFFERS - 2) * 2)
# TDM load for next tile
pred = (i + 1) - epilogue_lb
pred = (pred >> 31) & 1
producer = issue_loads(producer, a_desc, b_desc, 0, 0, a_buffer, b_buffer, BLOCK_K, NUM_BUFFERS, TRANSPOSE_B, pred=pred)

핵심 변경 사항은 issue_loads 호출을 루프의 하단부인 async_wait 직후로 옮긴 것입니다. 이를 통해 이전 루프 반복에서 발생한 데이터 로드 대기(async_wait)가 완료된 직후에 즉시 다음 데이터를 요청하게 됩니다. 결과적으로 루프 반복의 3/4 지점에서 로드하던 것을 전체 루프 반복만큼의 레이턴시를 숨길 수 있는 위치로 이동시킨 것입니다.

또한, 루프 진입 전 초기 로드(Prologue)를 명시적으로 수행하도록 코드를 추가하여 파이프라인이 비어있는 상태를 최소화했습니다.

왜 이게 좋은가

1. 메모리 레이턴시 은닉 (Latency Hiding)

GPU 연산에서 가장 큰 비용은 메모리에서 데이터를 가져오는 시간입니다. issue_loads를 async_wait 직후로 옮김으로써, 연산 유닛이 다음 데이터를 기다리는 유휴 시간(Idle time)을 획기적으로 줄였습니다. 기존에는 루프 반복의 일부만 레이턴시 은닉에 활용했다면, 이제는 전체 루프 반복 주기를 활용하여 메모리 대역폭을 더 효율적으로 사용합니다.

2. 파이프라인 효율성 증대

이러한 구조적 변경은 하드웨어의 TDM 엔진이 연산 유닛과 병렬로 동작하는 시간을 극대화합니다. 리뷰어 guacamoleo가 언급했듯이, 이 작은 변경만으로도 커널 성능이 유의미하게 향상되었으며, 이는 복잡한 스케줄링 로직을 수동으로 튜닝하는 것이 GPU 커널 최적화에서 얼마나 중요한지 보여줍니다.

교훈

• 데이터 의존성 분석: 메모리 로드와 연산 사이의 의존성을 파악하고, async_wait와 로드 명령 사이의 간격을 최대화하는 것이 중요합니다.
• Prologue/Epilogue 처리: 루프 최적화 시 루프 시작 전(Prologue)과 끝난 후(Epilogue)의 데이터 흐름을 명확히 제어해야 파이프라인 스톨을 방지할 수 있습니다.

결론

이번 최적화는 하드웨어 아키텍처의 특성을 깊이 이해하고, 데이터 흐름을 정교하게 제어함으로써 얻어낸 성과입니다. Triton과 같은 도구를 사용할 때도 하드웨어 레벨의 파이프라인 동작을 고려한 스케줄링이 성능의 차이를 만든다는 점을 다시 한번 확인하게 됩니다.

참고 자료

• https://github.com/triton-lang/triton

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