deepseek-ai/FlashMLA

해결하는 문제

FlashMLA는 대규모 언어 모델(LLM)의 추론 및 prefilling 단계를 가속화하기 위해 설계된 고도로 최적화된 어텐션 커널을 제공하며, 특히 DeepSeek-V3 및 V3.2 모델을 지원합니다. 밀집(dense) 및 희소(sparse) 어텐션 패턴 모두에 대한 특화된 구현을 제공함으로써 Multi-head Latent Attention (MLA)의 계산 병목 현상을 해결합니다.

작동 원리

이 라이브러리는 다양한 GPU 아키텍처(SM90 및 SM100)와 동작 모드에 맞춤화된 다양한 어텐션 커널을 구현합니다:

• Sparse Attention Kernels: DeepSeek Sparse Attention (DSA)을 구현하여, prefill 및 decoding 단계 모두에서 특정 토큰(indices 텐서 경유)에 대해서만 어텐션을 계산할 수 있도록 합니다. decoding 커널은 bfloat16으로 계산을 수행하는 동안 메모리 오버헤드를 줄이기 위해 FP8 KV cache를 활용합니다.
• Dense Attention Kernels: prefilling 및 decoding을 위한 표준 밀집 어텐션 구현을 제공합니다.
• 하드웨어 최적화: 커널은 NVIDIA H800 및 B200 GPU에 최적화되어 있으며, 특정 아키텍처 기능을 활용하여 높은 TFlops 성능을 달성합니다.

대상 사용자

이 프로젝트는 DeepSeek 모델을 사용하는 AI 연구자 및 엔지니어 또는 하이엔드 NVIDIA GPU에서 최대의 하드웨어 활용도와 추론 속도가 필요한 Multi-head Latent Attention (MLA) 구현자를 대상으로 합니다.

주요 특징

• 높은 처리량: B200 GPU에서 밀집 MLA decoding 시 최대 660 TFlops, 희소 MLA prefill 시 최대 1450 TFlops를 달성합니다.
• FP8 KV Cache 지원: 스케일 인자를 사용하는 양자화된 FP8 형식을 사용하여 희소 decoding 중 메모리 사용량을 줄입니다.
• 토큰 수준의 희소성: prefill 및 decoding 단계에서 계산을 최적화하기 위해 선택적 토큰 어텐션을 지원합니다.
• 폭넓은 하드웨어 지원: SM90 및 SM100 아키텍처에 최적화되어 있으며, 다양한 기타 GPU 가속기에 대한 커뮤니티 포트를 제공합니다.