FlashMLA is DeepSeek's library of optimized attention kernels designed to accelerate multi-head latent attention computations on modern GPUs. The repository contains implementations powering the DeepSeek-V3 and DeepSeek-V3.2-Exp large language models. Written primarily in C++, it provides both sparse and dense attention kernels optimized for different stages of transformer inference and training.
The library implements two main categories of kernels. Sparse attention kernels power DeepSeek Sparse Attention (DSA) with token-level sparse attention for both prefill and decoding stages, where the decoding variant supports FP8 KV cache quantization. Dense attention kernels provide standard attention computation for prefill and decoding stages. The sparse decoding kernel achieves up to 410 TFlops on H800 SXM5 GPUs and 350 TFlops on B200, while the dense decoding kernel reaches 660 TFlops in compute-bound configurations and 3000 GB/s in memory-bound configurations on H800 SXM5. For prefill operations, the sparse MLA kernel achieves 640 TFlops on H800 SXM5 and 1450 TFlops on B200, while the dense MHA kernel reaches 1460 TFlops in forward and 1000 TFlops in backward computation on B200.
The kernels support SM90 and SM100 GPU architectures with CUDA 12.8 and above, requiring PyTorch 2.0 or later. The FP8 KV cache format stores quantized key-value data in a specific structure with 512 bytes of quantized values, 16 bytes of scale factors, and 128 bytes of unquantized RoPE embeddings per token. Sparse attention is enabled through an indices tensor that specifies which tokens to attend to, allowing the kernel to skip unnecessary computations.
The repository shares contributors with related projects including sgl-project/sglang, vllm-project/vllm, and nvidia/tensorrt-llm, indicating integration within the broader ecosystem of LLM inference optimization tools.
Recent development includes a September 2025 release of sparse attention kernels coinciding with DeepSeek-V3.2 launch, an April 2025 performance update delivering 5 to 15 percent improvements for compute-bound workloads, and August 2025 additions of MHA kernels for SM100 architecture contributed by NVIDIA. The library provides detailed documentation including deep-dive technical blogs explaining kernel implementation details and FP8 sparse decoding optimization strategies.
FlashMLA supports community implementations across multiple GPU architectures through partnerships, with corresponding versions available for MetaX, Moore Threads, Hygon DCU, Intellifusion NNP, Iluvatar Corex, and AMD Instinct GPUs. The kernels are designed to be compatible with existing interfaces while providing performance improvements, allowing users to upgrade without modifying calling code. The implementation draws inspiration from FlashAttention 2 and 3 projects and NVIDIA's CUTLASS library.
