资讯动态

MNN LLM 的 attention_mode 怎么配,兼顾 FlashAttention 与 KV Cache 量化?

发布时间:2026/9/15 15:06:40 来源:尧图企业网站定制
MNN LLM 的 attention_mode 怎么配兼顾 FlashAttention 与 KV Cache 量化【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN在 MNN 上跑 LLM 推理时KV Cache 占用和 attention 算子的内存峰值是长上下文场景下的两个主要压力点。MNN 用一个整数配置项attention_mode同时控制这两件事高位是 FlashAttention 开关低位是 KV Cache 量化方式。本文基于 docs/transformers/llm.md 中“推理配置”和“LLM Benchmark 工具使用”两节的说明给出attention_mode的编码规则、按后端与模型规模的取法、写入config.json的方式以及用llm_bench做 A/B 验证的命令。前提是 MNN 已按文档开启 LLM 支持编译如../build_64.sh -DMNN_BUILD_LLMON或sh package_scripts/ios/buildiOS.sh -DMNN_BUILD_LLMON并拿到了llm_demo、llm_bench和导出好的模型目录。attention_mode 的编码规则attention_mode的编码规则为attention_mode flash_attention * 8 kv_quant_mode默认为8。旧的quant_qkv选项已废弃文档明确提示改用attention_mode。两个分量分别取值KV Cache 量化模式attention_mode % 8值含义0不量化key 和 value 均为 fp161key 使用 int8 量化value 不量化2key 和 value 均使用 int8 量化3key 使用 TQ33-bit量化value 不量化4key 和 value 均使用 TQ33-bit量化5key 使用 TQ44-bit量化value 不量化6key 和 value 均使用 TQ44-bit量化FlashAttention 开关attention_mode / 80表示不使用1表示使用。怎么选先看后端再看内存压力与模型规模文档给出的常用配置及定位如下attention_mode组合文档定位8FlashAttention不量化默认推荐10FlashAttention KV-INT8精度几乎无损12FlashAttention KV-TQ33-bit 量化极致压缩推荐 4B 及以上参数模型14FlashAttention KV-TQ44-bit 量化内存节省大于 30%推荐 4B 及以上参数模型按以下条件取舍默认或精度优先保持8FlashAttention 开启、KV 保持 fp16。想省内存但不牺牲 FlashAttention取10KV-INT8文档说明其精度几乎无损。需要进一步压缩 KV 内存取12或14TQ3/TQ4。文档说明 TQ3/TQ4 基于 TurboQuant 算法WHT 旋转 Lloyd-Max 码本建议在 4B 及以上参数模型上使用小模型1B精度损失较大。后端限制是选值时的硬约束Metal 后端的 KV Cache 量化目前仅支持 int8 通道即attention_mode % 8只能取0/1/2TQ3/TQ4 为 CPU 专属。Metal 上文档给出的常用组合是8长上下文默认推荐和10进一步降低 KV 显存。其他 GPU 后端OpenCL / Vulkan / CUDA目前仅遵循attention_mode中的默认行为不支持 FlashAttention 切换。Metal 的 FlashAttention 作用于 prefill 阶段启用融合 Flash Attention kernel 后跳过mTempQK/mTempSoftMax两块 O(seq²·B·H) 中间显存长上下文场景峰值内存显著下降。当前 kernel 支持 head_dim ∈ {64, 128, 256}、GQA group ∈ {1, 2, 4, 8}、causal ADD-mask 输入。如何写入配置运行时的模型目录中有一个config.json与llm.mnn、llm.mnn.weight、tokenizer.mtok等文件同目录attention_mode属于其中的推理配置直接加一项即可{ backend_type: cpu, thread_num: 4, attention_mode: 14 }上面示例表示开启 FlashAttention 并将 KV Cache 做 TQ4 量化只保留了完成本场景所需的最小字段模型文件字段llm_model、llm_weight等可按文档默认值省略。用llm_demo加载该目录验证配置是否生效./llm_demo model_dir/config.json prompt.txtC 代码路径下引擎同样从该 config 读取attention_mode见 llm.cpp 中对attention_mode的解析也可以用运行时的set_config合并配置例如文档中 iOS 侧的用法llm-set_config({\tmp_path\:\...\});用 llm_bench 做 A/B 验证改配置前最好有基线数据。llm_bench提供-qa和-fa两个独立参数控制这两个分量最终attention_mode flash * 8 quant-qa, --quant-attentionKV Cache 量化模式默认0不量化。取值与上表attention_mode % 8一致-fa, --flash-attention1开启 Flash Attention默认0关闭。文档给出的例子-qa 0 -fa 1等效于attention_mode8-qa 2 -fa 0等效于attention_mode2。因此llm_bench命令参数与config.json里的attention_mode可以直接换算对照。文档给出的 A/B 对比命令在 build 目录下运行./model/config.json换成你的模型 config 路径# Flash Attention 开启默认 ./llm_bench -m ./model/config.json -a metal -p 512 -n 128 -rep 5 # Flash Attention 关闭 ./llm_bench -m ./model/config.json -a metal -fa 0 -p 512 -n 128 -rep 5 # 开启 KV-INT8 量化KeyValue ./llm_bench -m ./model/config.json -a metal -qa 2 -p 512 -n 128 -rep 5测试模式方面-p 512是 prefill-only 测试表格中记为pp512-n 128是 decode-only 测试tg128-pg 512,128是先 prefill 512 个 token 再复用该 KV cache 生成 128 个 token 的联合测试pp512tg128。文档特别提醒评估真实推理性能尤其是长上下文下 decode 随 KV 增长变慢的效应应使用-pg它是唯一让 decode 阶段带着真实 KV cache 长度运行的模式-n得到的是 KV≈0 的理想值通常明显偏高。llm_bench的输出形如以下为文档中的示例结果不是固定预期| model | modelSize | backend | threads | precision | test | speed(tok/s) | | ------------------ | ---------: | ------- | ----: | ---------- | ------------- | ------------ | | qwen3-0.6b | 355.68 MiB | METAL | 4 | Low | pp512 | 5073.12 ± 8.61 | | qwen3-0.6b | 355.68 MiB | METAL | 4 | Low | tg128 | 345.85 ± 1.02 | | qwen3-0.6b | 355.68 MiB | METAL | 4 | Low | pp512tg128 | 5061.44 ± 9.30br293.40 ± 0.42 |判断方式对比开启/关闭 FlashAttention、不同-qa值各跑的ppN与tgN/ppNtgM速度行确认切换attention_mode后哪个阶段的速度变化符合预期例如 Metal 上 KV-INT8 对 prefill 的影响。结果也可以加-fp ./test_result以 markdown 追加写入文件便于多次对比。Metal 后端还有一个调试开关环境变量MNN_ENABLE_FLASH_ATTN_PREFILL1可强制开启 FA无视 config0强制关闭文档说明其用途是 A/B 基准。已知的权衡与边界Metal 上 KV-INT8 的代价文档说明attention_mode10时“短上下文 pp 会略降 5–14%”即 prefill 阶段速度有一定回退换取 KV 显存下降。如果主要瓶颈不在内存8更稳妥。TQ3/TQ4 只在 CPU 后端生效且面向 4B 及以上参数模型1B 小模型文档明确提示精度损失较大。OpenCL / Vulkan / CUDA 后端不支持 FlashAttention 切换在这类后端上调整attention_mode的高位不会有文档承诺的效果。配置入口已统一为attention_mode如果旧配置里还写着quant_qkv按文档应替换为attention_mode否则两处不一致时以attention_mode为准引擎解析顺序见 llm.cpp 中对quant_qkv的兼容读取。完成llm_benchA/B 验证后把测得最优的attention_mode值写回模型目录的config.json即可让llm_demo、CLlm接口和 PythonMNN.llm的所有入口使用同一份配置。【免费下载链接】MNNMNN: A blazing-fast, lightweight inference engine battle-tested by Alibaba, powering high-performance on-device LLMs and Edge AI.项目地址: https://gitcode.com/GitHub_Trending/mn/MNN创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价