深入 Android 端侧 AI 推理的内存管理策略：从模型加载的内存峰值优化到 KV Cache 的动态回收机制
做端侧大模型部署，最先撞上的瓶颈通常不是推理速度，而是内存。一个 3B 参数的量化模型，加载时内存峰值能冲到 4-5GB，在 8GB RAM 的 Android 设备上，系统直接 OOM kill 是常态。这里记录我在端侧推理工程化过程中积累的内存管理策略——从模型加载、张量生命周期到 KV Cache 回收，形成一套可落地的低内存压力方案。
模型加载：mmap 为什么是端侧的标配
常规做法是用 malloc 分配一块堆内存，把模型权重整个读进来。模型文件 2GB，堆上就占 2GB，I/O 缓冲区还会额外消耗几百 MB，峰值内存是模型大小的 1.5 倍。
换成内存映射（mmap），事情就不一样了。让内核直接把文件映射到虚拟地址空间，物理页按需加载。
// 传统加载方式 — 峰值 = 文件大小 + I/O 缓冲区
char* buffer = new char[file_size];
fread(buffer, 1, file_size, fp); // 物理内存 +2GB
// mmap 方式 — 虚拟地址映射，物理页按需分配
int fd = open("model.gguf", O_RDONLY);
void* mapped = mmap(nullptr, file_size, PROT_READ,
MAP_SHARED, fd, 0);
// 物理内存仅占用实际访问过的页mmap 带来两个好处：一是只有实际读取的页面才从磁盘载入物理内存，未访问的权重不占 RAM；二是系统内存紧张时，干净的映射页可以直接回收，需要时再换入。
但 mmap 不是银弹。如果模型结构导致推理时需要随机访问全部权重，它就退化成了和全部加载一样的效果。实测中，MLP 层的稠密矩阵运算页面局部性差，mmap 的按需特性发挥有限；Embedding 层和 Attention 层只有当前序列触及的部分权重被加载，效果明显。
踩过的一个坑是文件分片与 mmap 的对齐问题。模型导出时如果没按页边界（4KB）对齐，mmap 映射后跨页访问的性能损耗可达 30%。建议在导出阶段对权重 tensor 做 page-aligned padding，配合 madvise(MADV_SEQUENTIAL) 给内核预取提示。
// 对映射区域设置预取策略
madvise(mapped, file_size, MADV_SEQUENTIAL); // 顺序访问，提前预取
// 推理完成后释放提示
madvise(mapped, file_size, MADV_DONTNEED); // 主动回收物理页张量生命周期：从”全部持有”到”用完即弃”
多数推理框架的执行流程大致是：读入输入 → 逐层计算 → 保留全部中间张量 → 输出结果。每层产生的张量全部堆在内存里，直到整次推理结束才释放。
实际并不需要全留着。Attention 层的 KV 张量要保留（涉及后续 token 的上下文计算），FFN 层的中间结果算完当前 token 就是废数据。
我的做法是注入 张量生命周期标记，让 allocator 在算子调度层面感知哪些张量可以提前释放：
struct TensorLifecycle {
int last_consumer_op; // 最后一个消费该张量的算子索引
bool is_persistent; // 是否跨前向推理保留（如 KV Cache）
};
// 在算子调度器中注入内存回收逻辑
void on_op_complete(int op_index, TensorPool& pool) {
for (auto& tensor : pool.tensors) {
if (tensor.last_consumer_op == op_index
&& !tensor.is_persistent) {
pool.free(tensor); // 立即回收，不等整帧结束
}
}
}优化后，Llama 架构 3B 模型的中间张量峰值占用从 1.2GB 降到 360MB。
实现层面还有两个需要处理的问题。
算子融合（Operator Fusion）：把 LayerNorm + MatMul 合并为一个 kernel，中间结果直接走寄存器不落内存，对 Attention 前后的归一化操作尤其有效。
内存池的碎片控制：频繁 alloc/free 会产生碎片，实际占用可能比理论值高 20-30%。我用了 BFC（Best-Fit with Coalescing）算法，相邻空闲块自动合并，单次推理的碎片率控制在 5% 以内。
KV Cache：滑动窗口是低内存设备的生存线
长文本场景下，KV Cache 才是真正的内存杀手。以 3B 模型、32 层、头维度 128 为例：输入 4096 token 时，KV Cache 约占据 512MB；扩展到 32K token，直奔 4GB。
静态裁剪的问题
限制序列长度是最直接的做法，比如硬截断到 2048 token。对问答和代码补全场景，上下文截断意味着模型丢失关键信息，回答质量断崖式下降。
滑动窗口回收
滑动窗口（Sliding Window）是更实用的方案——只保留最近 N 个 token 的 KV，超出部分直接回收：
void sliding_window_reclaim(KVCache& cache, int window_size) {
if (cache.seq_len > window_size) {
int evict_count = cache.seq_len - window_size;
// 将 [0, evict_count) 的 KV 槽位直接标记为可复用
// 实际实现中是环形缓冲区的指针移动，不触发内存拷贝
cache.head = (cache.head + evict_count) % cache.capacity;
cache.seq_len = window_size;
}
}实现关键在于环形缓冲区（Ring Buffer）。如果每次回收都做 memmove 把数据往前挪，CPU 开销完全不可接受。用 head 指针标记当前窗口起点，回收只是指针算术，O(1) 完成。
窗口大小需要分场景权衡：代码补全 512-1024 token 就够用（局部上下文），文档问答需要 2048-4096 token（跨段落依赖）。我一般采用动态窗口策略——默认 1024，检测到注意力权重中跨长距离依赖增强时，临时扩大到 2048。
按层衰减回收
滑动窗口的一刀切回收有个硬伤：浅层 Attention 更关注局部信息，深层 Attention 更依赖全局上下文。所有层用同一个窗口太浪费了。
按层衰减（Layer-wise Decay）更精细：浅层窗口更小，深层窗口更大，被驱逐的 KV 不做硬丢弃，而是加权衰减后合并到相邻槽位：
void layer_wise_evict(KVCache& cache, int layer_id, int max_tokens) {
int layer_window = max_tokens * (0.5 + 0.5 * layer_id / total_layers);
// 浅层保留少，深层保留多
if (cache.seq_len > layer_window) {
// 对超出部分做加权衰减而非直接丢弃
float decay = exp(-(cache.seq_len - layer_window) * 0.01);
cache.evict_with_decay(cache.seq_len - layer_window, decay);
}
}在 LongBench 基准上，相对等长滑动窗口，长文本问答的 ROUGE-L 提升约 4.2%，内存峰值降低 35%。
工程落地中的三个取舍
**mmap 加载 vs 预热策略。**mmap 首次访问时的缺页中断会拖慢首 token 延迟，冷启动场景尤其明显。折中方案：模型加载时对前 3 层的权重做 madvise(MADV_WILLNEED) 预热，其余层保持按需加载。首 token 延迟仅增加 100-200ms，内存峰值不受影响。
**KV Cache 的量化存储。**将 KV 从 FP16 量化为 INT8，内存直接减半，代价是精度损失。代码生成场景下 INT8 KV Cache 的 pass@1 下降不足 1%，可以接受；复杂推理任务退化到 3-5%。我按任务动态切换：代码和摘要用 INT8，翻译和推理保持 FP16。
**进程优先级与 LMK。**即便做了上述优化，极端场景下 Android 的 Low Memory Killer 仍可能杀掉推理进程。应对策略：申请 android:largeHeap="true"，将推理 service 提升到 IMPORTANCE_FOREGROUND，同时监听 onTrimMemory(TRIM_MEMORY_RUNNING_CRITICAL) 回调，触发 KV Cache 主动缩减到 256 token 的保底窗口。
这套组合拳打下来，3B 模型在 6GB RAM 设备上可以稳定跑 8192 token 的长文本推理，4GB RAM 设备上支持 4096 token 的对话场景。把端侧 AI 从”能跑”推到”能用”，内存管理是绕不开的硬仗。