Dao-AILab/flash-attention

Fast and memory-efficient exact attention

해결하는 문제

FlashAttention은 Transformer 모델의 어텐션 메커니즘을 최적화하여 더 빠르고 메모리 효율적으로 만듭니다. 표준 어텐션 연산과 관련된 높은 메모리 및 계산 오버헤드를 해결하며, 특히 서로 다른 GPU 메모리 유형 간의 데이터 이동 횟수를 줄이는 방식(IO-awareness)을 사용합니다.

작동 원리

메모리 액세스를 최소화하는 특수 버전의 스케일드 닷 프로덕트 어텐션(scaled dot product attention)을 구현합니다. 이 프로젝트는 NVIDIA(Ampere, Ada, Hopper, Blackwell) 및 AMD(ROCm)를 포함한 다양한 GPU 아키텍처에 최적화된 여러 버전(FlashAttention, FlashAttention-2, 3, 4)을 제공합니다. 인과 마스킹(causal masking), 슬라이딩 윈도우 로컬 어텐션, 멀티 쿼리/그룹 쿼리 어텐션(MQA/GQA)과 같은 고급 기능을 지원합니다.

대상 사용자

대규모 언어 모델(LLMs)을 학습하거나 배포하는 머신러닝 엔지니어 및 연구원을 위해 설계되었으며, 어텐션 단계에서 GPU 처리량을 극대화하고 메모리 소비를 줄여야 하는 경우에 적합합니다.

주요 특징

  • 멀티 GPU 아키텍처 지원: NVIDIA Hopper (H100) 및 Blackwell (B200) GPU와 AMD ROCm 호환 GPU를 위한 최적화된 커널 제공.
  • 메모리 효율성: IO 오버헤드를 줄여 순방향 및 역방향 패스 속도를 향상.
  • 추론 최적화: KV 캐시를 지원하는 증분 디코딩을 위한 특수 함수 포함.
  • 유연한 어텐션 패턴: 인과 마스킹, ALiBi 슬로프, 슬라이딩 윈도우 어텐션 지원.

관련

  • 프로젝트
    flagos-ai/FlagAttention

    FlagAttention은 Triton 기반 라이브러리로, 다양한 가속기에서 대규모 모델의 학습 및 추론을 위한 빠르고 메모리 효율적인 어텐션 커널(FlashAttention, 조각별, 페이징, 희소, 양자화, 재귀적 등)을 제공합니다. 안정적인 밀집 연산자는 프로덕션 준비가 완료되었으며, 알파 버전은 KV 캐시 페이징, 블록 희소 디코딩, INT8 양자화와 같은 실험적 기능을 제공합니다.

  • 프로젝트
    HKUSTDial/flash-sparse-attention

    Flash Attention의 메모리 효율성과 스파스 계산을 결합하여 Transformer 모델에서 극히 긴 시퀀스를 처리할 수 있도록 하는 고성능 트레이너블 스파스 어텐션 구현.

  • 프로젝트
    lucidrains/x-transformers

    x-transformers는 구성 가능한 Transformer 모델(인코더-디코더, 디코더 전용, 인코더 전용, 비전)을 구축하기 위한 PyTorch 라이브러리로, 많은 연구 수준의 기능(flash attention, 메모리 토큰, 대체 정규화, 희소/top-k 어텐션, talking-heads, GLU 변형 등)을 갖추고 있습니다. 연구자와 실무자가 새로운 아이디어를 빠르게 프로토타이핑할 수 있도록 간결하고 확장 가능한 구현을 제공하는 것을 목표로 합니다.

  • Dispatch
    프로덕션 환경에서의 LLM 최적화: 정밀도, 어텐션 및 아키텍처

    Hugging Face는 저정밀도 양자화, 메모리 효율성을 위한 Flash Attention, 그리고 RoPE, ALiBi, MQA, GQA와 같은 아키텍처 최적화에 초점을 맞춘 효율적인 LLM 배포를 위한 핵심 기술들을 설명합니다.

  • 프로젝트
    vllm-project/vllm

    vLLM은 대규모 언어 모델의 높은 처리량과 낮은 지연 시간 서빙을 위한 오픈소스 라이브러리입니다. 효율적인 KV 캐시 처리를 위한 PagedAttention을 도입하고, 다양한 양자화 형식을 지원하며, 빠른 어텐션 커널을 통합하고, GPU, CPU 및 기타 가속기에서 실행되는 OpenAI 호환 API를 제공합니다.