Repolode
Dao-AILab

Dao-AILab/flash-attention

Transformer 모델의 GPU 학습과 추론을 위한 메모리 효율적인 exact attention 커널 FlashAttention의 공식 구현입니다.

스타
25,163
포크
3,175
열린 이슈
1,382
마지막 커밋
9시간 전

어떤 프로젝트인가요?

FlashAttention은 GPU 메모리 사용량을 줄이고 attention 연산 속도를 높이는 Transformer 모델용 exact attention 구현입니다. 일반 attention은 시퀀스 길이에 따라 메모리 사용량이 제곱으로 증가하지만, FlashAttention은 같은 결과를 메모리 사용량이 선형으로 증가하는 방식으로 계산해 GPU에서 훨씬 긴 시퀀스를 처리할 수 있습니다.

이 저장소에는 NVIDIA CUDA(Ampere, Ada, Hopper GPU) 및 AMD ROCm용 FlashAttention-2와 Hopper GPU용 FlashAttention-3, Hopper 및 Blackwell GPU용 CuTeDSL 기반 FlashAttention-4가 포함되어 있습니다. 주요 Python 함수는 표준 attention, 가변 길이 attention, 패킹된 QKV attention, KV-cache 추론용 attention을 지원하며, causal mask, sliding window attention, ALiBi, MQA/GQA, dropout, softcapping도 지원합니다.

패키지는 pip으로 설치하며 PyTorch와 CUDA 또는 ROCm 툴킷이 필요합니다. 주로 지원하는 플랫폼은 Linux입니다.

핵심 기능

  • 01시퀀스 길이에 비례해 메모리 사용량이 증가하는 방식으로 exact attention을 계산합니다.
  • 02causal masking, sliding window local attention, ALiBi를 지원합니다.
  • 03Q보다 head 수가 적은 KV 텐서를 통한 MQA와 GQA를 제공합니다.
  • 04증분 디코딩을 위한 paged KV cache와 제자리 KV cache 업데이트를 지원합니다.
  • 05NVIDIA GPU용 CUDA 빌드와 AMD GPU용 Triton 및 CK 백엔드를 제공합니다.

활용 사례

  • NVIDIA A100 또는 H100 GPU에서 긴 시퀀스로 대규모 Transformer 언어 모델을 학습합니다.
  • 증분 디코딩과 paged KV cache를 사용해 LLM 추론을 서빙합니다.
  • ROCm 백엔드를 통해 AMD Instinct GPU에서 attention을 실행합니다.

빠른 시작

  1. 1

    빌드 필수 패키지인 packaging, psutil, ninja를 설치합니다

  2. 2

    PyPI에서 FlashAttention을 설치합니다

    pip install flash-attn --no-build-isolation

사용 전 알아둘 점

  • 공식 지원은 Linux에만 제공됩니다. Windows 컴파일은 충분히 테스트되지 않았습니다.
  • CUDA 또는 ROCm 툴킷과 지원되는 GPU가 필요합니다. CUDA의 경우 Ampere, Ada, Hopper 등이 해당됩니다.
  • ninja를 사용하지 않고 병렬 작업 수를 제한하지 않으면 소스에서 컴파일하는 데 시간이 오래 걸리고 메모리를 많이 사용할 수 있습니다.

토픽

자주 묻는 질문

flash-attention은(는) 무엇인가요?
Transformer 모델의 GPU 학습과 추론을 위한 메모리 효율적인 exact attention 커널 FlashAttention의 공식 구현입니다. FlashAttention은 GPU 메모리 사용량을 줄이고 attention 연산 속도를 높이는 Transformer 모델용 exact attention 구현입니다. 일반 attention은 시퀀스 길이에 따라 메모리 사용량이 제곱으로 증가하지만, FlashAttention은 같은 결과를 메모리 사용량이 선형으로 증가하는 방식으로 계산해 GPU에서 훨씬 긴 시퀀스를 처리할 수 있습니다.
flash-attention은(는) 무료로 사용할 수 있나요? 라이선스는 무엇인가요?
flash-attention은(는) BSD-3-Clause 라이선스의 오픈소스입니다. 상업용 제품에 사용하기 전에 라이선스 조건을 확인하세요.
flash-attention은(는) 아직 유지보수되고 있나요?
활발히 유지보수 중. flash-attention의 마지막 커밋은 2026년 10월 20일입니다.
flash-attention은(는) 어디에 사용할 수 있나요?
NVIDIA A100 또는 H100 GPU에서 긴 시퀀스로 대규모 Transformer 언어 모델을 학습합니다. 증분 디코딩과 paged KV cache를 사용해 LLM 추론을 서빙합니다. ROCm 백엔드를 통해 AMD Instinct GPU에서 attention을 실행합니다.
flash-attention는 어떻게 설치하나요?
빌드 필수 패키지인 packaging, psutil, ninja를 설치합니다 PyPI에서 FlashAttention을 설치합니다: pip install flash-attn --no-build-isolation
flash-attention은(는) 어디서 받을 수 있나요?
소스 코드와 설치 방법은 GitHub에 있습니다: https://github.com/Dao-AILab/flash-attention

요약은 프로젝트 README를 바탕으로 AI가 작성했습니다. GitHub 통계는 매일 갱신됩니다.

관련 프로젝트