FlashMLA：DeepSeek發(fā)布的高效的MLA解碼內(nèi)核，優(yōu)化了變長序列的處理服務(wù)

映技派于2025-02-24發(fā)布在Ai產(chǎn)品

FlashMLA是什么？

FlashMLA是DeepSeek在2025年2月24日推出的一款針對NVIDIA Hopper架構(gòu)GPU（如H800）優(yōu)化的MLA（Multi-Head Latent Attention）解碼內(nèi)核，特別優(yōu)化了變長序列的處理服務(wù)。

DeepSeek發(fā)布FlashMLA.webp

FlashMLA的主要特性：

BF16支持：FlashMLA支持BF16（Bfloat16）數(shù)據(jù)類型，這使得它在計(jì)算和內(nèi)存使用上更加高效。
分頁KV緩存：通過分頁機(jī)制管理鍵值（KV）緩存，塊大小為64，這使得它能夠高效處理大規(guī)模序列。
高性能：FlashMLA的內(nèi)存帶寬可達(dá)3000 GB/s（在內(nèi)存瓶頸場景下），計(jì)算性能可達(dá)580 TFLOPS（在計(jì)算瓶頸場景下，基于BF16數(shù)據(jù)類型）。

FlashMLA的技術(shù)背景：

FlashMLA的出現(xiàn)是為了解決大型語言模型在推理過程中面臨的計(jì)算和內(nèi)存瓶頸問題。傳統(tǒng)的多頭注意力機(jī)制（MHA）在處理長序列時，需要大量的內(nèi)存來存儲鍵值對（KV）緩存，這限制了模型在有限硬件資源上的部署。MLA通過引入潛在注意力機(jī)制，減少了KV緩存的大小，同時保持了模型的性能。

FlashMLA的應(yīng)用場景：

FlashMLA特別適用于需要高效解碼的自然語言處理（NLP）任務(wù)，如大語言模型（LLM）的推理。它針對變長序列進(jìn)行了優(yōu)化，并在實(shí)際生產(chǎn)環(huán)境中經(jīng)過了驗(yàn)證，特別適合高性能計(jì)算需求。

FlashMLA開源.webp

FlashMLA的技術(shù)實(shí)現(xiàn)

低秩壓縮：MLA通過低秩矩陣分解實(shí)現(xiàn)KV緩存的有效壓縮，減少了內(nèi)存占用。
KV緩存優(yōu)化：優(yōu)化KV緩存機(jī)制，顯著降低了硬件資源需求，從而降低了推理成本。
并行解碼：引入并行解碼機(jī)制，允許同時處理多個token，顯著提升推理速度。

FlashMLA的性能提升

采用FlashMLA后，DeepSeek在自然語言處理任務(wù)中的準(zhǔn)確率提升了約5%，推理速度提高了20%，計(jì)算資源消耗降低了15%。這些改進(jìn)使得DeepSeek在實(shí)時交互場景（如對話ai、實(shí)時翻譯）中表現(xiàn)更優(yōu)。

FlashMLA安裝使用

環(huán)境要求：

Hopper 架構(gòu) GPU（如 NVIDIA A100）
CUDA 12.3 及以上版本
PyTorch 2.0 及以上版本

1. 首先，你需要安裝 FlashMLA 庫。你可以通過以下命令進(jìn)行安裝：

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

或者如果你已經(jīng)克隆了倉庫并且想要重新構(gòu)建：

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

2. 獲取 MLA 元數(shù)據(jù)

在使用 FlashMLA 之前，你需要獲取 MLA 的元數(shù)據(jù)。這通常涉及準(zhǔn)備輸入張量和其他必要的參數(shù)。

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# 假設(shè)你已經(jīng)有了 cache_seqlens 和其他相關(guān)變量
cache_seqlens = [...]  # 每個序列的長度列表
s_q = ...              # 查詢維度
h_q = ...              # 頭數(shù)量
h_kv = ...             # 鍵值頭數(shù)量

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

3. 執(zhí)行 MLA 解碼

接下來，你可以執(zhí)行 MLA 解碼操作。假設(shè)你已經(jīng)有查詢矩陣 q_i、鍵值緩存 kvcache_i、塊表 block_table 等必要組件。

dv = ...               # 輸出維度

for i in range(num_layers):  # 循環(huán)遍歷每一層
    o_i, lse_i = flash_mla_with_kvcache(
        q_i[i],                # 當(dāng)前層的查詢矩陣
        kvcache_i[i],           # 當(dāng)前層的鍵值緩存
        block_table,            # 塊表
        cache_seqlens,          # 緩存序列長度
        dv,                     # 輸出維度
        tile_scheduler_metadata,# MLA 元數(shù)據(jù)
        num_splits,             # 劃分?jǐn)?shù)目
        causal=True             # 是否因果掩碼
    )
    # 繼續(xù)處理輸出結(jié)果 o_i 和 lse_i

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