注意力缓冲)
如何理解PicoLM的Flash Attention在线Softmax节省O(n)注意力缓冲【免费下载链接】picolmRun a 1-billion parameter LLM on a $10 board with 256MB RAM项目地址: https://gitcode.com/gh_mirrors/pi/picolmPicoLM 是一个纯 C、零依赖的LLM 推理引擎能在 256MB 内存的 $10 开发板上运行 10 亿参数大模型。本文带你搞懂 PicoLM 如何实现Flash Attention通过在线 SoftmaxOnline Softmax单遍计算注意力彻底去掉 O(seq_len) 的注意力分数缓冲——这正是小内存设备上长上下文推理的关键优化。问题诊断传统注意力为什么浪费内存标准 Transformer 注意力要分3 遍完成打分开对每个历史位置t计算score[t] dot(Q, K_t) / sqrt(d)Softmax 归一化遍历全部分数求指数和加权求和用归一化权重对 Value 向量加权累加三遍计算意味着必须先把完整分数数组落到内存大小为n_heads × max_seq_len个 float。对 TinyLlama1.1B就是约64KB的att[]缓冲且序列越长占用越大O(n)更糟的是要对 KV 缓存做三次完整读取缓存命中差。对内存仅剩 45MB 的 PicoLM 来说每一 KB 都要精打细算 。核心思路Online Softmax 单遍求注意力在线 Softmax 的洞察Softmax 的归一化其实可以边算边改。只需维护三个量状态量作用max_score已见过的最大分数运行时最大值sum_exp指数权重累计和acc[]按权重累加的 Value 累加器算法流程伪代码级理解即可遍历每个缓存位置t算出当前分数score若score max_score出现新最大值用修正因子exp(旧max - 新max)把之前所有累加按比例缩回再计入当前项否则直接以权重exp(score - max_score)累加全部结束后输出 acc[] / sum_exp与标准 Softmax数值完全等价数学上等价、且只计算exp(x)x ≤ 0数值稳定性还有保障——不会溢出。源码走读PicoLM 在哪里实现 Flash Attention核心实现位于前向传播的注意力阶段注释里写明了整个单遍算法Flash Attention 主循环model.c —— 每个 Query 头维护max_score / sum_exp / acc[]逐位置读取 FP16 KV 缓存并增量更新缓冲删除注释model.c 标注att buffer removed (flash attention)结构体定义处的移除说明model.h ——att buffer REMOVED — flash attention uses online softmax深度技术博客含逐行讲解BLOG.md 中 Optimization 6: Flash Attention / Online Softmax 一节优化总览README.md 的 5. Flash Attention (Online Softmax)配合GQA 分组查询注意力32 个 Q 头共享 4 个 KV 头见 model.c单头只需遍历head_dim大小的 K/V 切片进一步缩小内存带宽压力。优化收益省了多少、快在哪收益项效果内存彻底删除att[]分数缓冲TinyLlama 省 64KB模型越大省得越多访存KV 缓存1 遍读取原来是 3 遍缓存行为显著更好数值与标准 Softmax 注意力数值等价不影响输出质量长上下文注意力开销不再产生 O(seq_len) 额外分配2048 上下文更从容在 256MB RAM 的预算表里见 README.md Memory Budget 一节运行时内存被压到约45MB2048 上下文——Flash Attention 是这 9 项优化中唯一直接砍掉随序列增长分配的成员。与其他优化的配合Flash Attention 并非孤军奋战PicoLM 的 9 项优化互相咬合FP16 KV 缓存model.hK/V 用 16 位浮点存储缓存体积减半Flash Attention 的fp16_to_fp32()在点积时即时转换预计算 RoPE 表查表替代 2.5 万次三角函数调用融合反量化 点积矩阵乘内存流量降约 50%最终 x86 上生成速度从 1.6 tok/s 提到13.5 tok/sARM NEON 预期更高数据来自 BLOG.md 性能汇总。新手上手3 分钟跑起来git clone https://gitcode.com/gh_mirrors/pi/picolm cd picolm/picolm make native # 树莓派用户: make pi make model # 下载 TinyLlama 1.1B Q4_K_M638MB构建脚本见 Makefile一条命令安装见 install.sh。一句话总结Flash Attention 的本质是用运行时最大值 修正因子把三遍注意力压成一遍。PicoLM 用不到 40 行 C 代码实现它换掉 64KB 缓冲、砍掉两次 KV 扫描——这就是极简 LLM 推理引擎的生存之道。【免费下载链接】picolmRun a 1-billion parameter LLM on a $10 board with 256MB RAM项目地址: https://gitcode.com/gh_mirrors/pi/picolm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考