FlashMLA

Overview

FlashMLA is DeepSeek's library of optimized attention kernels, powering inference of DeepSeek-V4.1 on NVIDIA GPUs and Huawei Ascend NPUs. It includes token-level sparse attention kernels for prefill and decoding with an FP8 or FP4 KV cache, a fused kernel that combines Q-norm, Q-RoPE, attention, O-RoPE, and the cast to FP8, and dense attention kernels for prefill and backward. The sparse kernels power DeepSeek Sparse Attention.

Key Features

• Token-level sparse attention kernels for prefill and decoding with FP8 or FP4 KV cache
• Fused kernel combining Q-norm, Q-RoPE, core attention, O-RoPE, and cast to FP8
• Dense attention kernels for prefill and backward
• Sparse kernels that power DeepSeek Sparse Attention
• Runs on NVIDIA GPUs and Huawei Ascend NPUs

Use Cases

• Inference of DeepSeek-V4.1 on NVIDIA GPUs and Huawei Ascend NPUs
• Token-level sparse attention for DSA decoding and prefill
• Dense multi-head attention forward and backward computations
• Fused attention pipelines for DeepSeek-V4.1 attention with permuted Q_b and Wv weights

Getting Started

• Clone the repository, initialize submodules, and install with pip using --no-build-isolation.
• For DSA decoding, call get_mla_metadata once before the decoding loop, then call flash_mla_with_kvcache in each decoding step.
• For DSA prefill, call flash_mla_sparse_fwd directly with the query, key-value, indices, scale, and optional parameters.
• For dense MHA prefill, use flash_attn_varlen_func, flash_attn_varlen_qkvpacked_func, or flash_attn_varlen_kvpacked_func.

Deployment & Requirements

• For NVIDIA, requires SM100 or SM103 GPU, CUDA 13.1 or above, and PyTorch 2.10 or above.
• For Huawei, requires Ascend 950 NPU, CANN 9.2.0 or above, torch_npu, and PyTorch 2.10 or above.
• Installation requires --no-build-isolation because setup.py imports torch during the build and the repository does not declare it as a PEP 518 build requirement.

Before You Adopt

• License: MIT. Review its terms before using, modifying, or distributing the project.
• The 2026.09.30 release removed Hopper architecture and earlier DeepSeek model support and changed the FP8/FP4 KV cache format, making it incompatible with previous versions.
• The fused norm-RoPE-attn-RoPE-cast kernel supports CUDA only and is not supported on Ascend.
• DSA decoding supports only FP8 and FP4 KV cache formats; unquantized bfloat16 KV cache is not supported.
• The DSA prefill kernel does not support a batch dimension.