FlashMLA

Introduction

FlashMLA is DeepSeek's library of optimized attention kernels, powering the DeepSeek-V3 and DeepSeek-V3.2-Exp models. This repository contains the following implementations:

Sparse Attention Kernels

These kernels power DeepSeek Sparse Attention (DSA), as introduced in this paper.

Dense Attention Kernels

News

Performance

Test & benchmark MLA decoding (Sparse & Dense):

python tests/test_flash_mla_dense_decoding.py
python tests/test_flash_mla_sparse_decoding.py

The dense MLA decoding kernel achieves up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5 with CUDA 12.8. The token-level sparse MLA decoding kernel (which uses an FP8 KV cache while performing the matrix multiplication in bfloat16) achieves 410 TFLOPS in compute-bound configuration on H800 SXM5 with CUDA 12.8, and achieves up to 350 TFlops on B200 (which is not really optimized yet).

Test & benchmark MHA prefill (Dense):

python tests/test_fmha_sm100.py

It achieves up to 1460 TFlops in forward and 1000 TFlops in backward computation on B200, as reported by NVIDIA.

Test & benchmark MLA prefill (Sparse):

python tests/test_flash_mla_sparse_prefill.py

It achieves up to 640 TFlops in forward computation on H800 SXM5 with CUDA 12.8, and achieves up to 1450 TFlops on B200, CUDA 12.9.

Requirements

Support matrix:

| Kernel | GPU Architecture | MLA Mode [2] | KVCache Format | | :---: | :---: | :---: | :---: | | Dense Decoding | SM90 | MQA | BF16 | | Sparse Decoding | SM90 & SM100 | MQA | FP8 [1] | | Dense Prefill | SM100 | MHA | | | Sparse Prefill | SM90 & SM100 | MQA | |

[1]: For more details on using FP8 KV cache, see documents below.

[2]: Here "MLA Mode" refers to the mode used for MLA calculation. MQA stands for Multi-Query Attention mode (i.e. head_dim_k = 576 with head_dim_v = 512), while MHA stands for Multi-Head Attention mode (i.e. head_dim_k = 192 / 128 with head_dim_v = 128). For a detailed explanation of these modes, please refer to the appendix of DeepSeek V3.2's Paper.

Installation

git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla git submodule update --init --recursive pip install -v .

Usage

MLA Decoding

To use the MLA decoding kernels, call get_mla_metadata once before the decoding loop to get the tile scheduler metadata. Then, call flash_mla_with_kvcache in each decoding step. For example:

from flash_mla import get_mla_metadata, flash_mla_with_kvcache
tile_scheduler_metadata, num_splits = get_mla_metadata(     cache_seqlens,     s_q * h_q // h_kv,     h_kv,     h_q,     is_fp8,     topk, )
for i in range(num_layers):     ...     o_i, lse_i = flash_mla_with_kvcache(         q_i, kvcache_i, block_table, cache_seqlens, dv,         tile_scheduler_metadata, num_splits,         is_causal, is_fp8_kvcache, indices,     )     ...

Where

FP8 KV Cache: If is_fp8_kvcache is set to True, the kernel reads the KV cache in the "FP8 with scale" format (described below). It dequantizes the cache to bfloat16 and performs attention computation in bfloat16. The output is also in bfloat16.

In the "FP8 with scale" format, each token's KV cache is 656 Bytes, structured as:

See tests/quant.py for quantization and dequantization details.

Sparse Attention (indices tensor): The indices tensor (if provided) enables token-level sparse attention by instructing the kernel to compute attention only for specified tokens.

Return Values: The kernel returns (out, lse), where:

See tests/test_flash_mla_decoding.py for a complete example.

Sparse MLA Prefill

For the sparse MLA prefill kernel, call flash_mla_sparse_fwd directly with the following parameters:

Note on batching: This kernel does not support a batch dimension. For multi-batch inference, reshape the input tensors and adjust the indices parameter to simulate batch processing.

Invalid indices: Set invalid entries in indices to -1 or any number >= s_kv.

Return Values and Equivalent PyTorch Code: The kernel returns (out, max_logits, lse). This is equivalent to the following PyTorch operations:

Q: [s_q, h_q, d_qk], bfloat16
kv: [s_kv, h_kv, d_qk], bfloat16 indices: [s_q, h_kv, topk], int32
kv = kv.squeeze(1)  # [s_kv, d_qk], h_kv must be 1 indices = indices.squeeze(1)    # [s_q, topk] focused_kv = kv[indices]    # For the i-th sequence (s_q), the corresponding KV tokens are selected from the KV cache based on indices[i, :]. This operation results in a tensor of shape [s_q, topk, d_qk].
P = (Q @ focused_kv.transpose(-1, -2)) * sm_scale * math.log2(math.e)    # [s_q, h_q, topk] max_logits = P.max(dim=-1) # [s_q, h_q] lse = log2sumexp2(P, dim=-1, base=2)   # [s_q, h_q]，"log2sumexp2" means that the exponentiation and logarithm are base-2 S = exp2(P - lse)      # [s_q, h_q, topk] out = S @ focused_kv  # [s_q, h_q, d_qk]
return (out, max_logits, lse)

See tests/test_flash_mla_prefill.py for a complete example.

Dense MHA Prefill

This kernel implements the standard dense Multi-Head Attention (MHA) forward and backward operations. It can be called using:

The usage is similar to the flash_attn package. See tests/test_fmha_sm100.py for a complete example.

Acknowledgement

FlashMLA is inspired by FlashAttention 2&3 and cutlass projects.

Community Support

MetaX

For MetaX GPUs, visit the official website: MetaX.

The corresponding FlashMLA version can be found at: MetaX-MACA/FlashMLA

Moore Threads

For the Moore Threads GPU, visit the official website: Moore Threads.

The corresponding FlashMLA version is available on GitHub: MooreThreads/MT-flashMLA.

Hygon DCU

For the Hygon DCU, visit the official website: Hygon Developer.

The corresponding FlashMLA version is available here: OpenDAS/MLAttention.

Intellifusion

For the Intellifusion NNP, visit the official website: Intellifusion.

The corresponding FlashMLA version is available on Gitee: Intellifusion/tyllm.

Iluvatar Corex

For Iluvatar Corex GPUs, visit the official website: Iluvatar Corex.

The corresponding FlashMLA version is available on GitHub: Deep-Spark/FlashMLA

AMD Instinct

For AMD Instinct GPUs, visit the official website: AMD Instinct.

The corresponding FlashMLA version can be found at: AITER/MLA

Citation

@misc{flashmla2025,
      title={FlashMLA: Efficient Multi-head Latent Attention Kernels},       author={Jiashi Li, Shengyu Liu},       year={2025},       publisher = {GitHub},       howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, }