[triton] Triton AMD 커널 최적화: 루프 언롤링(Loop Unrolling)을 통한 성능 향상

PR 링크: triton-lang/triton#8998 상태: Merged | 변경: +55 / -18

들어가며

고성능 GPU 커널 개발에서 루프(Loop)는 성능 병목의 핵심 지점입니다. 특히 AMD GPU 환경에서 동작하는 Triton 기반의 FlashAttention 커널은 복잡한 파이프라이닝과 레지스터 관리가 필수적입니다. 이번 PR은 attn_fwd_pipelined_kernel의 메인 루프를 unroll_factor=2로 언롤링하여, 반복적인 iter_id 계산과 레지스터 회전(rotating registers)에 필요한 산술 연산 오버헤드를 제거하는 최적화를 수행했습니다.

코드 분석

1. 메인 루프 언롤링 (Loop Unrolling)

기존 코드는 매 반복마다 iter_id를 업데이트하고 모듈로 연산을 통해 버퍼 인덱스를 계산했습니다. 이를 2배 언롤링함으로써 연산 횟수를 절반으로 줄였습니다.

Before:

for block_id in range(block_min, block_max, BLOCK_N):
    # ... (생략)
    v = pgm.tdm_shared_load_v(iter_id % NUM_BUFFERS, wait_count=2)
    pgm.tdm_load_global_to_shared_k([t_3, 0], (iter_id + 1) % NUM_BUFFERS)
    # ... (생략)
    iter_id += 1

After:

for block_id in range(block_min, block_max, 2 * BLOCK_N):
    # 1/2 of unrolled loop (iter_id 고정)
    v = pgm.tdm_shared_load_v(0, wait_count=2)
    pgm.tdm_load_global_to_shared_k([t_3, 0], 1)
    # ...
    # 2/2 of unrolled loop (iter_id 고정)
    v = pgm.tdm_shared_load_v(1, wait_count=2)
    pgm.tdm_load_global_to_shared_k([t_3, 0], 0)

이 변경을 통해 컴파일러가 레지스터 할당을 더 효율적으로 수행할 수 있게 되었으며, 매 루프마다 발생하는 iter_id 증가 및 모듈로 연산 비용이 제거되었습니다.

2. 테스트 환경 개선

테스트 코드에서도 run_attention 함수를 분리하여 유연성을 높였습니다. 특히 waves_per_eu=1 옵션을 추가하여 AMD GPU의 스케줄링 효율을 최적화했습니다.

왜 이게 좋은가

1. 연산 오버헤드 감소: 루프 제어 변수(iter_id)와 인덱스 계산을 위한 산술 연산이 절반으로 줄어들어, GPU 사이클을 실제 데이터 연산(GEMM 등)에 더 집중할 수 있게 합니다.
2. 레지스터 회전 최적화: unroll_factor=2를 통해 컴파일러가 고정된 버퍼 인덱스를 사용할 수 있게 되어, 레지스터 간 데이터 이동(rotating registers) 시 발생하는 불필요한 의존성을 제거했습니다.
3. 일반적 교훈: GPU 커널 최적화 시 루프 언롤링은 단순히 코드 길이를 늘리는 것이 아니라, 컴파일러가 레지스터를 정적으로 할당할 수 있는 범위를 넓혀주어 성능 향상에 큰 기여를 합니다. 특히 파이프라이닝이 적용된 커널에서는 이러한 최적화가 필수적입니다.

결론

이번 최적화는 AMD GPU 아키텍처의 특성을 고려하여 루프 구조를 개선한 좋은 사례입니다. 하드웨어의 레지스터 활용도를 극대화하기 위해 루프를 언롤링하고, 이를 통해 제어 로직의 복잡도를 낮춘 점이 핵심입니다.

참고 자료

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

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