AMD GPU에서도 동작하는 FlashAttention 구현체, Aule-Attention

출처: GitHub - AuleTechnologies/Aule-Attention

• AMD GPU에서 쓸 수 있는 FlashAttention 구현체 aule-attention이 공개됨. 설치 시 별도 컴파일 없이 바로 쓸 수 있는 게 핵심임

• 백엔드를 세 가지로 자동 선택함:

  • Triton: AMD ROCm + NVIDIA CUDA (학습/추론 모두 지원)
  • Vulkan: Intel, Apple Silicon, AMD 소비자용 GPU 등 Vulkan 지원 기기 (추론 전용)
  • CPU: GPU 없으면 NumPy로 폴백

• 설치는 간단함. AMD 데이터센터 GPU(MI300X, MI250, MI100) 기준:

  pip install torch --index-url https://download.pytorch.org/whl/rocm6.2
  pip install aule-attention

  NVIDIA는 그냥 pip install torch && pip install aule-attention

• Llama 2, Mistral, Qwen 같은 LLM에서 쓰는 Grouped Query Attention 지원함. heads_q가 heads_kv의 배수면 됨

• API는 PyTorch 친화적:

  • query/key/value: [batch, heads, seq_len, head_dim] 형태의 텐서
  • causal: 기본값 True (미래 토큰 마스킹)
  • scale: 기본값 1/sqrt(head_dim)
  • Triton 백엔드는 head_dim <= 128, Vulkan은 <= 64 제약 있음
  • float16, bfloat16, float32 지원

• 메모리 효율: 기존 어텐션은 O(N²), aule-attention은 O(N). Attention 행렬 자체를 메모리에 올리지 않아서 긴 시퀀스 처리가 가능해짐

• 수치 정확도는 PyTorch scaled_dot_product_attention 대비:

  • float32: 최대 절대 오차 < 1e-6, 상대 오차 < 1e-5
  • float16: 최대 절대 오차 < 1e-3, 상대 오차 < 1e-2

• 지원 GPU 현황:

  • Triton(학습+추론): AMD Instinct MI300X/MI250/MI100 (테스트 완료), NVIDIA H100/A100/RTX 4090 등
  • Vulkan(추론): AMD RDNA3/RDNA2, Intel Arc/내장 12~14세대, Apple M1/M2/M3(MoltenVK 경유), NVIDIA Vulkan 1.2+ 기기

• 보안 정책이 눈에 띔. git 레포에 빌드된 바이너리를 일절 포함하지 않음. 공급망 공격 방지 목적이고, 소스 빌드 또는 GitHub Releases에서 공식 배포본(체크섬 포함)을 받도록 유도함

• 내부는 Zig 0.13.0으로 작성됐고, Vulkan SDK + Python 3.8+ 필요. FlashAttention(Tri Dao), Triton(OpenAI), Vulkan(Khronos), Zig(Andrew Kelley) 등 오픈소스 생태계 위에 구축됨

• 라이선스는 MIT, 2025 Aule Technologies