deepseek-ai/FlashMLA

FlashMLA: Efficient Multi-head Latent Attention Kernels

FlashMLA – High‑performance attention kernels for DeepSeek models

What it is – FlashMLA is a C++/CUDA (and Huawei Ascend) library that implements highly‑optimized attention operators used by DeepSeek’s large language models, most recently the DeepSeek‑V4.1‑Flash model. It provides:

  • Sparse‑attention kernels (token‑level “DeepSeek Sparse Attention”) for both the prefill phase (when the model processes the initial prompt) and the decoding phase (generating one token at a time). These kernels work with FP8 or FP4‑quantized key‑value caches, dramatically reducing memory bandwidth while keeping accuracy.
  • Dense multi‑head attention (MHA) kernels for standard (non‑sparse) prefill and backward passes.
  • A fused kernel that combines Q‑norm, RoPE (rotary positional encoding), the core attention computation, the conjugate RoPE on the output side, and the final cast to FP8, all in a single launch. This eliminates the overhead of many tiny kernels.
  • Support for NVIDIA GPUs (SM100/SM103, CUDA 13.1+) and Huawei Ascend 950 NPUs (via CANN 9.2+ and torch_npu).

Why it matters – Large language models spend most of their inference time in the attention operation. By rewriting attention in low‑precision (FP8/FP4) and exploiting sparsity, FlashMLA can reach >1 000 TFlops on modern hardware, which is close to the theoretical peak of the devices. This translates into faster response times and lower cost for serving models like DeepSeek‑V4.1.

How to use it

  1. Install – Clone the repo, pull submodules, and run pip install -v . --no-build-isolation. The build script detects whether you are on CUDA or Ascend and compiles the appropriate extensions.
  2. Prefill (sparse) – Call flash_mla_sparse_fwd(q, kv, indices, …) with your query tensor, KV cache, and a 3‑D indices tensor that tells the kernel which KV tokens each query token should attend to.
  3. Decoding (sparse) – First obtain scheduler metadata via get_mla_metadata(), then repeatedly call flash_mla_with_kvcache(...) for each generation step. The function expects the FP8/FP4 KV cache layout described in the README.
  4. Dense MHA – Use the familiar flash_attn_varlen_func / flash_attn_varlen_qkvpacked_func APIs (compatible with the popular flash_attn package) but note the extra head_dim_qk argument and that the return value includes both the output and the log‑sum‑exp (lse).
  5. Fused kernel – After permuting the Q‑projection and Wv‑projection weights into the layout required by the fused kernel (helpers are provided in fused_norm_rope_attn_rope_cast.permute_*), call prefill(...) or decode(...) to get FP8 outputs ready for a downstream FP8‑einsum.

Performance highlights (as reported in the repo)

  • Sparse prefill: up to 1350 TFlops on an NVIDIA B200 GPU (CUDA 13.3). On Ascend 950 NPU, 410 TFlops (≈95 % of hardware peak).
  • Sparse decoding: up to 1024 TFlops on B200; 360 TFlops on Ascend 950 (≈83 % of peak).
  • Fused norm‑RoPE‑attn‑RoPE‑cast: up to 1460 TFlops (prefill) and 950 TFlops (decoding) on B200, with a 10‑15 % extra boost over the unfused version.
  • Dense MHA: up to 1460 TFlops forward and 1000 TFlops backward on B200.

Key technical details

  • KV cache formats – Only FP8 (float8_e4m3) or FP4 (e2m1) quantized caches are supported. The layout stores raw quantized values followed by per‑token scale factors, enabling fast dequantization on‑the‑fly.
  • Indices tensor – A 3‑D integer tensor that encodes which KV “pages” each query token should attend to. Invalid entries are marked with -1 and ignored by the kernel.
  • Register‑spill safety – After building the CUDA extension, the build system checks for register spilling; the build will fail unless you set FLASH_MLA_SKIP_REG_SPILL_CHECK=1.
  • Compatibility note – The September 2026 release drops support for NVIDIA Hopper GPUs and older DeepSeek models (V3/V3.2/V4.0) and changes the KV‑cache format. Users needing those older configurations must checkout the legacy commit linked in the README.

Where to learn more – The repository contains deep‑dive markdown reports (docs/…) that explain the algorithms and low‑level optimizations, as well as test scripts (tests/*.py) that can be run to verify performance on your hardware.

FlashMLA is essentially a performance‑critical building block for DeepSeek’s LLM inference pipeline, targeting both cutting‑edge NVIDIA GPUs and Huawei Ascend NPUs.

Related

  • Project
    QwenLM/FlashQLA

    FlashQLA is a Python‑package kernel library that speeds up the Gated‑Delta‑Rule linear‑attention used in Qwen models. Built with TileLang, it offers 2‑3× faster forward and ~2× faster backward passes on modern NVIDIA GPUs, automatic intra‑card parallelism, and a drop‑in API compatible with flash‑linear‑attention.

  • Dispatch
    FlashQLA: CP-/Bwd-Friendly Fused Linear Attention Kernels for GDN

    Qwen has open-sourced FlashQLA, a high-performance linear attention kernel library built on TileLang that achieves 2-3x forward and 2x backward speedups for Gated Delta Network (GDN) layers on NVIDIA Hopper GPUs.

  • Project
    flashinfer-ai/flashinfer

    FlashInfer is a GPU‑kernel library (with a Python package) that provides ultra‑fast, low‑precision attention, GEMM, MoE, and sampling operations for LLM inference. It automatically selects the best backend (FlashAttention‑2/3, cuDNN, CUTLASS, TensorRT‑LLM) for a wide range of NVIDIA GPUs, from Turing up to Blackwell, and is already used by popular inference servers like vLLM and TensorRT‑LLM.

  • Project
    fla-org/flash-linear-attention

    A hardware-efficient library of subquadratic kernels and layers for linear attention, state space models, and hybrid LLM architectures across multiple hardware backends.

  • Dispatch
    vLLM DeepSeek V4 Support: Efficient Long-context Attention

    vLLM now supports DeepSeek V4-Pro and V4-Flash, implementing a new attention mechanism that enables context lengths up to one million tokens with significant KV cache memory savings.