[triton] AMD: LLVM 백엔드에 커스텀 스케줄러 옵션 추가로 메모리 바운드 커널 최적화

PR 링크: triton-lang/triton#8700 상태: Merged | 변경: +60 / -7

들어가며

GPU 커널의 성능은 명령어 스케줄링에 크게 좌우됩니다. 특히 메모리 바운드 Flash Attention 변형 커널처럼 elementwise 연산이 많은 경우, 기본 LLVM 스케줄러보다 더 적합한 전략이 필요합니다. 이번 PR은 AMD HIP 백엔드에 schedule_hint 옵션을 확장하여 LLVM의 iterative-ilp 스케줄러를 사용할 수 있도록 합니다.

핵심 코드 분석

Before: 단일 hint 처리

# compiler.py
if options.schedule_hint.lower() != "none":
    amd.passes.ttgpuir.insert_instruction_sched_hints(pm, options.schedule_hint)

After: 쉼표 구분 다중 hint 지원 + LLVM 스케줄러 속성

# compiler.py
if options.schedule_hint.lower() != "none":
    for hint in options.schedule_hint.split(","):
        amd.passes.ttgpuir.insert_instruction_sched_hints(pm, hint)

# LLVM IR 단계에서 스케줄러 전략 적용
if "memory-bound-attention" in options.schedule_hint.split(','):
    fns[0].add_fn_attr("amdgpu-sched-strategy", "iterative-ilp")

schedule_hint="attention,memory-bound-attention" 처럼 쉼표로 여러 힌트를 조합할 수 있으며, memory-bound-attention 힌트가 있으면 LLVM 백엔드에 iterative-ilp 스케줄러를 지정합니다. 기존의 sink-insts-to-avoid-spills 플래그는 제거되었습니다.

왜 이게 좋은가

1. Thread-safe 구현: PR #8326의 thread-safe 버전으로, 여러 커널을 동시 컴파일할 때 안전합니다.
2. 조합 가능한 힌트: 쉼표 구분으로 여러 최적화 힌트를 동시에 적용할 수 있습니다.
3. LLVM 수준 최적화: Triton IR이 아닌 LLVM IR 수준에서 스케줄링 전략을 변경하므로 기존 파이프라인에 영향을 주지 않습니다.

정리

이 PR은 컴파일러 백엔드의 스케줄링 전략을 사용자가 선택할 수 있게 하여, 특정 워크로드에 맞는 코드 생성을 가능하게 합니다. 실험적 기능이지만 메모리 바운드 커널에서 유의미한 성능 향상을 기대할 수 있습니다.

참고 자료

• triton-lang/triton#8700

---

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