ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

03 ·纯 C11 在 MCU 上写 Transformer 推理:无 SIMD 的标量内核全解析

03 ·纯 C11 在 MCU 上写 Transformer 推理:无 SIMD 的标量内核全解析 03 · 纯 C11 在 MCU 上写 Transformer 推理无 SIMD 的标量内核全解析English version:en/03-scalar-inference-kernel.md本篇对应源码main/kmcu.c·main/kmcu.h·main/main.c目标理解kmcu.c/h如何在一个 32 位 RISC-V MCU 上、用纯标量 C跑通Mistral 家族的 Transformer decode以及为什么 32 MB PSRAM 能装下 12 MB 权重。1. 整体架构main.c入口 ├─ shs_selftest() SHS 公理算子自检S_0/UPA ├─ probe_psram() PSRAM 容量/带宽探针 ├─ probe_sdcard() TF 卡探针默认关 ├─ pie_bench() PIE 单算子对拍 测速 └─ model_test() 模型加载 decode kmcu.c/h核心推理 ├─ km_open() 解析 KMCU header 目录 ├─ km_fast_build_ex() 把 q4 从 Flash 分块展开为 int8 fp32 scalePSRAM ├─ km_decode_step() 单步 decode输入 token → 输出下一个 token ├─ gemv_q8() GEMVint8 权重 × 激活 ├─ rms_norm()/rope_apply() 归一化 / 位置编码 └─ q8_row_get() 读嵌入行fp32 反量化2. 两条权重路径路径说明慢路径直接从 q4 原图逐元素反量化正确性基准已不用于 decode快路径装载时把 q4展开为int8 qfp32 scale d默认快路径的核心结构typedefstruct{constint8_t*q;/* 展开后的 int8行对齐行尾含零填充 */constfloat*d;/* 每 32 元素一块的 fp32 scale块号 r*nblk b */constfloat*f32;/* 非 q4 张量norm的 fp32 副本 */uint32_tn;uint8_tis_q4;}km_ft_t;为什么「分块展开」而不是整份载入12.26 MB 镜像 23 MB 展开后的 arena会超过 32 MB PSRAM 预算。所以用km_read_fn回调从 Flash 分块读逐块展开展开完就释放 Flash 原图typedefint(*km_read_fn)(void*ctx,uint32_toff,void*dst,uint32_tlen);快路径不引入新的量化近似q和d完全来自原 q4 数据只把半字节提取从「热循环」移到「装载期」。3. decode 单步流程km_decode_step()输入一个 token id输出下一个 token贪心 argmax输入嵌入: x tok_embed[token] (q8_row_get, fp32) for L in 0..n_layers: ├─ h RMSNorm(x, in_ln) ├─ q GEMV(q_proj, h); k GEMV(k_proj, h); v GEMV(v_proj, h) ├─ RoPE(q); RoPE(k) ├─ 写 KV cache ├─ attentionGQA causal softmax→ ao ├─ x GEMV(o_proj, ao) ├─ h RMSNorm(x, post_ln) ├─ g GEMV(gate, h); u GEMV(up, h) ├─ act silu(g) * u └─ x GEMV(down, act) h RMSNorm(x, final_norm) logits h tok_embed^Ttied lm_head return argmax(logits)各算子的标量实现要点RMSNormMistral 用 RMSNorm非 LayerNormssΣ x[i]^2/n out[i]x[i]/sqrt(sseps)*w[i]RoPE与 HF 一致out1 x1·cos − x2·sin; out2 x2·cos x1·sinforh in heads:fori in half:inv1/theta^(2i/hd)angpos*inv v[i]a·cos(ang)− b·sin(ang)v[ihalf]b·cos(ang)a·sin(ang)GQA attention12 个 q head 共享 4 个 kv headkv_rep 3每 kv head 服务 3 个 q head。softmax 前先减去 max数值稳定。SwiGLUact silu(gate) * upsilu(x) x / (1 exp(-x))。4. scratch 布局内存预算的关键decode 的工作区不含权重是一个scratch数组布局严格对齐[float 区] x[dim_p] h[dim_p] q[dim_p] kk[kvw] vv[kvw] ao[dim_p] g[ffn] u[ffn] act[ffn_p] [int16 区] xq[align32(max(dim,ffn))] ← PIE GEMV 的激活量化缓冲dim_p align32(dim)ffn_p align32(ffn)。对齐 padding 区在km_state_init里清零后不再被写保证 GEMV 尾部整块读取数值正确。KV cache 另算2 × n_layers × n_ctx × n_kv_heads × head_dim个 float。本例KV_CTX128时 KV cache 2×520 KB。5. M1 能力探针硬件能力基准接入模型前先摸清硬件能力这些数据不随量化对齐改变是稳定的硬件边界项实测值芯片ESP32-P4 rev v3.1双核 LP 核400 MHzPSRAM 容量32768 KB32 MBhex 模式 200 MHzPSRAM 写带宽83–84 MB/s16 MBu32 顺序写PSRAM 读带宽84 MB/su32 顺序/90 MB/s×4 展开/106–107 MB/s-O2下Flash16 MB NORDIO 80 MHz关键判读4 路展开只让读带宽 7%84→90说明 PSRAM带宽受限而非延迟受限。这直接决定了后面「纯计算 21× 的 PIE 被带宽墙压到 2.11×」。6. M2 实测结果未行对齐时点序列已作废项实测decode1.63 s/token0.61 tok/s19 步 31.0 sFlash→PSRAM 装载12.26 MB / 1.61 s 7.8 MB/sPSRAM 占用权重 12.26 MB KV 2×520 KB 目录算力利用率22.8M MAC / 1.63 s ≈ 14 MMAC/s相对 PIE 潜力 ~100× 余量M2 时点未行对齐的对拍数据single-token[1] top1: id684 logit5.229595 (PC: 5.229599, 差 4e-6) 4-token prompt top1: id23624 logit5.184734 (PC: 5.184731, 差 3e-6) 20-token 序列: 1,450,2217,4996,23624,1199,8752,23624,1199,20925,3161,4865,442,5102,2672,3161,13040,17574,28394,1628⚠️ 这条 20-token 序列是未行对齐的 q4 布局下的结果M3-2 引入行对齐后量化块边界改变序列已变。当前基准序列见 02 篇 的完整输出。瓶颈标量逐元素 q4 反量化每权重一次半字节提取 fp16 转换 输出头32002×312 10M元素的流式 argmax占 44%。7. TF 卡诊断结论被推翻的一次本板 TF 卡SDIO引脚main.c实测定义SD_PWR_GPIO 45 ← 负载开关低有效 SD_PIN_CLK 43 SD_PIN_CMD 44 SD_PIN_D0 39 D1 40 D2 41 D3 42第一次6 种组合供电极性 × 总线宽度 × 时钟全报ESP_ERR_TIMEOUT send_op_cond怀疑卡坏或电气层问题。第二次复测错误变为ESP_FAIL / failed to mount card (13) FatFSFR_NO_FILESYSTEM。卡电气链路其实正常有效组合GPIO450负载开关低有效只差 FAT 分区表。结论0.032B 常驻不依赖 TF 卡权重直接放 PSRAM/FlashTF 卡只在后续 0.1B 分层驻留才需要。TF 卡两个坑esp_vfs_fat_sdmmc_mount失败时内部已自行 deinit调用方再sdmmc_host_deinit()会双重 deinit → Instruction access fault 崩溃。默认SDMMC_SLOT_CONFIG_DEFAULT()的 width 允许 8-bit会把GPIO45 当 D4 抢走它正是 TF 卡电源开关必须显式slot.width 4。8. M2 阶段的踩坑测试脚手架 bug非模型 bug贪心循环曾把生成结果写回seq[0..PROMPT_LEN-1]覆盖 prompt导致 MCU 从未处理完整 prompt。修正为「先预处理整条 prompt 再自回归」后 20/20 一致。计算型任务单步 1.6 s 不喂狗需关任务看门狗CONFIG_ESP_TASK_WDT_ENn否则每步打寄存器转储。riscv32 上uint32_t是unsigned longprintf 必须用PRIu32或显式转unsigned。对应源码文件关键符号 / 位置支撑本文哪部分main/kmcu.ckm_open、km_fast_build_ex、km_decode_step、gemv_q8、rms_norm、rope_apply、q8_row_get第 1–4 节内核架构、两条权重路径与 decode 单步流程main/kmcu.hkm_ft_t、km_read_fn、km_state_init第 2–4 节快路径结构、分块读回调与 scratch 对齐main/main.cshs_selftest、probe_psram、probe_sdcard、model_test第 5、7 节 M1 能力探针与 TF 卡诊断host_verify.ckm_open、km_decode_step、pie_gemv_row标量 stub第 1 节 PC 端对拍框架仓库https://gitee.com/pei-xiaoguang/kestrel-llm-mcu
返回列表