56M 参数大模型在三块 ESP32-S3 上分布式推理

一位开发者将一个 56M 参数的微型语言模型拆分到三块 ESP32-S3 N16R8 微控制器开发板上，通过 ESP-NOW 无线协议协同完成端到端文本生成。项目灵感来自单板运行 TinyStories 的方案，作者进一步将其扩展为多板分布式架构，并加入了 KV 缓存，使模型第一次能按设计方式运行带上下文的注意力计算。

项目概览

模型在 PC 端基于 WikiText-103 语料训练，采用了来自 Google Gemma 架构的 Per-Layer Embeddings（PLE）设计，并做了 4-bit group-wise 量化。整张约 50.3M 参数的 PLE 表被拆成两半，分别放在两块板上以适配每板 16 MB flash 的限制。最终生成的模型规模约 56M 参数，能在浏览器中实时看到逐 token 流式输出。

三板分工与推理流程

三块板各有明确职责：

• 板 A：负责 BPE 分词、token 嵌入查询、PLE 投影、输出归一化与 LM head，承担「入口」与「出口」角色。
• 板 B：承担核心 transformer 计算，6 层注意力 + 6 层 FFN，并维护 KV 缓存。
• 板 C：保存 PLE 表的另一半、做采样，并通过自带 WiFi AP 托管网页服务器，对外提供 SSE 流式接口。

端到端流程是：浏览器通过 WiFi 连入板 C，发送 prompt；板 C 经 ESP-NOW 转发给板 A；板 A 完成分词与本地投影，再向板 C 请求 PLE 表行；接着把组合后的表示发给板 B，板 B 跑完 transformer 后返回隐藏状态；板 A 计算 logits 并把 top-40 token 发回板 C；板 C 采样、解码并通过 SSE 推回浏览器。

Split-PLE 与 KV 缓存设计

为了把 PLE 表塞进 16 MB flash，作者把它拆成两个 128 维半张：

• PLE_A（约 12.5 MB，4-bit group=64）放在板 C，运行时由板 A 通过 ESP-NOW 按需请求。
• PLE_B（约 12.6 MB，4-bit group=128）常驻板 B，本地使用。

板 A 只保留很小的 ple_model_proj（约 50 KB）与 ple_proj_norm，并把原本放在板 A 的 12.5 MB PLE 表挪到板 C，腾出的空间用来把 tok_emb 升到 8-bit 以改善推理质量。

KV 缓存则是项目里质量提升最显著的一处改动。此前每个 token 都以 seq_len=1 通过板 B 的注意力，模型实际上看不到历史信息；改造后板 B 在 PSRAM 中维护约 1.5 MB 的 KV 缓存，最多支持 256 个位置，RoPE 也使用真实位置索引。这是项目自述中提到的「单一最大质量改进」。

通信协议与项目结构

板间通信使用乐鑫 ESP-NOW 点对点协议，无需路由器，单包最大 250 字节并自动分片，作者在协议层加了序号与 ACK，估计单次往返延迟在 1–5 ms 区间。项目仓库目录清晰地分成了两大部分：

• src/：PC 端的训练与导出流水线，包含 PyTorch 模型定义、数据加载、训练循环、量化与多板固件导出脚本。
• firmware/：基于 Arduino IDE 的三块板固件，共享 llm.h 分布式 C 推理运行时与 espnow_protocol.h 协议定义，每块板各自一个 sketch 与分区表。

整体而言，这是一个把「在 MCU 上跑 LLM」从单板演示推到多板协同、并认真处理了 KV 缓存与权重复用等关键工程问题的开源作品，适合对嵌入式端侧推理感兴趣的开发者参考。