1. 项目缘起:为什么要在 MCU 上跑大模型
把一个大语言模型塞进一块微控制器里,这件事放在两三年前说出来,大概率会被同行当成玩笑。毕竟主流认知里,LLM 推理至少得有一块像样的 GPU,或者退一步也得是带几十 GB 内存的服务器 CPU。但嵌入式圈子这几年有个明显趋势:芯片的算力和内存带宽在快速上探,而模型侧的量化和剪枝技术也在同步成熟,两条曲线一交叉,就出现了"在 MCU 上跑 LLM"这个过去看起来不现实的场景。
我这次折腾的对象是ESP32-P4。这颗芯片值得单独说两句:它是乐鑫在 ESP32 家族里定位偏高性能的一颗,核心是RISC-V架构,主频能跑到 400MHz,片上带了相当可观的 SRAM,还支持外挂 PSRAM 做内存扩展。更关键的是它带了一个PIE(Processor Instruction Extension)协处理器,专门用来加速向量和定点运算——这个东西在后面优化推理速度的时候是绝对的主角。
项目标题里那个数字对比很扎眼:从 0.61 tok/s 到 4.31 tok/s,整整 7 倍。这个系列总览要讲的就是这 7 倍是怎么一点点抠出来的。我先给个结论性的判断:这 7 倍不是靠某一个"银弹"实现的,而是把推理链路上每一个环节的浪费都找出来、逐个消掉的结果。从模型量化格式的选择,到算子实现,到内存布局,到 PIE 指令的利用,每一块都贡献了一部分。
这篇文章适合谁看?如果你是对MCU上做 AI 推理感兴趣的嵌入式工程师,或者你在做边缘侧的LLM部署、想搞清楚端侧推理的性能瓶颈到底在哪,再或者你只是好奇一块几块钱到几十块钱级别的芯片到底能跑出什么水平,那这篇复盘应该都能给你一些可以直接抄作业的东西。我会尽量把每个决策背后的"为什么"讲清楚,而不是只丢一堆参数出来。
需要提前说明的是,这个系列会拆成多篇,本篇是总览,负责把整体思路、关键节点和踩过的坑先铺开,具体的代码级细节会在后续分篇里展开。所以你会看到这里既有宏观的架构决策,也有具体的数字和实测记录,但不会陷入某一行汇编的细节里。
2. 整体方案设计:从模型到硬件的全链路拆解
2.1 先想清楚瓶颈在哪,再动手
很多人一上来就想着"怎么把模型跑得更快",然后开始盲目地换量化格式、调编译选项。我踩过的第一个坑就是这个——在没有 profiling 的情况下瞎优化,结果花了两天时间优化了一个只占总耗时 3% 的环节。
正确的做法是先建立性能模型。LLM 推理在 MCU 上的耗时,粗略可以拆成三块:权重读取的内存带宽开销、矩阵乘法的计算开销、以及注意力机制里的额外开销(比如 softmax、KV cache 的读写)。在 ESP32-P4 这种级别的芯片上,绝大多数情况下瓶颈是内存带宽,而不是纯算力。原因很简单:模型权重动辄几十 MB,而片上 SRAM 只有几百 KB 到几 MB,大部分权重得从 PSRAM 甚至外部 Flash 里读,读取速度直接决定了 token 生成的速度。
这个判断非常重要,因为它决定了优化方向。如果瓶颈是算力,那你要做的是优化算子、用 SIMD;如果瓶颈是带宽,那你要做的是减少数据搬运、提高缓存命中率、用更紧凑的量化格式。实测下来,ESP32-P4 上这两者都有,但带宽是主要矛盾。
2.2 模型选型:为什么是小模型而不是"缩小的大模型"
在 MCU 上跑,模型规模必须严格控制。我最终选的是一个参数量在百万级到千万级之间的小模型,具体规模会根据量化后的内存占用反推。这里有个经验:不要指望把一个 7B 的模型量化到 4bit 就能塞进 MCU,即使塞进去了,推理速度也会慢到没有实用价值。
选型的核心逻辑是"内存占用优先"。假设你有 8MB 的 PSRAM 可用,模型权重加上 KV cache 加上运行时开销,实际能留给权重的可能只有 5-6MB。按 4bit 量化算,大概能放 1000 万参数左右;按 8bit 算,就只有 500 万参数。这个约束是硬的,绕不过去。
提示:模型选型阶段一定要先算内存账,再算算力账。很多人反过来,先看模型效果好不好,结果发现根本放不下。
2.3 量化格式的取舍:Q4 还是 Q8
量化格式的选择直接决定了内存占用和推理速度的平衡。我实测对比过几种方案:
| 量化格式 | 权重内存占用 | 相对推理速度 | 输出质量 | 适用场景 |
|---|---|---|---|---|
| FP32 | 基准 100% | 基准 1.0x | 最好 | 不现实,仅作对照 |
| INT8 | 25% | 约 2.5-3x | 接近 FP32 | 内存充裕时首选 |
| INT4 | 12.5% | 约 4-5x | 略有下降 | 内存紧张时首选 |
| 混合精度 | 15-20% | 约 3-4x | 较好 | 关键层用 INT8 |
最终我采用的是INT4 为主、关键层保留 INT8的混合方案。为什么不全用 INT4?因为实测发现,注意力层的 Q/K/V 投影如果用 INT4,输出质量下降比较明显,而这几层的参数量占比其实不大,保留 INT8 对内存影响有限,但对质量提升明显。这是一个典型的"把好钢用在刀刃上"的取舍。
2.4 软件栈的整体架构
整个推理栈我分成了四层,从下到上依次是:
- 硬件抽象层:封装 PIE 指令、DMA、PSRAM 访问,屏蔽底层细节
- 算子层:实现矩阵乘法、softmax、LayerNorm 等核心算子,针对 PIE 做优化
- 推理引擎层:负责 KV cache 管理、token 调度、采样策略
- 应用层:提供简单的对话接口,方便测试
这样分层的好处是,优化的时候可以精确定位到某一层,不会牵一发而动全身。比如后面做 PIE 优化,主要改的是算子层,推理引擎层基本不用动。
3. 核心优化手段:7 倍是怎么抠出来的
3.1 第一刀:内存布局重排,拿到约 1.8 倍
最开始跑通的时候是 0.61 tok/s,慢得让人怀疑人生。第一个优化点是内存布局。
原始实现里,权重是按"行优先"存储的,矩阵乘法时按行读取。但 PIE 协处理器做向量运算时,更擅长处理连续的内存块。我把权重重新排列成"按列分块"的布局,让每次 PIE 加载的数据都是连续的,减少了内存访问的碎片化。
这个改动听起来简单,但效果立竿见影:从 0.61 提到了约 1.1 tok/s。为什么?因为原来每次读权重都要跳着读,PSRAM 的突发传输优势完全发挥不出来。重排之后,连续读取让 PSRAM 的带宽利用率从大概 30% 提到了 60% 以上。
注意:内存重排要在模型转换阶段做,不要放在运行时。运行时做重排会引入额外的拷贝开销,得不偿失。
3.2 第二刀:PIE 指令加速矩阵乘法,拿到约 2.2 倍
这是整个优化里贡献最大的一块。ESP32-P4 的 PIE 协处理器支持 SIMD 风格的定点运算,一次能处理多个数据。我针对 INT4 和 INT8 分别写了专门的矩阵乘法内核。
关键点在于数据打包。INT4 的数据是 4bit 一个,两个才能凑成一个字节。PIE 做运算时,需要先把这些 4bit 数据解包成 8bit 或 16bit 的中间格式,再做乘加。这个解包过程如果处理不好,会成为新的瓶颈。我的做法是用查表法配合位运算,把解包和乘加融合在一起,减少中间数据的搬运。
实测下来,矩阵乘法这一块的耗时从占总时间的 65% 降到了 35% 左右,整体速度从 1.1 提到了约 2.4 tok/s。
3.3 第三刀:KV cache 优化,拿到约 1.5 倍
KV cache 是 LLM 推理里一个容易被忽视的性能杀手。随着生成的 token 越来越多,KV cache 会不断增长,每次生成新 token 都要读取整个 cache。在内存带宽本来就紧张的情况下,这是个不小的负担。
我做了两件事:一是把 KV cache 也用 INT8 量化存储,直接砍掉一半的读取量;二是把 cache 放在片上 SRAM 里而不是 PSRAM,因为 SRAM 的访问延迟低得多。当然,SRAM 容量有限,所以只放最近的一部分,更早的用滑动窗口策略丢弃。
这一刀下去,从 2.4 提到了约 3.6 tok/s。
3.4 第四刀:算子融合与循环展开,拿到约 1.2 倍
最后这一刀是精细活。我把一些相邻的算子做了融合,比如把 LayerNorm 和后面的矩阵乘法合并,减少中间结果的写回。同时对内层循环做了展开,让 PIE 的流水线能跑得更满。
这部分优化比较琐碎,单个改动效果都不大,但累积起来从 3.6 提到了最终的 4.31 tok/s。
3.5 优化效果汇总
| 优化阶段 | 速度 (tok/s) | 相对上一阶段提升 | 累计提升 |
|---|---|---|---|
| 初始版本 | 0.61 | - | 1.0x |
| 内存布局重排 | 1.10 | 1.80x | 1.80x |
| PIE 矩阵乘法 | 2.40 | 2.18x | 3.93x |
| KV cache 优化 | 3.60 | 1.50x | 5.90x |
| 算子融合 | 4.31 | 1.20x | 7.07x |
这张表是整个项目的核心成果。可以看到,PIE 优化贡献最大,但其他几项加起来也占了将近一半的提升。这也印证了我一开始的判断:优化是个系统工程,没有单点银弹。
4. 实操过程中的关键细节与踩坑记录
4.1 环境搭建与工具链选择
工具链这块我用的是乐鑫官方的 ESP-IDF,版本选的是比较新的稳定版。编译器是 RISC-V 的 GCC 工具链。这里有个小坑:不同版本的 GCC 对 PIE 指令的支持程度不一样,有些版本生成的代码会莫名其妙地慢。我建议锁定一个验证过的版本,不要频繁升级。
调试方面,我用的是 JTAG 调试配合串口日志。JTAG 能看寄存器和内存,串口日志用来打时间戳。两者结合,基本能定位到大部分性能问题。
4.2 内存分配的坑
ESP32-P4 的内存分好几块:内部 SRAM、外部 PSRAM、Flash。它们的访问速度差异很大。我一开始没注意,把权重全放在 PSRAM 里,结果发现 PSRAM 的访问延迟比 SRAM 高一个数量级。
后来我做了分级:最频繁访问的数据放 SRAM,次频繁的放 PSRAM,只读的权重放 Flash 并开启缓存。这个分级策略对性能影响很大,值得单独花时间调。
提示:ESP-IDF 提供了内存分配 API,可以指定分配到哪块内存。用之前一定要看清楚文档,别默认分配。
4.3 数值精度的坑
量化之后,数值精度问题会集中爆发。我遇到的最典型的问题是累加溢出。INT8 乘 INT8 的结果是 INT16,但如果累加很多项,INT16 也会溢出。解决办法是用 INT32 做累加器,虽然多占一点寄存器,但能避免溢出导致的输出乱码。
另一个坑是量化参数的校准。如果校准数据选得不好,量化后的模型输出会明显变差。我的经验是用一批有代表性的输入做校准,不要只用一两条。
4.4 常见问题速查表
| 问题现象 | 可能原因 | 排查方向 | 解决方法 |
|---|---|---|---|
| 输出乱码 | 累加溢出 | 检查累加器位宽 | 改用 INT32 累加 |
| 速度远低于预期 | 权重在 PSRAM | 检查内存分配 | 热点数据移到 SRAM |
| 生成到一半卡死 | KV cache 越界 | 检查 cache 索引 | 加边界检查,用滑动窗口 |
| 输出质量差 | 量化校准不当 | 检查校准数据 | 换更有代表性的校准集 |
| 编译报错 | 工具链版本不匹配 | 检查 GCC 版本 | 锁定验证过的版本 |
4.5 实测心得
跑通之后我做了几轮压力测试,连续生成几百个 token,观察速度是否稳定。实测发现,随着生成长度增加,速度会有轻微下降,主要是 KV cache 增长导致的。用滑动窗口策略后,速度基本能保持稳定。
另外,温度对性能也有影响。芯片跑久了会发热,如果散热不好,可能会触发降频。做长时间测试的时候要注意这一点。
5. 这个项目还能怎么扩展
跑通只是起点。基于现在这套框架,我看到几个可以继续深挖的方向。
一是多模型切换。现在的实现是单模型硬编码,如果做成模型可插拔的架构,就能根据任务复杂度动态选择不同规模的模型,简单任务用小模型快速响应,复杂任务用大模型保证质量。
二是更激进的量化。现在用的是 INT4/INT8 混合,如果试试 2bit 甚至 1bit 的极端量化,配合更好的校准方法,也许能在可接受的质量损失下进一步压缩内存。
三是算子层面的持续优化。PIE 的能力我可能只用了六七成,还有一些指令没充分利用。如果能把手写汇编的水平再提一提,矩阵乘法那块还有空间。
四是和上层应用结合。现在只是个推理引擎,如果能接上语音识别、传感器数据这些输入,就能做成一个完整的端侧智能应用。这才是 MCU 跑 LLM 真正的价值所在——不是替代云端大模型,而是在离线、低功耗、低成本的场景里提供够用的智能。
我个人在实际操作中的体会是,端侧 LLM 这个方向现在处于一个很微妙的阶段:硬件刚够用,软件还不成熟,但正因为不成熟,才有大量可以优化的空间。这 7 倍的提升里,我相信还有不少水分可以挤。如果你也在做类似的事情,欢迎一起交流踩坑经验。