What is FlashMLA?

FlashMLA is an MLA (Multi-Head Latent Attention) decoding kernel launched by DeepSeek on February 24, 2025, optimized for NVIDIA Hopper architecture GPUs (such as the H800), with special optimization for variable-length sequence processing services.

Key Features of FlashMLA:

• BF16 Support: FlashMLA supports the BF16 (Bfloat16) data type, making it more efficient in computation and memory usage.

• Paged KV Cache: Manages the key-value (KV) cache through a paging mechanism with a block size of 64, enabling efficient handling of large-scale sequences.

• High Performance: FlashMLA achieves memory bandwidth of up to 3000 GB/s (in memory-bound scenarios) and computational performance of up to 580 TFLOPS (in compute-bound scenarios, based on the BF16 data type).

Technical Background of FlashMLA:

FlashMLA was created to address the compute and memory bottlenecks that large language models face during inference. Traditional Multi-Head Attention (MHA) requires substantial memory to store key-value pair (KV) caches when processing long sequences, which limits model deployment on hardware with limited resources. MLA reduces the size of the KV cache by introducing a latent attention mechanism while maintaining model performance.

Application Scenarios of FlashMLA:

FlashMLA is particularly well-suited for natural language processing (NLP) tasks that require efficient decoding, such as large language model (LLM) inference. It is optimized for variable-length sequences and has been validated in real production environments, making it especially suitable for high-performance computing needs.

Technical Implementation of FlashMLA

• Low-Rank Compression: MLA achieves effective compression of the KV cache through low-rank matrix factorization, reducing memory usage.

• KV Cache Optimization: Optimizes the KV cache mechanism, significantly reducing hardware resource requirements and thereby lowering inference costs.

• Parallel Decoding: Introduces a parallel decoding mechanism that allows multiple tokens to be processed simultaneously, significantly boosting inference speed.

Performance Improvements of FlashMLA

After adopting FlashMLA, DeepSeek saw approximately a 5% improvement in accuracy on natural language processing tasks, a 20% increase in inference speed, and a 15% reduction in computational resource consumption. These improvements enable DeepSeek to perform better in real-time interactive scenarios (such as conversational ai and real-time translation).

FlashMLAInstallation and Usage

Environment Requirements:

• Hopper architecture GPU (such as NVIDIA A100)

• CUDA 12.3 or later

• PyTorch 2.0 or later

1. First, you need to install the FlashMLA library. You can install it using the following commands:

git clone https://github.com/deepseek-ai/FlashMLA.git
cd FlashMLA
python setup.py install

Or if you have already cloned the repository and want to rebuild:

python setup.py clean --all && python setup.py build_ext --inplace

2. Get MLA Metadata

Before using FlashMLA, you need to obtain the MLA metadata. This typically involves preparing input tensors and other necessary parameters.

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# Assume you already have cache_seqlens and other related variables
cache_seqlens = [...]  # list of lengths for each sequence
s_q = ...              # query dimension
h_q = ...              # number of heads
h_kv = ...             # number of key-value heads

tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)

3. Perform MLA Decoding

Next, you can perform MLA decoding operations. Assume you already have the necessary components such as the query matrix q_i, key-value cache kvcache_i, and block table.

dv = ...               # output dimension

for i in range(num_layers):  # iterate over each layer
    o_i, lse_i = flash_mla_with_kvcache(
        q_i[i],                # query matrix for the current layer
        kvcache_i[i],           # key-value cache for the current layer
        block_table,            # block table
        cache_seqlens,          # cache sequence lengths
        dv,                     # output dimension
        tile_scheduler_metadata,# MLA metadata
        num_splits,             # number of splits
        causal=True             # whether to use causal masking
    )
    # Continue processing the output results o_i and lse_i

FlashMLA GitHub:https://github.com/deepseek-ai/FlashMLA