FlashMLA

FlashMLA is an attention library for developers building DeepSeek-V4.1 inference systems on their own hardware. It requires NVIDIA SM100 or SM103 GPUs, or Huawei Ascend 950 NPUs.

• Sparse attention: DeepSeek Sparse Attention kernels handle both prompt processing and token generation, computing attention over selected tokens.
• Quantized caches: Decoding supports FP8 and FP4 key-value caches to reduce the memory used by cached attention data.
• Fused computation: A CUDA-only kernel combines attention, rotary position operations and FP8 conversion to reduce the overhead of separate operations.
• Dense attention: CUDA kernels provide dense multi-head attention for forward and backward computation.

The NVIDIA backend uses CUDA and PyTorch. The Huawei backend uses CANN, torch_npu and PyTorch, and supports sparse prompt processing and decoding. The fused and dense attention kernels aren't available on Ascend.

FlashMLA is open source under MIT. It doesn't support NVIDIA Hopper GPUs or earlier DeepSeek models, including V3, V3.2 and V4.0.