RoPE(Rotary Position Embedding)을 Triton으로 작성한 프로젝트입니다.
먼저 CUDA가 설치되어 있어야 합니다. 그 다음 필요한 라이브러리를 아래 명령어로 설치합니다.
pip install -r requirements.txt
src/rope_triton.py에 있는 apply_rotary_pos_emb_triton 함수를 사용하면 됩니다.
torch와 triton이 같은 아웃풋을 내는지 확인하기 위해서 아래 스크립트 실행합니다.
PYTHONPATH=. pytest tests/test_rope.py
CUDA 버젼과 Triton 버젼의 성능 비교를 하기 위해서 아래 스크립트 실행합니다.
PYTHONPATH=. python3 benchmarks/benchmark_rope.py
