大模型上小芯片:内存、算力与功耗的三重压缩工程实录

cover

一、256KB SRAM 上的推理困局:大模型与微控制器的物理鸿沟

把语言模型塞进 Cortex-M4,听起来像天方夜谭,但这其实是边缘 AI 的真实需求。工业场景中,设备端需要做关键词检测、异常文本分类、简单指令解析——这些任务用传统规则引擎很难覆盖,但完整的大模型又远超硬件承载能力。

以一块典型的 STM32F407 为例:256KB SRAM,1MB Flash,Cortex-M4F @ 168 MHz,没有 FPU 的向量运算。一个 5M 参数的 INT8 量化模型,权重存储需要 5MB——光 Flash 就放不下。这不是优化能解决的问题,而是需要从模型架构、推理框架、内存管理三个维度同时做极致压缩。本文记录将一个轻量文本分类模型部署到 STM32 上的完整工程路径,以及每个环节的取舍决策。

二、从模型到芯片:三层压缩的工程流水线

大模型部署到微控制器,需要经历三层压缩:模型压缩(结构裁剪与蒸馏)、数值压缩(量化与编码)、运行时压缩(内存复用与算子融合)。每一层都有精度损失,必须逐层验证。

flowchart TD
    A[原始模型 50M 参数] --> B[结构裁剪: 剪枝+蒸馏]
    B --> C[轻量模型 5M 参数]
    C --> D[INT8 量化: 权重+激活]
    D --> E[量化模型 5MB]
    E --> F[权重编码: 行程编码+码本]
    F --> G[压缩模型 1.2MB]
    G --> H[TFLite Micro 转换]
    H --> I[C 字节数组: 1.2MB Flash]
    I --> J[运行时: 张量复用+算子融合]
    J --> K[峰值 SRAM: 80KB]

    style A fill:#f99,stroke:#900
    style K fill:#9f9,stroke:#090

关键约束:STM32F407 的 Flash 访问速度约 25 MB/s(启用缓存后),SRAM 访问速度约 168 MHz 零等待。权重从 Flash 加载到 SRAM 的带宽是推理延迟的瓶颈。因此,运行时的核心优化不是减少计算量,而是减少 Flash 读取次数。

三、生产级微控制器推理流水线:TFLite Micro 的深度定制

标准 TFLite Micro 的内存分配器是线性分配,不支持张量复用。对于 SRAM 严重受限的微控制器,必须替换为自定义的环形缓冲区分配器。

#include <stdint.h>
#include <string.h>
#include <stdbool.h>

/* 环形张量缓冲区:在 SRAM 中预分配一块固定区域,
 * 所有中间张量共享此区域,按生命周期复用。
 * 为什么不用 malloc?嵌入式系统没有堆管理器,
 * 即使有,碎片化也会导致运行一段时间后分配失败。 */
#define TENSOR_ARENA_SIZE (80 * 1024)  /* 80KB,STM32F407 可承受 */

/* 静态分配张量竞技场:放在 SRAM 的特定段,
 * 避免 C 运行时堆管理器的干扰 */
__attribute__((section(".sram_tensor")))
static uint8_t tensor_arena[TENSOR_ARENA_SIZE];

/* 权重常量:放在 Flash 中,通过 DMA 加速读取
 * __attribute__((section(".flash_weights")))
 * 实际项目中由转换工具自动生成 */
extern const uint8_t g_model_data[];
extern const uint32_t g_model_data_len;

/* 简易张量生命周期管理器 */
typedef struct {
    uint8_t *base;
    uint32_t offset;
    uint32_t capacity;
} tensor_arena_t;

static tensor_arena_t arena = {
    .base     = tensor_arena,
    .offset   = 0,
    .capacity = TENSOR_ARENA_SIZE,
};

/**
 * @brief 从竞技场分配张量空间
 *
 * 为什么用线性分配而非空闲链表?
 * 推理过程中张量的分配/释放顺序是确定的,
 * 不存在随机分配模式。线性分配 O(1) 时间复杂度,
 * 且零碎片。每次推理前重置 offset 即可复用整个区域。
 */
static void *arena_alloc(tensor_arena_t *a, uint32_t size)
{
    /* 4 字节对齐:Cortex-M4 的 LDRD/STRD 要求对齐访问,
     * 未对齐访问会触发 UsageFault 或性能严重下降 */
    size = (size + 3) & ~3u;

    if (a->offset + size > a->capacity) {
        /* 竞技场溢出:这是致命错误,不能静默失败 */
        return NULL;
    }

    void *ptr = a->base + a->offset;
    a->offset += size;
    return ptr;
}

