ik_llama.cpp 超长上下文实战用-amb参数将注意力计算缓冲区从 42.7GB 降到 2GiB【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp导读本文围绕 ik_llama.cpp 仓库中 PR #237「Reduce size of compute buffers」 展开剖析长上下文推理中计算缓冲区compute buffer的真正瓶颈——注意力K*Q张量并完整讲解其引入的-amb / --attn-max-batch参数及其分步注意力计算原理。读完本文你将掌握如何在不牺牲 prefill 速度的前提下为 DeepSeekV3/R1 等 163k 上下文大模型配置合理的注意力批次上限并在 CPU / GPU 上跑满完整上下文。为什么 KV cache 变小了显存却依然不够用ik_llama.cpp 长期致力于压缩 KV cache但在 PR #237 的描述 中作者明确指出真正阻碍超长上下文实际运行的是计算缓冲区的大小而不是 KV cache。以 DeepSeekV3/R1 为例模型声称支持 163k token 的最大上下文。如果按n_head 128、默认 u-batch-ub为 512 来估算注意力部分的K*Q张量尺寸为n_ctx × n_ubatch × n_head × sizeof(float)代入数值163000 × 512 × 128 × 4 bytes ≈ 42.7 GB。也就是说仅仅为了保存一个K*Q中间结果就需要为每块 GPU 预留超过 40GB 的 CUDA compute buffer。即便完全不使用 GPU、纯 CPU 推理42.7GB 也绝非可以忽略的开销。一个看似可行的规避方式是调小 u-batch把-ub降到 64K*Q缓冲区就能缩到约 5GB从而给 KV cache 和注意力张量腾出显存。但代价是prefill提示词处理速度下降约 3 倍——这正是很多长上下文部署要么跑不满、要么跑得慢的两难根源。缓解手段之一复用缓冲区与 Softmax 原地计算K*Q之后还有一个等尺寸的softmax(K*Q)。幸运的是后端back-end足够聪明会复用同一个缓冲区来存放 softmax 结果因此实际占用的峰值并不是两份 42.7GB而是一份。即便如此单份K*Q的尺寸在长上下文下仍然直接决定了 compute buffer 的上限因为它会在模型加载时的测试计算图test graph中被完整记录。核心方案-amb参数与分步注意力计算PR #237 的解法是新增一个命令行参数用来指定可容忍的最大K*Q尺寸-amb size_in_MiB 或 --attn-max-batch MiB将其记为M_max。推理过程中在执行K*Q乘法之前先计算本次所需尺寸M若M ≤ M_max照常一次性完成K*Q → softmax → V*softmax(K*Q)若M M_max将V*softmax(K*Q)拆成n ⌈M / M_max⌉步执行M与M_max均按 MiB 取整。假设注意力头数为K则每一步只计算K/n个头每一步产生的K*Q张量缩小为原来的1/n。每一步与 V 相乘后得到的结果张量只含n_embd × n_tokens个元素相比超长上下文下的K*Q可以忽略不计。最后通过**拼接concatenation**把 n 步的V*softmax(K*Q)结果组装成完整输出。DeepSeekV3/R1 全量 163k 上下文下的量化分析PR 作者以 DeepSeekV3/R1、完整 163k 上下文、-amb 2048即 2 GiB为例给出了三组关键数字生成阶段TGu_batch 1K*Q尺寸为163000 × 128 × 1 × 4 79 MiB ≤ 2048 MiB计算照常进行无任何分步开销模型加载时的测试图u_batch 512K*Q尺寸为163000 × 128 × 512 × 4 40750 MiB因此拆成20 步执行每步只处理 6 或 7 个头。后端把K*Q张量的尺寸记录为 2 GiBcompute buffer 只需略大于 2 GiB再留出其他中间结果的空间实际提示词处理PPKV cache 中 token 少于 8k 时2 GiB 上限不会被触及8k–16k token 期间V*softmax(K*Q)拆 2 步16k–24k 拆 3 步依此类推。对于如此大的K和Q矩阵矩阵乘法的耗时远高于多启动几次矩阵乘法和 softmax 的开销因此对性能的影响可以忽略。源码级验证参数解析与计算图实现参数定义与命令行解析attn_max_batch首先在公共参数结构体中定义默认值为 256定义common/common.h中int attn_max_batch 256; // Max batch size to use when computing attention命令行解析common/common.cpp中-amb/--attention-max-batch分支std::stoi解析后若值在(0, 128)之间会被强制抬升到 128并打印警告amb %d is too low. Changing to 128帮助信息common/common.cpp中max batch size for attention computations (default: %d)上下文参数传递common/common.cpp中cparams.attn_max_batch params.attn_max_batch;并最终写入llama_cparams见src/llama-cparams.h与include/llama.h中int attn_max_batch; // maximum batch size for attention computations [EXPERIMENTAL]。从src/llama.cpp的上下文构建日志LLAMA_LOG_INFO(%s: attn_max_b %d\n, ...)可以看到加载时该值会以attn_max_b N的形式打印出来便于核对实际生效值。通用注意力路径的分步实现在通用非 MLA注意力路径中src/llama-build-context.cpp展示了完整的分步逻辑auto kq_size k-ne[1]*q-ne[1]*q-ne[2]*sizeof(float)/(1024*1024); if (cparams.attn_max_batch 0 || cparams.attn_max_batch kq_size || k-ne[2] ! q-ne[2] || v-ne[2] ! q-ne[2] || sinks) { // 常规路径一次性 K*Q - softmax - V*softmax(K*Q) } else { int n_step (kq_size cparams.attn_max_batch - 1)/cparams.attn_max_batch; n_step std::min(n_step, int(k-ne[2])); int n_per_step (q-ne[2] n_step - 1)/n_step; ... for (int i12 0; i12 q-ne[2]; i12 n_per_step) { // 视图切分 K/Q - K*Q - softmax - V*softmax(K*Q) // 首次迭代直接赋值后续迭代用 ggml_concat 拼接 } }关键细节kq_size以 MiB 为单位与-amb的单位一致当attn_max_batch 0未启用或 kq_size未超限时走常规一次性路径当k-ne[2] ! q-ne[2]GQA/MQA 场景、v-ne[2] ! q-ne[2]或存在 sink token 时为保证正确性也会回退到常规路径分步路径下每步通过ggml_view_3d切出子张量逐个头组完成mul_mat → soft_max_ext → mul_mat并用ggml_concat(ctx, kqv, kqv_i, 2)沿 head 维度拼接见 src/llama-build-context.cpp。DeepSeek MLA 路径的迭代注意力对于 DeepSeek 这类 MLAMulti-head Latent Attention模型分步逻辑体现在src/graphs/build_deepseek2.cpp中先根据wkv_b与 KV cache 计算kv_f32_sizeMiB若超过attn_max_batch则尝试寻找一个能被n_head整除的迭代数niter使kv_f32_size/niter attn_max_batch从而确定每轮处理的头数n_max_head见 src/graphs/build_deepseek2.cpp随后在for (int iter 0; iter n_head/n_max_head; iter)循环中逐轮执行flash_attn_ext并通过ggml_concat沿 head 维度拼装结果见 src/graphs/build_deepseek2.cpp。此外src/llama.cpp的llama_context::max_nodes会为 MLA 迭代注意力预留额外的计算图节点数当mla_attn 1、n_tokens 128、attn_max_batch 0且存在wkv_b时按每层每迭代 10 个节点的估算追加节点预算避免分步展开导致图节点数超限见 src/llama.cpp。从源码结构可以推断max_nodes的这套逻辑正是为 PR #237 的分步方案配套设计的。真实部署案例15×RTX 3090 全量跑满 163k 上下文PR 评论区中用户davidsyoung给出了极具参考价值的实测配置。在 15×RTX 3090合计 360GB 显存上以 Q2_K 量化约 230GB的 DeepSeek-R1 模型完整上下文启动命令为-m /models/gghfez_DeepSeek-R1-11446-Q2_K/DeepSeek-R1-11446-Q2_K-00001-of-00030.gguf -mla 2 -fmoe -b 2048 -ub 1024 -amb 1024 --tensor-split 37,25,25,25,24.5,24,24,24,24,25,24,25,24,24.5,31 --temp 0.5 --ctx-size 163840 --seed 3407 --n-gpu-layers 100 --host 0.0.0.0 --port 8080注意这里把-ub从默认的 512 提高到了1024、-amb设为1024 MiB——这正是 PR 带来的关键红利因为K*Q缓冲区上限被-amb锁死u-batch 可以放心调大prefill 吞吐随之大幅提升。加载日志中的几个关键行印证了效果n_ctx 163840、n_batch 2048、n_ubatch 1024、mla_attn 2、attn_max_b 1024、fused_moe 1模型元数据deepseek2.attention.head_count 128、context_length 163840、model params 672.050 B、model size 227.689 GiB (2.910 BPW)KV cache每卡 720 MiBCUDA0 为 1080 MiB总计约 10980 MiBcompute buffer每卡仅 5088 MiBCUDA0–CUDA13 为 5088.02 MiBCUDA14 为 5088.03 MiB外加CUDA_Host compute buffer size 2588.05 MiB——相比不使用-amb时动辄数十 GB 的K*Q缓冲缩减幅度非常可观。用户报告的实际吞吐-ub 1024下prompt eval time 3086.56 ms / 645 tokens ( 4.79 ms per token, 208.97 tokens per second) generation eval time 91155.54 ms / 1587 runs ( 57.44 ms per token, 17.41 tokens per second)prompt eval time 105435.60 ms / 14645 tokens ( 7.20 ms per token, 138.90 tokens per second) generation eval time 86701.60 ms / 1041 runs ( 83.29 ms per token, 12.01 tokens per second)从这些记录看prefill 保持在 138–212 t/s 区间即使 KV cache 内 token 已超过 14k、进入多步分阶段也未见明显崩塌用户还提到换回-ub 512时 prefill 约降到 120–140 t/s从侧面说明-amb释放了提高 u-batch 的空间。作者同时补充分步计算只影响注意力乘法的执行方式对 CPU 与 GPU 推理均适用。过程中的两个坑GGML_OP_CONCAT与后端拷贝PR 作者在实现中踩了两个值得记录的坑CUDA 端GGML_OP_CONCAT的 bug拼接各步结果依赖ggml_concat而 CUDA 实现当时存在缺陷。PR 顺带修复了所需用例连续张量、第二个张量直接追加到第一个末尾这也是该 PR 标题之外附带的价值后端对先分配最终结果、再按偏移写入方案的抗拒作者曾试图避免拼接所需的2n次拷贝改为预分配最终结果、把各步结果直接存到对应偏移处但后端因此崩溃在空指针访问上最终只能回归每步计算 concat 拼接的稳妥方案。参数速查与使用建议参数说明默认值备注-amb, --attn-max-batch MiB注意力计算中单个K*Q张量允许的最大尺寸MiB256传入(0, 128)会被强制修正为 1280表示不限制-ubu-batchmicro-batch大小512在-amb约束下可放心调大以提升 prefill-mlaMLA 注意力模式如2—与-amb配合作用于 DeepSeek 系模型-faFlash Attention 开关onattn_max_batch注释标明仅适用于flash_attn false的路径见 common/common.h使用建议基于 PR 描述与实测案例归纳目标值是让 compute buffer 与 KV cache 之和能塞进显存因此-amb应结合--ctx-size、量化等级与显存总量一起权衡典型取值在 1024–2048 MiB 之间不要为了省显存把-ub调得过小如 64那会付出约 3 倍 prefill 减速的代价正确姿势是用-amb锁住K*Q峰值再尽量调大-ub若上下文较短、K*Q本身不会触顶分步逻辑不会触发-amb不会带来额外开销观察启动日志中的attn_max_b N与CUDAx compute buffer size即可验证参数是否生效。总结PR #237 解决了KV cache 已压缩、compute buffer 却仍是超长上下文拦路虎的最后一公里问题。其核心贡献可归纳为三点一是引入-amb / --attn-max-batch参数把注意力K*Q张量的峰值显存需求从随上下文线性增长变为由用户显式封顶二是以按头分组、多步执行、ggml_concat拼接的方式实现 CPU/GPU 通用的分步注意力且对矩阵乘法为主的注意力计算几乎没有性能损失三是顺带修复了 CUDA 端GGML_OP_CONCAT的缺陷。对希望在有限显存上跑满 163k 上下文 DeepSeekV3/R1 等大模型的部署者而言-amb与-ub的组合已成为一份性价比极高的标准配置。【免费下载链接】ik_llama.cppllama.cpp fork with additional SOTA quants and improved performance项目地址: https://gitcode.com/GitHub_Trending/ik/ik_llama.cpp创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考