66RING/tiny-flash-attention

해결하는 문제

FlashAttention은 현대 기계 학습에서 사용되는 빠르고 메모리 효율적인 어텐션 알고리즘이지만, 원래 구현은 복잡합니다. 이 프로젝트는 이를 여러 프로그래밍 언어로 작고 읽기 쉬운 버전으로 분해하여 학습자가 블랙박스로 취급하는 대신 작동 방식을 이해할 수 있도록 합니다.

작동 방식

이 저장소는 FlashAttention의 작은 구현을 제공합니다:

• 순수 Python(naive) – 알고리즘을 단계별로 따라갈 수 있는 가장 간단한 버전.
• Triton – GPU 지향적인 Python 유사 언어로, 병렬 하드웨어에서 동일한 아이디어를 표현하는 방법을 보여줍니다.
• C – 핵심 어텐션 계산의 naive C 구현.
• CUDA – 독립형 CUDA 커널과 Python 바인딩으로 래핑된 버전.
• CUTLASS – CUTLASS 라이브러리를 사용한 CUDA 구현으로, 영어와 중국어로 된 상세한 튜토리얼이 포함되어 있습니다.
• Rust – 나열되어 있지만 아직 완성되지 않았습니다(빈 체크박스로 표시).

모든 코드는 성능보다 명확성에 초점을 맞춰 최소한으로 유지됩니다. CUTLASS 버전에는 저자의 환경(CUTLASS v3.4, PyTorch 1.14, CUDA 12.4)에서 구현을 설명하는 단계별 튜토리얼이 포함되어 있습니다.

대상 독자

FlashAttention이 다양한 프로그래밍 환경에서 어떻게 구현되는지 배우고자 하는 개발자와 학생, 특히 GPU 프로그래밍, CUDA 또는 성능 지향 어텐션 메커니즘을 공부하는 사람들.

주요 특징

• 여러 언어를 다룹니다: Python, Triton, C, CUDA, CUTLASS 및 Rust(작업 중).
• 각 버전은 의도적으로 작고 교육적이며, 프로덕션용으로 최적화되지 않았습니다.
• 독립형 CUDA 코드와 Python 바인딩 CUDA 버전을 포함합니다.
• 영어와 중국어로 된 상세한 CUTLASS/CuTe 튜토리얼을 제공합니다.
• 프로젝트는 아직 진행 중이며, 더 많은 버전이 추가될 수 있습니다.