/**
 * @brief 重置竞技场:推理结束后调用
 */
static void arena_reset(tensor_arena_t *a)
{
    a->offset = 0;
}

/* 推理结果结构 */
typedef struct {
    int8_t   top_class;       /* 最高概率类别 */
    uint8_t  confidence_pct;  /* 置信度百分比 */
    int32_t  inference_ms;    /* 推理耗时 */
} inference_result_t;

/**
 * @brief 单次推理:输入文本特征,输出分类结果
 *
 * 完整流程:特征提取 -> 模型前向传播 -> 后处理
 * 特征提取在 MCU 上用查表法实现,不依赖浮点运算
 */
static inference_result_t run_inference(const int8_t *input_tokens,
                                         uint16_t seq_len)
{
    inference_result_t result = {0};
    uint32_t start_tick, end_tick;

    /* 重置张量竞技场:上一轮推理的中间张量全部失效 */
    arena_reset(&arena);

    start_tick = HAL_GetTick();

    /* === 输入张量分配 === */
    int8_t *input_tensor = (int8_t *)arena_alloc(
        &arena, seq_len * sizeof(int8_t));
    if (!input_tensor) {
        /* 竞技场溢出:序列长度超出设计上限 */
        result.top_class = -1;
        return result;
    }

    /* 拷贝输入数据到张量竞技场
     * 为什么不直接用 input_tokens 指针?
     * 因为 TFLite Micro 的算子要求输入张量
     * 位于竞技场连续区域内,否则 DMA 传输会越界 */
    memcpy(input_tensor, input_tokens, seq_len);

    /* === 中间张量分配 ===
     * 以下为简化示例,实际项目中由 TFLite Micro 的
     * 内存规划器自动计算各层张量大小和生命周期 */

    /* Embedding 层输出:seq_len * embed_dim */
    int8_t *embed_out = (int8_t *)arena_alloc(
        &arena, seq_len * 32);  /* embed_dim=32,压缩后 */

    /* 全局平均池化输出:embed_dim */
    int8_t *pool_out = (int8_t *)arena_alloc(
        &arena, 32);

    /* 分类头输出:num_classes */
    int8_t *logits = (int8_t *)arena_alloc(
        &arena, 8);  /* 8 个类别 */

    if (!embed_out || !pool_out || !logits) {
        result.top_class = -2;  /* 内存不足 */
        return result;
    }

    /* === 前向传播(伪代码,实际由 TFLite Micro 算子执行) ===
     * 1. Embedding 查表:input_tokens -> embed_out
     * 2. 全局平均池化:embed_out -> pool_out
     * 3. 全连接层:pool_out -> logits
     * 4. Softmax/ArgMax:logits -> top_class
     */

    /* === 后处理:找到最大 logit === */
    int8_t max_val = -128;
    int8_t max_idx = 0;
    for (int i = 0; i < 8; i++) {
        if (logits[i] > max_val) {
            max_val = logits[i];
            max_idx = i;
        }
    }

    /* 将 INT8 logit 转换为百分比置信度
     * 为什么用线性映射而非 softmax?
     * softmax 需要浮点 exp 运算,Cortex-M4F 的 FPU
     * 不支持指数函数,软件实现耗时约 2us/次。
     * 线性映射精度够用,延迟降低 90% */
    result.top_class      = max_idx;
    result.confidence_pct = (uint8_t)((max_val + 128) * 100 / 255);
    result.inference_ms   = HAL_GetTick() - start_tick;

    return result;
}

/**
 * @brief 功耗感知推理调度器
 *
 * 为什么需要调度器?电池供电设备不能持续运行推理,
 * 必须根据电池电量和业务优先级决定是否执行推理。
 * 低电量时降低推理频率,高优先级事件立即触发。
 */
typedef enum {
    PWR_MODE_FULL,    /* 全速推理:100ms 间隔 */
    PWR_MODE_NORMAL,  /* 正常推理:500ms 间隔 */
    PWR_MODE_LOW,     /* 低功耗推理:2000ms 间隔 */
    PWR_MODE_SLEEP,   /* 休眠:不推理,仅中断唤醒 */
} power_mode_t;

typedef struct {
    power_mode_t mode;
    uint32_t     interval_ms;
    uint32_t     battery_mv;    /* 电池电压,mV */
    bool         urgent_event;  /* 紧急事件标志 */
} inference_scheduler_t;

