DeepSeek-V3 모델을 위한 고속 어텐션 커널

FlashMLA는 FP8 KV 캐시를 사용하여 메모리 대역폭(Memory Bandwidth)을 절감하고, bfloat16으로 연산을 수행하여 정확도를 유지한다. 특히, 토큰 레벨의 Sparse 어텐션을 지원하며, H800 SXM5 GPU에서 410 TFLOPS, B200 GPU에서 350 TFlops의 성능을 달성한다. 이는 계산 집약적(Compute-bound) 작업 부하에서 높은 성능을 제공하며, FP8 형식의 장점을 활용하여 메모리 사용량(Memory Usage)을 최적화한다.

FlashMLA는 SM90 및 SM100 아키텍처를 지원하며, NVIDIA H800, B200 GPU에서 최적화된 성능을 제공한다. 또한, MetaX, Moore Threads, Hygon DCU, Intellifusion, Iluvatar Corex, AMD Instinct 등 다양한 GPU 플랫폼에 대한 지원을 제공한다. 이는 광범위한 하드웨어 환경(Hardware Environment)에서 FlashMLA를 사용할 수 있도록 하여, 모델 배포(Model Deployment)의 유연성을 높인다.

Sparse 어텐션은 `indices` 텐서를 사용하여 토큰 레벨의 희소성을 구현한다. 이 텐서는 각 토큰이 위치한 페이지 블록의 인덱스와 오프셋 정보를 포함하며, 불필요한 계산을 줄여(Reduce Unnecessary Computation) 성능을 향상시킨다. `indices` 텐서의 유효하지 않은 항목(-1 설정)을 통해 특정 토큰에 대한 어텐션 계산을 비활성화할 수 있다. 이는 메모리 접근 패턴(Memory Access Pattern)을 최적화하고, 계산 효율성(Computational Efficiency)을 높이는 데 기여한다.

FlashMLA는 다양한 벤치마크를 통해 성능을 측정하며, H800 SXM5 GPU에서 최대 660 TFLOPS, B200 GPU에서 1450 TFLOPS의 성능을 달성했다. 특히, Dense MLA 디코딩 커널은 메모리 대역폭 제한 환경에서 최대 3000 GB/s, 계산 제한 환경에서 660 TFLOPS를 기록했다. 이러한 벤치마크 결과는 FlashMLA가 고성능 컴퓨팅(High-Performance Computing) 환경에서 효율적인 어텐션 연산(Efficient Attention Computation)을 수행함을 보여준다.