FEATURED · 精选文章

ik_llama.cpp 中的 MLA 矩阵吸收:基于压缩 KV 缓存的长上下文预填充优化原理与实现

发布时间 / 2026/9/18 1:36:48
来源 / 创域科博编辑部
栏目 / 资讯中心
ik_llama.cpp 中的 MLA 矩阵吸收:基于压缩 KV 缓存的长上下文预填充优化原理与实现 ik_llama.cpp 中的 MLA 矩阵吸收基于压缩 KV 缓存的长上下文预填充优化原理与实现【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp本文围绕 ik_llama.cpp 社区讨论《MLA matrix absorption》展开系统讲解 DeepSeek 系列模型所采用的 MLAMulti-head Latent Attention多头潜在注意力中矩阵吸收这一数学优化技巧它如何在长上下文预填充prefill阶段避免显式解压压缩后的 KV 缓存从而同时压低计算量与中间显存占用并在此基础上结合仓库源码说明该思想在 ik_llama.cpp 中的工程落地——FlashMLA 系列实现、-mla与-amb参数的作用以及注意力分数分裂Move Elision等配套优化。读完本文你将掌握 MLA 吸收的数学原理、FLOPs/缓存收益的量化对比以及如何在 ik_llama.cpp 中通过命令行参数实际控制 MLA 推理路径。一、背景MLA 与长上下文预填充的矛盾MLA 是 DeepSeek-V2 引入的注意力机制其核心思想是不直接为每个 token 缓存高维的 K/V而是先通过低秩投影把 K/V 压缩为潜在表示latent representationc_t^KV再在注意力计算时用上投影矩阵W^{UK}、W^{UV}恢复出完整的 K/V。这样 KV 缓存可以大幅缩小典型压缩率高达 98.6%但也带来了新问题缓存的压缩形态内存占用极小但每次注意力计算前都必须解压缩decompress而解压缩本身是高开销的大矩阵乘法。在长上下文场景下预填充需要为海量 token 反复执行解压缩这使得 MLA 从内存友好退化为计算密集 中间显存爆炸。讨论中给出的数学表达如下k_t^C W^{UK} · c_t^KV 昂贵的操作 v_t W^{UV} · c_t^KV 昂贵的操作如果每次注意力计算都要把整段历史 KV 从压缩态解压到完整维度预填充的算力和显存都会成为瓶颈。这正是 ktransformers 等长上下文实现中 prefill 优化的目标也是本讨论帖的出发点。二、关键洞察矩阵乘法的结合律讨论帖给出的核心洞察非常朴素矩阵乘法满足结合律。既然解压缩是用矩阵去乘压缩表示那么完全可以把这些解压缩矩阵吸收absorb进相邻的运算中K 吸收把 K 的解压缩矩阵吸收进 Q 投影V 吸收把 V 的解压缩矩阵吸收进输出投影。于是整条计算链不再显式出现解压后的完整 K/V所有计算都直接在压缩表示上进行只在最后一步才把结果还原到需要的维度。三、K 吸收注意力分数直接算在压缩空间以非 RoPE 部分的注意力分数为例原始计算是{q_t^C}^T · k_t^C (W^{UQ} c_t^Q)^T · (W^{UK} c_t^KV)利用结合律改写为{q_t^C}^T · k_t^C (c_t^Q)^T · ((W^{UQ})^T · W^{UK}) · c_t^KV其中(W^{UQ})^T · W^{UK}可以预计算于是运行时只剩两个小规模矩阵乘避免了显式计算完整维度的k_t^C用更小的运算替代了原来巨大的矩阵乘法带来的 FLOPs 缩减幅度惊人从 33.64 MFLOP/token 降到 0.28 MFLOP/token。需要说明的是RoPE 部分无法参与吸收旋转位置编码破坏结合律因此实际实现中 RoPE 部分与非 RoPE 部分要分开处理这正是下文Move Elision优化的由来。四、V 吸收延迟解压缩 加权求和融合V 分支的思路与 K 分支对称但更激进不显式解压 V而是把注意力加权求和直接做在压缩空间最后再投影到完整维度。讨论帖用 PyTorch 的 einsum 给出了等价写法# 不再显式解压 V attn_output torch.einsum(bhql,blc-bhqc, attn_weights, compressed_kv) # 压缩空间内加权求和 attn_output torch.einsum(bhqc,hdc-bhqd, attn_output, W_UV) # 延迟解压缩 u torch.einsum(hdD,bhqd-bqD, W_O, attn_output) # 输出投影这样做的收益完全避免了物化完整的v_t其维度约为压缩表示c_t^KV的 128 倍把解压缩与加权求和融合成一步中间张量的内存占用被压到最低。仓库中的对应实现在 ik_llama.cpp 中V 分支的延迟解压缩体现在 src/graphs/build_deepseek2.cppFlashMLA 先在压缩潜在缓存上完成注意力得到kqv_compressed随后才用wv_b投影还原维度auto wv_b model.layers[il].wv_b; // ... reshape 到 {kv_lora_rank, n_embd_head_v, n_head} ... kqv ggml_mul_mat(ctx0, wv_b, kqv_compressed); // 延迟解压缩压缩空间注意力结果 - 完整 V 维度 cur ggml_reshape_2d(ctx0, kqv, n_embd_head_v*n_head, n_tokens); cur llm_build_lora_mm(lctx, ctx0, model.layers[il].wo, cur); // 输出投影也就是说W^{UV}源码中的wv_b被推迟到注意力输出之后才应用与讨论帖中torch.einsum(bhqc,hdc-bhqd, attn_output, W_UV)的晚期解压完全对应。而 K 分支的吸收则体现在 src/graphs/build_deepseek2.cpp用wk_b直接对q_nope做投影把 K 解压缩吸收进 Q 侧q_nope2 ggml_mul_mat(ctx0, wk_b, q_nope); // (W^UQ)^T · W^UK 吸收进 Q q ggml_concat(ctx0, ggml_permute(ctx0, q_rope, 0, 2, 1, 3), q_nope2, 0);随后把q与压缩后的 KV 缓存直接送入 FlashMLA 注意力核src/graphs/build_deepseek2.cpp全程不物化解压后的完整 K。五、Move Elision分裂注意力分数消灭拼接临时张量吸收解决了解压缩开销但还有一个隐患RoPE 部分无法吸收必须保留完整维度的 RoPE K/Q如果照旧把 RoPE 与非 RoPE 部分拼接成大张量会制造大量高维临时缓冲区讨论帖中 Q/K 为 192 维DeepSeek 注意力头维度达 576/512临时张量尤为昂贵。讨论帖给出的解法是注意力分数分两部分计算后相加attn_weights ( torch.matmul(q_pe, k_pe.transpose(2, 3)) # RoPE 部分走常规路径 torch.einsum(bhqc,blc-bhql, q_nope, compressed_kv) # 非 RoPE 部分吸收路径 )收益消除内存密集的拼接concat操作避免存储完整维度的 Q/K 张量192 维让 RoPE 与非 RoPE 各自走最合适的计算路径。仓库中的对应实现在 ik_llama.cpp 中RoPE 与压缩缓存的分流同样清晰q_rope走ggml_repeat 拼接后与 KV 缓存的 RoPE 段交互kv_cache_lora仅潜在部分则直接作为 FlashMLA 的 V 输入参与注意力src/graphs/build_deepseek2.cpp即压缩空间内加权求和与讨论帖的torch.einsum(bhql,blc-bhqc, ...)一一对应。六、为什么不把所有矩阵都预先吸收理论上还可以更激进预先计算W_new^UQ (W^UQ)^T · W^UK和W_new^O W^O · W^UV把吸收彻底固化进权重。讨论帖明确指出这并不高效W_new^UQ会是一个巨大的低秩矩阵如 1536×512W_new^O更加庞大如 5120×512 × 每个注意力头实际的顺序小运算实现即上文展示的 K/V 吸收路径比预吸收大矩阵更快。这也是为什么工程上选择运行时用预计算小矩阵、顺序执行小矩阵乘而不是离线把所有矩阵融合成一个大矩阵。七、性能影响从内存受限到计算受限讨论帖给出的量化对比表格不同优化路径在相同模型规模下的表现优化路径KV 缓存大小计算量 (FLOPs/token)相对基线加速基线缓存已解压81.92 KB0.08 M1.0×压缩缓存无吸收1.15 KB33.64 M0.3×吸收 Move Elision1.15 KB0.28 M1.2×关键收益KV 缓存比解压版本小约71 倍计算量比无吸收的压缩方案低约120 倍更小的缓存与中间张量意味着更大的 batch 可以塞进显存这对 GPU 利用率至关重要该方案在保持约 98.6% KV 缓存压缩率的同时把 MLA 从内存受限问题转化为计算受限问题更好地利用现代 GPU 张量核心。八、ik_llama.cpp 中的工程落地与历史演进讨论帖作者 ikawrakow 在回复中明确指出这条技术路线在 ik_llama.cpp 中早有更早的独立探索讨论 #246 比外部文档早约一个月提出了减少 MLA 乘法累加量的基本思路#260 解释了-amb选项#273 则是当时仓库中最好的 MLA 版本——MLA2 用于预填充prompt processing原始 MLAMLA1用于 token 生成。作者还提到迟迟未能更快完成的主要原因是在 CUDA 上高效实现 MLA1 的挣扎——这源于 DeepSeek 异常大的注意力头维度576 和 512。8.1 FlashMLA 时间线从 README.md 可以看到完整的演进足迹2025-02-09为 DeepSeek 模型引入 MLAPR 1882025-02-13MLA 支持Q8_0量化 KV 缓存PR 2062025-02-27MLA 无转置缓存PR 2352025-03-03 起FlashMLAMLA Flash Attention在 CPUPR 240/243、CUDAPR 247上陆续落地并持续优化PR 253、PR 265、PR 2732025-05-07FlashMLA-3DeepSeek 模型在 CUDA 上的最快实现PR 386需 Ampere 或更新的 NVIDIA GPU。其中 FlashMLA-3PR 273正是讨论帖所说MLA2 用于预填充、MLA1 用于 token 生成的集大成版本而吸收路径在 src/graphs/build_deepseek2.cpp 中被称为FlashMLA-3 absorb path。8.2-mla选择 MLA 推理变体命令行参数-mla, --mla-use用于选择 MLA 的实现方式参数解析见 common/common.cpp默认值与语义见 common/common.hint mla_attn 3; // MLA 0: standard, 1: MLA with K and V^T cache, // 2: MLA with just K cache, 3: the best of both worlds0标准 MLA显式解压用于对照/兜底1MLA缓存 K 和 V 的转置FlashMLA 路径用于 token 生成2MLA只缓存 K适合预填充3默认两者结合预填充用 MLA2、生成用 MLA1即两全其美。该参数同样适用于其他采用 MLA 的较新模型如 Qwen3、Kimi-K2 等详见 docs/parameters.md。8.3-amb限制注意力中间计算缓冲区讨论帖 #260 提到的-amb选项-amb, --attention-max-batch用于限制 MLA 注意力计算时临时 compute 缓冲区的用量参数解析见 common/common.cpp默认值见 common/common.hint attn_max_batch 256; // Max batch size to use when computing attention其语义是容忍的最大 K×Q 中间张量大小MiB。在 src/graphs/build_deepseek2.cpp 中可以看到它的实际作用当kv_f32_size attn_max_batch时把一次性参与解压缩/注意力的头数n_max_head逐步降为 1再按n_head % niter 0找到一个能整除、且规模不超过上限的分组数把大计算拆成多次迭代执行auto kv_f32_size model.layers[il].wkv_b-ne[1] * kv_cache_nope-ne[1] * sizeof(float) / (1024*1024); int n_max_head n_head; if (cparams.attn_max_batch 0 kv_f32_size cparams.attn_max_batch) { n_max_head 1; for (int niter 2; niter n_head; niter) { if (n_head % niter 0 kv_f32_size/niter cparams.attn_max_batch) { n_max_head n_head/niter; break; } } }对于 512K 上下文、单 token 的极端场景这种分头迭代可以避免 GB 级的临时缓冲区源码注释在 src/graphs/build_deepseek2.cpp 给出非 FlashMLA 路径的分步执行逻辑。-amb的默认值从 0不限制改为 256正是为了避免 MLA 模型上出现过大的临时计算缓冲区docs/parameters.md。8.4 预填充的另一种选择pp_opt 物质化路径值得一提的还有预填充阶段的pp_opt分支src/graphs/build_deepseek2.cpp当mla 1、n_tokens 128且n_kv 1024时实现会反其道而行之——从潜在缓存中物质化出每个 rank 的完整 K/V改用标准ggml_flash_attn_ext而非 FlashMLA-3 吸收路径。注释表明这是为了在长上下文预填充下获得更优的性能/精度权衡说明吸收并不是唯一的答案而是工程上依据上下文长度、batch 大小动态选择的多条路径之一。九、总结矩阵吸收的本质可以概括为四步利用矩阵结合律把解压缩矩阵推入相邻运算K 解压缩并入 Q 投影V 解压缩并入输出投影V 延迟解压缩注意力加权求和直接在压缩潜在表示上进行最后才投影到完整维度分裂注意力分数RoPE 部分与非 RoPE 部分分开计算再相加避免高维张量拼接全程保持压缩表示让压缩 KV 缓存的红利贯穿整个计算链。由此MLA 从每 token 33.64 MFLOP 的解压开销 巨大临时缓冲区变成每 token 0.28 MFLOP 71 倍更小的 KV 缓存从内存受限问题转化为计算受限问题。在 ik_llama.cpp 中这套思想经由 FlashMLA → FlashMLA-2 → FlashMLA-3 的迭代README 时间线可见一斑最终固化在 src/graphs/build_deepseek2.cpp 的多条 MLA 推理路径中用户则可通过-mla 0/1/2/3选择推理变体、通过-amb调节注意力中间缓冲区上限在预填充吞吐、显存占用与生成延迟之间找到适合自己硬件配置的平衡点。【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