vipshop/ffpa-attn

해결하는 문제

FFPA(Fast and Memory-Efficient Exact Attention)는 큰 헤드 차원(D > 256, 최대 1024)을 다룰 때 표준 어텐션 메커니즘의 성능 병목 현상을 해결합니다. 표준 구현은 이러한 규모에서 메모리와 레지스터 압력으로 인해 상당한 속도 저하를 겪는 경우가 많습니다. FFPA는 PyTorch의 기본 Scaled Dot Product Attention(SDPA)보다 상당한 속도 향상을 제공하는 고도로 최적화된 커널을 제공합니다.

작동 방식

FFPA는 어텐션을 최적화하기 위해 두 가지 주요 기술을 사용합니다.

• Split-D: 이 타일링 전략은 헤드 차원(D) 축을 분할하여 SRAM 사용량을 일정하게 유지하여 메모리 병목 현상을 효과적으로 해결합니다.
• TiledMMA: 이 접근 방식은 행렬 곱셈의 N 차원을 분할하여 레지스터 사용을 최적화하고 레지스터 스필을 방지하여 기존 구현보다 훨씬 큰 헤드 차원을 커널이 처리할 수 있게 합니다.

이 라이브러리는 CUDA, Triton, CuTe-DSL을 포함한 여러 백엔드를 지원하며, 한 줄의 코드로 기존 PyTorch 워크플로우에 통합할 수 있습니다.

대상 사용자

이 프로젝트는 대규모 모델 훈련 및 추론에 종사하는 AI 연구자와 엔지니어를 대상으로 하며, 특히 큰 헤드 차원을 가진 모델이나 높은 처리량의 FP8/FP4 정밀도가 필요한 모델의 어텐션 레이어를 최적화해야 하는 사람들을 위한 것입니다.

주요 특징

• 고성능: 표준 PyTorch SDPA보다 1.5배에서 15배까지 속도 향상을 달성합니다.
• 큰 헤드 차원 지원: 정밀도 손실 없이 최대 1024까지 헤드 차원을 지원합니다.
• 양자화 지원: 더 빠른 추론을 위한 FP8 및 FP4 어텐션 커널의 실험적 지원을 포함합니다.
• 유연한 통합: torch.nn.functional.scaled_dot_product_attention을 몽키 패치하여 즉각적인 성능 향상을 제공하는 간단한 API를 제공합니다.
• 광범위한 호환성: 자기 어텐션, GQA/MQA, 교차 어텐션, 인과/마스크 어텐션 등 다양한 어텐션 유형을 지원합니다.