deepseek-ai/FlashMLA

What it solves

FlashMLA provides highly optimized attention kernels designed to accelerate the inference and prefilling stages of Large Language Models, specifically powering the DeepSeek-V3 and V3.2 models. It addresses the computational bottlenecks of Multi-head Latent Attention (MLA) by providing specialized implementations for both dense and sparse attention patterns.

How it works

The library implements a variety of attention kernels tailored for different GPU architectures (SM90 and SM100) and operational modes:

• Sparse Attention Kernels: These implement DeepSeek Sparse Attention (DSA), allowing the model to compute attention only for specific tokens (via an indices tensor) during both prefill and decoding stages. The decoding kernel utilizes an FP8 KV cache to reduce memory overhead while performing calculations in bfloat16.
• Dense Attention Kernels: These provide standard dense attention implementations for prefilling and decoding.
• Hardware Optimization: The kernels are optimized for NVIDIA H800 and B200 GPUs, leveraging specific architectural features to achieve high TFlops performance.

Who it’s for

This project is intended for AI researchers and engineers working with DeepSeek models or those implementing Multi-head Latent Attention (MLA) who need maximum hardware utilization and inference speed on high-end NVIDIA GPUs.

Highlights

• High Throughput: Achieves up to 660 TFlops for dense MLA decoding and 1450 TFlops for sparse MLA prefill on B200 GPUs.
• FP8 KV Cache Support: Reduces memory footprint during sparse decoding by using a quantized FP8 format with scale factors.
• Token-Level Sparsity: Supports selective token attention to optimize computation during the prefill and decoding stages.
• Broad Hardware Support: Optimized for SM90 and SM100 architectures with community ports for various other GPU accelerators.