FlashMLA
Optimized attention kernels for DeepSeek model inference on GPUs and NPUs
What it is
DeepSeek's library of Multi-head Latent Attention kernels, used to serve DeepSeek-V4.1 on NVIDIA Blackwell GPUs and Huawei Ascend 950 NPUs. It offers token-level sparse attention for both the prefill and decode stages, using KV caches stored in FP8 or FP4, along with a fused norm-RoPE-attention-RoPE-cast kernel and dense attention kernels for forward and backward passes. Engineers building inference engines or training stacks call it from PyTorch to speed up attention on supported hardware.