vAttention 구현 — microsoft/vattention 저장소

vAttention 구현 — microsoft/vattention 저장소

오픈소스 구현은 Microsoft 저장소 microsoft/vattention에 있어요. vAttention을 LLM 서빙 시스템 Sarathi-Serve(OSDI'24에 발표)에 통합한 형태로 제공됩니다.

출처: https://github.com/microsoft/vattention

저장소 구조

  • vattention — vattention 메모리 할당기의 소스 코드
  • sarathi-lean — PagedAttention과 vAttention 스타일의 메모리 관리를 모두 지원하도록 수정한 Sarathi-Serve
  • scripts — 실험 실행용 스크립트
  • nvidia-vattn-uvm-driver — 수정한 NVIDIA UVM 드라이버
  • microbenchmarks — 유용한 마이크로벤치마크 실행 스크립트

설치

vAttention과 Sarathi-Serve 설치는 conda 환경을 만들어 진행해요.

conda create -n vattn python=3.10
conda activate vattn

libtorch를 내려받아 압축 해제해야 합니다(vattention 메모리 할리기 빌드에 필요). 그런 다음 sarathi-serve와 vattention을 차례로 빌드해요.

# build sarathi-serve
cd sarathi-lean/
pip install -e . --extra-index-url https://flashinfer.ai/whl/cu121/torch2.3/
cd ../

# build vattention
cd vattention/
LIBTORCH_PATH=<path to libtorch dir> python setup.py install
cd ../

벤치마크

저장소는 벤치마크 러너를 제공해 동적/정적, 데이터셋/합성 등 다양한 워크로드를 여러 어텐션 백엔드·스케줄러로 실행할 수 있어요. 구성 노브는 default.yml에 정리돼 있고, 실제 성능 평가용 맞춤 벤치마크 스크립트도 두 개 들어 있습니다.

커스텀 UVM 드라이버

NVIDIA CUDA 드라이버는 메모리를 2MB 이상의 큰 페이지 단위로만 할당해요. vAttention을 64KB·128KB·256KB 같은 더 작은 페이지 크기로 쓰려면 기본 CUDA UVM 드라이버를 이 저장소의 커스텀 드라이버(nvidia-vattn-uvm-driver)로 교체해야 합니다.

성능

vAttention은 PagedAttention 방식보다 성능이 좋은 경우가 많아요. 논문에 따르면 FlashAttention의 비페이지드 커널을 사용할 때 decode 처리량에서 vLLM 대비 최대 1.97배, 입력 프롬프트 처리에서는 PagedAttention 변형 FlashAttention·FlashInfer 대비 각각 최대 3.92배·1.45배 빠르다고 보고됩니다.

더 알아보기