当前位置：首页 > news >正文

FlashMLA：把 KV Cache 压缩到原来的八分之一

news 2026/5/24 16:19:54

标准 MHA 的 KV Cache 是推理显存的第一大户。LLaMA-7B，32 层，每层 32 头，HeadDim=128，SeqLen=128K——KV Cache 吃 40GB。MLA（Multi-head Latent Attention）用低秩分解把 KV 映射到一个远小于 HeadDim 的潜在空间，存压缩版。FlashMLA 是这套机制在昇腾NPU上的高性能实现。

MLA 的压缩原理

MLA 的核心是 KV Compression——不是存原始 K/V，而是存它们的低秩表示。原始 K 和 V 分别通过下投影矩阵 W_DK 和 W_DV 映射到一个很小的潜在空间（比如 16 维），存下来的就是这个 16 维的压缩向量。

推理时再通过上投影矩阵把压缩向量恢复到原始 HeadDim 做 Attention 计算。

标准 MHA 的 K/V 存储： K = X @ W_K → [B, H, S, 128] ← 每层每 Token 存 128 个 FP16 V = X @ W_V → [B, H, S, 128] MLA 的 KV 存储（DeepSeek-V2 参数）： C_KV = X @ W_DKV → [B, H, S, 16] ← 每层每 Token 只存 16 个 FP16 使用时解压： K = C_KV @ W_UK → [B, H, S, 128] V = C_KV @ W_UV → [B, H, S, 128] 压缩比：128 / 16 = 8:1 128K 上下文 KV Cache：40GB → 5GB

FlashMLA 的计算流程

FlashMLA 在 FlashAttention 的分块框架上加了两步：解压和 RoPE 融合。不是一次性把所有 K/V 解压出来——而是每次算一个 Tile 的 Attention 时才解压这个 Tile 对应的 K/V。这样解压的中间结果不占显存。

FlashMLA 的 Tile 级执行流程： for each Tile in Q: for each Tile in KV: 1. 解压：K_tile = C_KV_tile @ W_UK (16×128 矩阵乘，在 Cube Unit 上算) 2. 解压：V_tile = C_KV_tile @ W_UV 3. Attention：S = Q_tile @ K_tile^T (在 Cube Unit 上) 4. Online Softmax（FlashAttention 的分块累积） 5. 累积：O += Softmax(S) @ V_tile

三个关键优化：

解压矩阵 W_UK/W_UV 很小（16×128），刚好塞进 Cube Unit 单次运算。
RoPE 融进解压步骤——传统 Attention 里 RoPE 是独立 Kernel，FlashMLA 在解压的同时算 RoPE，省一次 L1 写回。
解压在 Cube Unit 上跑时，Vector Unit 同时做下一个 Tile 的 Softmax 累积——双单元并行。

昇腾NPU上的 FlashMLA 实现

ops-transformer 仓库里 FlashMLA 的核心代码：

// FlashMLA Kernel——简化版核心循环__aicore__voidflash_mla_kernel(GlobalTensor<fp16>q,// [B, H, S_q, D]GlobalTensor<int8>kv_cache,// 压缩 KV [B, H, S_kv, 16]GlobalTensor<fp16>W_UK,// 解压矩阵 [H, 16, D]GlobalTensor<fp16>W_UV,// 解压矩阵 [H, 16, D]GlobalTensor<fp16>output// [B, H, S_q, D]){// L1 上分配 Tile BufferLocalTensor<fp16>q_tile,k_tile,v_tile,scores,attn,out_acc;for(intq_block=0;q_block<num_q_blocks;q_block++){// 加载 Q 的一个 Tile 到 L1DataCopy(q_tile,q[q_block]);// Online Softmax 状态fp32 m_prev=-INFINITY,l_prev=0;for(intkv_block=0;kv_block<num_kv_blocks;kv_block++){// 1. 解压 K/V——Cube Unit 做小矩阵乘MatMul(k_tile,kv_cache[kv_block],W_UK);// [T, 16] @ [16, D]MatMul(v_tile,kv_cache[kv_block],W_UV);// 2. RoPE 融合——直接在 k_tile 上原地做ApplyRoPEInPlace(k_tile,kv_block*block_size);// 3. Attention 计算MatMul(scores,q_tile,k_tile,/*transB=*/true);// 4. Online Softmax + 累积fp32 m_curr=row_max(scores);fp32 m_new=max(m_prev,m_curr);fp32 l_curr=row_sum(exp(scores-m_new));out_acc=out_acc*exp(m_prev-m_new);out_acc+=exp(scores-m_new)*v_tile;m_prev=m_new;l_prev=exp(m_prev-m_new)*l_prev+l_curr;}// 最终归一化DataCopy(output[q_block],out_acc/l_prev);}}

显存收益

DeepSeek-V2 配置下，MLA 的 KV Cache 从每 Token 128KB（标准 MHA，128 头 × 128 HeadDim × 2）降到约 16KB。128K 上下文从 16GB 降到 2GB。省下的 14GB 直接把 Batch Size 从 4 推到 32——在线推理吞吐涨 8 倍。

Tensor缓存也跟着受益——压缩后的 KV 在 L1 上占的空间从 2MB/Tile 降到 256KB/Tile，同一块 L1 能缓存的上下文 Tile 数翻 8 倍，FlashAttention 的分块数减少，整体延迟更低。