static uint32_t get_inference_interval(inference_scheduler_t *sched)
{
    /* 紧急事件立即触发,不受模式限制 */
    if (sched->urgent_event) {
        sched->urgent_event = false;
        return 0;
    }

    /* 根据电池电压自动调整模式
     * 为什么用电压而非百分比?锂电池放电曲线在
     * 3.0V-3.6V 平台期后急速下降,百分比无法
     * 准确反映剩余可用能量,电压更直接 */
    if (sched->battery_mv > 3600) {
        sched->mode = PWR_MODE_FULL;
    } else if (sched->battery_mv > 3400) {
        sched->mode = PWR_MODE_NORMAL;
    } else if (sched->battery_mv > 3200) {
        sched->mode = PWR_MODE_LOW;
    } else {
        sched->mode = PWR_MODE_SLEEP;
    }

    switch (sched->mode) {
        case PWR_MODE_FULL:   sched->interval_ms = 100;  break;
        case PWR_MODE_NORMAL: sched->interval_ms = 500;  break;
        case PWR_MODE_LOW:    sched->interval_ms = 2000; break;
        case PWR_MODE_SLEEP:  sched->interval_ms = 0;    break;
    }

    return sched->interval_ms;
}

四、微控制器推理的物理极限:Flash 带宽、SRAM 碎片与热失控

在微控制器上跑推理,有三个物理极限无法通过软件优化绕过。

Flash 读取带宽。 STM32F407 的 Flash 在 168 MHz 主频下需要 5 个等待周期,有效读取带宽约 33 MB/s。一个 1.2MB 的模型,仅权重加载就需要 36 ms——这还没算计算时间。优化手段:启用 ART 加速器(Flash 预取指缓存),将热点权重复制到 CCM(Core Coupled Memory,零等待 RAM),但 CCM 只有 64KB,且 DMA 无法访问。

SRAM 碎片化。 推理过程中,不同层的张量大小差异巨大。Embedding 层输出可能是 32KB,而池化层输出只有 32B。如果用通用堆分配器,频繁分配释放不同大小的张量会导致严重碎片。解决方案是前文展示的竞技场分配器——线性分配、整体重置、零碎片。代价是所有张量必须共享同一块连续内存,不能跨推理调用保持中间状态。

热稳定性。 Cortex-M4F 在 168 MHz 满负荷运行时,芯片结温可能上升 20-30 度。在工业环境(环境温度 85 度)中,结温可能超过 105 度的规格上限。推理是计算密集型操作,持续推理会导致温度累积。解决方案:推理后插入空闲周期让芯片降温,或降低主频运行推理。前者增加延迟,后者降低吞吐——这是功耗与性能的经典权衡。

flowchart LR
    A[模型部署需求] --> B{Flash 容量}
    B -->|模型 > Flash| C[权重压缩: 码本量化]
    B -->|模型 < Flash| D{SRAM 容量}
    C --> D
    D -->|峰值 > SRAM| E[张量复用: 竞技场分配]
    D -->|峰值 < SRAM| F{推理延迟}
    E --> F
    F -->|延迟 > 要求| G[算子融合 + CCM 加速]
    F -->|延迟 < 要求| H{功耗预算}
    G --> H
    H -->|功耗超标| I[降频推理 + 间歇调度]
    H -->|功耗达标| J[部署完成]
    I --> J

还有一个容易被忽视的问题:INT8 量化的累加溢出。Cortex-M4 的 32 位 MAC(乘累加)单元在做 INT8 卷积时,如果卷积核较大(如 7x7),累加值可能超出 int32 范围。TFLite Micro 默认不做溢出检查,一旦溢出,推理结果完全错误且无任何报错。防护措施:在模型转换时指定 per_channel_quantization,缩小每通道的 scale 值,降低累加溢出风险。

五、总结

把大模型塞进微控制器,不是"能跑就行"的 Demo 工程,而是需要在模型架构、数值精度、内存管理、功耗调度四个维度同时做极致优化的系统工程。落地的关键步骤:

  1. 模型先裁剪再量化:先通过蒸馏将模型压缩到 Flash 可容纳的尺寸,再量化到 INT8。不要试图量化一个放不下的模型。
  2. 竞技场分配器替代堆:所有推理张量使用线性分配+整体重置,消除碎片化风险。
  3. Flash 带宽是瓶颈:权重读取延迟远大于计算延迟,优先优化权重编码和缓存策略。
  4. 功耗调度必须做:电池供电设备必须根据电量动态调整推理频率,低电量时降频或休眠。
  5. 累加溢出必须防:INT8 卷积的累加值可能溢出 int32,转换时用逐通道量化缩小 scale。

微控制器上的推理,每一 KB 内存、每一 mW 功耗、每一 ms 延迟都是硬约束。没有银弹,只有逐项优化。

更多推荐