[Triton] AMD FAv3 pingpong에서 s_xxx 명령어 배치 최적화

PR 링크: triton-lang/triton#8535 상태: Merged | 변경: +143 / -31

들어가며

AMD GPU의 FlashAttention v3(FAv3) 커널은 pingpong 스케줄링을 사용한다. 두 warp가 번갈아 compute cluster(MFMA)와 memory cluster(로드/스토어)를 실행하며, 이 두 phase가 겹치도록 스케줄링하는 것이 핵심이다. 이 PR은 s_setprio, s_waitcnt, s_barrier 같은 스칼라 명령어들의 배치를 최적화하여 파이프라인 효율을 높인다.

핵심 코드 분석

Before

// scf.for
rocdl.s.setprio 0    // compute cluster에서 낮은 우선순위
tt.dot               // MFMA
rocdl.s.setprio 1    // memory cluster에서 높은 우선순위
ttg.async_wait
rocdl.sched.barrier 0
// memory ops...
rocdl.s.barrier

After

// scf.for
rocdl.s.barrier           // 이전 iteration 완료 대기
rocdl.sched.barrier 0
tt.dot                    // compute cluster
rocdl.sched.barrier 0
ttg.async_wait
rocdl.s.setprio 1         // memory cluster 시작 시 높은 우선순위
rocdl.sched.barrier 0
// memory ops...
rocdl.sched.barrier 0
rocdl.s.setprio 0         // compute cluster 전환 시 낮은 우선순위
rocdl.s.waitcnt -7937     // ds_read 완료 대기

핵심 변경점: (1) memory cluster에 높은 우선순위를 부여하여 VALU co-issue 보장, (2) s_waitcnt lgkmcnt(0)를 memory cluster 끝에 배치하여 compute cluster에서 불필요한 대기 제거, (3) s_xxx를 memory cluster에 배치하여 다른 warp의 VALU와 co-issue.

왜 이게 좋은가

1. VALU co-issue 극대화: 서로 다른 warp의 s_xxx와 v_xxx가 동시 실행 가능
2. compute cluster 무중단: 스칼라 명령어가 모두 memory cluster에 위치하여 MFMA가 중단 없이 실행
3. DS read 대기 최적화: s_waitcnt lgkmcnt(0)를 cluster 경계에 배치하여 LLVM이 추가 waitcnt를 삽입하지 않음

정리

GPU 스케줄링 최적화는 명령어 하나의 위치가 전체 커널 성능에 큰 영향을 미친다. 이 PR의 코멘트에 포함된 상세한 설계 노트는 pingpong 스케줄링의 원리를 이해하는 데 훌륭한 참고 자료다.

참고 자료

• PR #8535

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.