flagos-ai/FlagAttention

A collection of memory efficient attention operators implemented in the Triton language.

FlagAttention – 빠르고 메모리 효율적인 어텐션 커널

소개 – FlagAttention은 Triton(고성능 GPU 커널 작성을 위한 언어)으로 작성된 사용자 지정 어텐션 연산자 모음을 제공하는 Python 라이브러리입니다. 이 커널들은 널리 사용되는 FlashAttention 알고리즘과 동일한 기능을 구현하며, 희소, 페이징 캐시, 조각별, 양자화, 재귀적 어텐션과 같은 많은 특수 변형을 추가합니다. 이들은 다양한 가속기(NVIDIA GPU, AMD, Ascend 등)에서 대규모 언어 모델이나 기타 Transformer 스타일 네트워크의 학습 및 추론을 위해 설계되었습니다.

주요 연산자 제품군

제품군	일반적인 사용 사례	API 진입점	그래디언트 지원	성숙도
FlashAttention	표준 밀집 어텐션 (scaled-dot-product)	`flag_attn.flash_attention`	순방향 + 역방향	안정
Piecewise Attention	토큰 거리에 따른 두 Q/K 쌍 간 전환 (예: NLPE)	`flag_attn.piecewise_attention`	순방향 + 역방향	안정
Split-KV FlashAttention	긴 KV 디코딩, 낮은 쿼리 병렬성	`flag_attn.flash_attention_split_kv`	순방향 전용	알파
Paged Attention	페이징된 KV 캐시를 통한 단일 토큰 쿼리 (디코더 전용 모델에서 일반적)	`flag_attn.paged_attention`	순방향 전용	알파
MiniMax M3 Sparse Attention	128 토큰 페이지를 사용한 블록 희소 프리필/디코딩	`flag_attn.minimax_m3_*`	추론 전용	알파
Chunk GLA / GDN2 / KDA	긴 시퀀스를 위한 재귀적 게이트 선형 또는 델타 규칙 어텐션	`flag_attn.chunk_*`	순방향 (+ GLA는 역방향)	알파
SageAttention	블록별 INT8 양자화 Q/K (선택적 마스크 포함)	`flag_attn.sage_attention.forward`	순방향 전용	알파
Parallel NSA	네이티브 희소 어텐션 + 압축	`flag_attn.parallel_nsa.parallel_nsa`	순방향 + 역방향	알파

설치 (빠른 가이드)

# Clone the repo
git clone https://github.com/flagos-ai/FlagAttention.git
cd FlagAttention

# Install the editable package with Triton and test extras
pip install -e "[triton,test]"

가속기에 호환되는 PyTorch 빌드(예: torch==2.x+cu118)가 이미 설치되어 있어야 합니다. 올바른 빌드는 GPU 공급업체마다 다르므로 이 패키지는 특정 PyTorch 휠에 의존하지 않습니다.

최소 사용 예제 – FlashAttention

import torch
from flag_attn import flash_attention

B, Hq, Hkv, M, N, D = 2, 16, 4, 2048, 4096, 128
q = torch.randn(B, Hq, M, D, device='cuda', dtype=torch.float16, requires_grad=True)
k = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)
v = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)

out = flash_attention(q, k, v, causal=True)
out.sum().backward()   # gradient flow works

함수 시그니처는 원래의 FlashAttention API를 반영하며, 드롭아웃, 로그 정규화 항목 반환, 총 어텐션 질량 및 Philox RNG 상태를 위한 선택적 플래그가 추가되었습니다.

연산자 선택 가이드

시나리오	권장 연산자
단일 GPU에서의 표준 밀집 학습 또는 추론	FlashAttention (안정)
거리 임계값 이후 어텐션 행렬을 전환해야 하는 경우 (예: 위치 임베딩 트릭)	Piecewise Attention (안정)
쿼리가 적은 매우 긴 KV 캐시 (디코딩)	Split-KV FlashAttention (알파)
페이징된 KV 캐시를 유지하는 디코더 전용 모델 (예: LLM 서빙)	Paged Attention (알파)
vLLM 스타일의 토큰 희소 디코딩과 같은 블록 희소 모델	MiniMax M3 (알파)
스트리밍 또는 매우 긴 시퀀스를 위한 재귀적 선형 어텐션	Chunk GLA / GDN2 / KDA (알파)
메모리 대역폭을 절약하기 위해 블록별 INT8 양자화를 원하는 경우	SageAttention (알파)
선택적 압축이 포함된 네이티브 희소 어텐션이 필요한 경우	Parallel NSA (알파)

성숙도 및 테스트

저장소는 conf/operators.yaml에서 연산자 성숙도를 추적합니다. FlashAttention과 Piecewise Attention만 안정으로 표시되며, 나머지는 모두 알파로 API 시그니처가 변경될 수 있습니다.
전체 pytest 제품군(pytest tests)은 FP32 PyTorch 참조를 기준으로 수치 정확도를 검증합니다. 정확도 허용 오차는 FP32 오차의 약 2배이며, 작은 절대 마진이 있습니다.
벤치마크는 benchmark/ 아래에 제공되며, 도우미 스크립트 tools/run_tests.py를 통하거나 직접 실행(예: python benchmark/flash_benchmark.py)할 수 있습니다. 결과는 FlagGems 플랫폼과 호환되는 JSON으로 출력됩니다.

지원되는 하드웨어 및 백엔드

기본: NVIDIA CUDA GPU (Triton 경유). 이 라이브러리에는 AMD, Ascend, Enflame, MetaX, MThreads, Intel 등을 위한 공급업체별 백엔드도 포함되어 있습니다. 백엔드 선택은 감지된 PyTorch 장치에 따라 자동으로 이루어지지만, 모든 연산자가 모든 공급업체에 대해 구현된 것은 아닙니다. 자세한 내용은 conf/operators.yaml 및 src/flag_attn/runtime/backend 디렉토리를 참조하십시오.
CPU: 패키지를 가져와 메타데이터를 검사할 수는 있지만, GPU 호환 런타임 없이는 커널을 실행할 수 없습니다.

추가 정보

문서 – README 자체에 전체 API 참조, 설치 노트 및 examples/ 아래의 예제 스크립트가 포함되어 있습니다.
테스트 및 벤치마크 – 구체적인 사용 패턴은 tests/를, 성능 수치는 benchmark/를 참조하십시오.
공급업체별 확장 – 저수준 구현은 src/flag_attn/runtime/backend/_<vendor>를 참조하십시오.

요약

FlagAttention은 연구자와 엔지니어에게 다양한 Transformer 워크로드를 위한 빠르고 메모리 절약적인 어텐션 커널을 제공하는 실제 프로덕션급 라이브러리입니다. 안정적인 FlashAttention 및 PiecewiseAttention API는 학습 루프의 기본 PyTorch 어텐션을 대체할 수 있으며, 알파 연산자는 고급 LLM 서빙 및 연구를 위해 페이징 캐시, 블록 희소 디코딩, 양자화 커널 및 재귀적 선형 어텐션과 같은 실험적 기능을 제공합니다.