MNN llm_demo深度解析：Omni多模态模型端侧推理全链路

MNN 的 llm_demo 里跑通一个 Omni 多模态模型比想象中要绕得多。我一开始以为只是把视觉编码器接到语言模型前面输入从文本换成图像加文本就完事了真正打开源码之后才发现从 token 层面的拼接规则、注意力掩码的生成到 KV Cache 的预分配策略再到算子层的融合方式每一个环节都跟纯文本 LLM 推理有本质区别。这篇就围绕 llm_demo 的源码结构把 Omni 模型推理的整条链路拆开来讲重点是输入怎么处理、三路模态怎么汇合、采样和 KV Cache 是怎么设计的、算子侧做了哪些取舍最后会带一点部署实测的经验和坑。适合正在做端侧多模态推理、或者想通过 llm_demo 快速理解 MNN 推理范式的开发者。1. llm_demo到底在解决什么问题纯文本LLM之外的第二个战场1.1 为什么MNN要单独维护一个llm_demo传统 CNN 模型的推理在 MNN 里很成熟模型是静态计算图输入尺寸基本固定一次推理前向就能出结果。LLM 完全不同它是自回归的一次前向只产生一个 token然后把这个 token 重新拼进输入序列继续下一次前向。这决定了它不能像 CNN 那样“一张图跑一次就完”而是需要一个专门的外层循环来管理 prompt 的初始推理、后续 decode、采样、停止条件、上下文窗口等一系列状态。llm_demo 就是这一层外层逻辑的载体。它不负责底层算子优化也不参与模型训练它做的事是把用户在终端里输入的话转换成模型需要的 token 序列交给 MNN 的 runtime 执行一次次前向再把每一轮吐出来的 logits 采样成具体的 token最后解码成文字打印出来。整个流程在代码里是可控、可读的这就给了做端侧部署的人一个很好的参照起点。1.2 Omni模型给推理框架带来的新变化如果你只跑过 7B 级别的纯文本模型可能会觉得 llm_demo 已经够用了一个 tokenizer、一个 LlmSession、一个采样器完事。但 Omni 模型不是这样它是多模态输入、多模态输出的通常包含视觉编码器、音频编码器和语言模型主干三块。以常见的 Omni 架构为例视觉部分是 ViT 加一个投影层音频部分可能是单独的编码器加可学习的降采样模块文本部分走 embedding。这三个分支的输出长度、特征维度、语义空间都不一样必须在送入语言模型之前完成对齐。这就意味着 llm_demo 的代码里不只是“文本 token ids”而需要引入一种更通用的“输入序列”概念一段序列里既有文本 token又有图像特征片段还可能插入音频特征片段。语言模型的注意力计算必须知道哪些位置是文本、哪些位置是图像、哪些位置是音频并且在生成掩码的时候对这些位置做正确的可见性处理。1.3 源码目录的核心文件职责我梳理代码的时候习惯先把目录看一遍搞清楚每个文件大概管什么再顺着调用链去读。llm_demo 的目录通常包含下面这些模块文件/目录主要职责llm.cpp主入口负责命令行参数解析、对话循环、调用 pipelinellm_session会话管理加载模型、初始化线程池、保存对话历史tokenizer把文本字符串转成 token ids支持特殊占位符prompt_processor处理多模态输入拼接图像/音频特征与文本sampler根据 logits 决定下一个 token支持贪心/top_k/top_p/temperaturekv_cacheKV Cache 的分配、读写、长度管理llm_export导出脚本把训练好的权重转换成 MNN 可加载的格式config读取模型的 json 配置比如 hidden_size、layer 数、量化位宽光看文件列表就看得出来llm_demo 已经不只是一个“demo”它本质上是一个可扩展的 LLM 推理框架只是代码量控制得比较精简适合学习和二次开发。2. 从一行输入到首tokenOmni推理主链路的源码追踪2.1 输入侧的处理文本、图像、音频如何变成同一批特征我在看代码时最先找的是入口函数看它接收了什么参数。llm_demo 在交互模式下用户输入通常是一个字符串加可选的文件路径比如描述一下这张图片图片路径 /tmp/a.jpg。prompt_processor 做的事就是把这个混合输入解析成标准的模型输入。文本直接走 tokenizer会在句子开头加上 chat 模板里的系统提示和用户标记。图像路径被摘出来之后先用 MNN 的图像解码接口读成 tensor再做 resize、归一化。这一步有个很关键的点图像输入大小不是随意定的要跟视觉编码器预训练时使用的分辨率保持一致通常是 336×336 或 448×448 这样的固定尺寸如果输入是其他尺寸要在预处理里 resize。音频分支相对少见但 Omni 模型里通常会把音频切成固定长度的帧经过编码器之后产生一系列帧级特征。在 llm_demo 里音频通常被当作和图像并列的特征段以特殊 token 开头后面跟随编码后的特征向量序列。最终输入被构造成一个统一的结构体里面至少包含两部分token ids 序列和 aligned features 序列。两者的组合顺序由模型的 chat template 决定典型顺序是系统 prompt 文本 token、用户文本 token、图像特征起始符、图像特征、换行符、用户问题文本 token。// 简化后的prompt_processor核心流程 std::vectorint tokenIds; std::vectorstd::vectorfloat visualFeatures; // 1. 文本部分常规编码 auto textTokens tokenizer-encode(userText, /*addSpeical*/true); tokenIds.insert(tokenIds.end(), textTokens.begin(), textTokens.end()); // 2. 如果存在图像输入在文本特定位置插入视觉特征 if (imagePath ! ) { auto imageTensor imageProcessor-loadAndPreprocess(imagePath); auto visionEmbeds visionEncoder-encode(imageTensor); tokenIds.push_back(tokenizer-imageStartId()); visualFeatures.push_back(visionEmbeds); tokenIds.push_back(tokenizer-imageEndId()); }这段代码虽然做了一定简化但体现了 llm_demo 里最核心的设计语言模型的输入序列永远是整数 token ids视觉和音频特征通过在对应位置插入特殊 token 来“占位”真正的浮点特征向量单独保存在 embedding 阶段拼进去。2.2 模型加载与计算图初始化模型文件不是直接加载一个 MNN 模型就完事llm_demo 走的是“权重分发 按需构图”的路线。在 llm_export 阶段原始权重会被转换成分片格式每个文件对应模型中若干层的参数配置文件里记录了所有分片的位置和数据类型。代码里通常会有类似load_model_config()的函数读取 JSON 配置拿到 hidden_size、num_attention_heads、num_hidden_layers、vocab_size 这些关键参数然后根据这些参数预先分配所有 tensor 的 buffer。这个设计在端侧特别重要因为手机内存是受限的TensorFlow Lite 和 ONNX Runtime 那种“用多少分配多少”的做法在这里会造成频繁的内存碎片和额外开销。session 初始化完成之后MNN 会构建一个针对当前输入动态 shape 的计算图。LLM 的输入长度每一轮都在变化所以图的 shape 是动态的但在 backend 里 MNN 会尽量复用已分配的内存。2.3 prefill与decode的两条路径一次完整的生成过程在 llm_demo 里被拆成两个阶段。prefill 阶段处理整段 prompt输入比较长需要并行计算所有位置的注意力。decode 阶段每次只输入一个新 token但 KV Cache 已经包含了之前所有的中间状态只需要计算最后一个位置。源码里我看到的典型结构是这样的LlmSession::run()内部先判断当前是否是第一轮如果是就执行一次 prefill把整个输入序列跑完得到第一个 logits之后每一轮都走 decode 路径。两个路径共用同一个计算图区别在于提供给图的输入 shape 和 mask 不同。prefill 的耗时一般远大于 decode因为它要处理整个 prompt 的注意力矩阵而 decode 每一步只有一行 query 要和全部 KV 做 attention。llm_demo 里会把 prefill 和 decode 的算子安排到不同的执行策略上prefill 更看重并行度decode 更看重单步延迟。3. 视觉、音频、文本三路特征在源码里是怎么汇合的3.1 特殊token与tokenizer层面的拼接规则Omni 模型在 tokenizer 层面就做了不少文章。视觉特征不是一个单独的 embedding 表而是通过几个“占位 token”告诉模型接下来的若干向量来自图像编码器不是来自词表。我在源码里经常看到的几个特殊 token id 是 image_start、image_end、audio_start、audio_end。它们的 id 是在导出模型时固定下来的llm_demo 的 tokenizer 模块会维护这些 id并在编码过程中自动插入或跳过。如果用户输入是“图片里有什么文字”而 tokenizer 配置里标注了图片文件路径那么编码结果可能是[CLS] 图片里有什么[/INST]其中image是一个单独的 token id后面跟着视觉特征。这看起来像模板字符串替换但在源码里它是一套严格的三段式处理先切文本、再放占位符、最后在 embedding 阶段覆盖。3.2 视觉分支的编码与特征映射视觉分支做的事情比我想象的多。它不是一个简单的 ViT forward通常包含一个图像编码器和一个 resampler/投影层。图像先被切成 patches经过 ViT 的若干 transformer 层输出 patch 级别的特征序列。这个序列往往很长比如 448×448 图像切成 16×16 patch得到 784 个 patch再乘以通道数算力开销不小。接着进入投影层把 ViT 的输出维度映射到语言模型的 hidden_size。llm_demo 源码里通常会封装一个VisionEncoder类内部持有两个 MNN 子模型一个代表 ViT一个代表投影层。调用encode()时先跑 ViT再跑投影层输出可以直接拼进语言模型 embedding 的二维矩阵。这里有个性能关键点视觉特征可能与文本特征拼接时长度非常长如果每次 decode 都重新编码图像效率会非常低。llm_demo 的做法是只在 prefill 阶段编码一次把得到的视觉特征缓存起来后续 decode 直接复用。3.3 注意力掩码的构造图像特征对整个序列可见纯文本 LLM 的注意力掩码是一个左下三角矩阵当前位置只能看到它前面的位置。但在多模态模型里图像特征段通常是整段可见的因为一张图的信息是全局的不存在“只看图的前半部分”这种说法。图像起始 token 和图像结束 token 之间所有的特征位置都会与整段文本建立注意力连接。源码里这个掩码的构造逻辑通常放在attention_mask的生成函数中。它对所有属于视觉特征的 token 位置赋值为 0表示不遮挡对文本部分保持因果掩码对当前正在生成的位置只允许看到已经生成的上下文。如果音频流是流式的掩码还会更复杂因为可能当前音频帧已经在播放后面的音频还没有到达。3.4 音频流与多模态融合的工程处理支持音频输入的 Omni 模型在 llm_demo 里通常是走离线帧编码的路线。音频文件先按固定窗口切帧每帧经过音频编码器得到若干特征向量这些特征向量再与文本特征拼接。因为音频特征往往包含时间先后关系它的位置编码不能乱插需要保证送入 attention 时顺序与时间轴一致。在代码层面我见过一种很实际的处理方式音频特征被当作图像特征的一种“变体”统一放到 visual_features 结构里只通过特殊的起始 token 区分。这样做的好处是语言模型主干完全不用区分输入来自哪个模态只要 token 位置和特征位置对应上attention 机制天然就能处理。4. decode阶段的采样与KV Cache源码里的工程取舍4.1 logits后处理与采样参数每轮前向结束后模型输出一个 shape 为[1, 1, vocab_size]的 logits。llm_demo 会把 logits 取出来经过温度缩放、top_k 截断、top_p 筛选最后采样得到下一个 token。源码里常见的函数签名看起来类似这样int Sampler::sample(const float* logits, int vocabSize, const GenConfig cfg) { // 温度缩放 if (cfg.temperature 0) { for (int i 0; i vocabSize; i) { logits[i] / cfg.temperature; } } // top_k if (cfg.topK 0 cfg.topK vocabSize) { // 找出第k大的值小于该值的logits置为-inf } // top_p if (cfg.topP 0 cfg.topP 1.0) { // 按概率从大到小累加累计超过阈值的后面全部置为-inf } // softmax // 随机采样或取最大 }这个模块看起来不复杂但实际调参的时候影响很大。我在实际跑模型时发现多模态模型对 temperature 的敏感度比纯文本模型更高图像输入之后模型很容易陷入重复输出“好的”“然后”这类词的循环。llm_demo 里实现了一个 repeat_penalty 机制会在 logits 后处理阶段对已经在上下文里出现过的 token 做惩罚这个对多模态场景很关键。4.2 KV Cache的组织方式与复用策略KV Cache 是 LLM 推理里最占内存的部分。一个 7B 模型层数 32hidden_size 4096如果上下文长度是 2048KV Cache 的总量大约是 32 × 2 × 4096 × 2048 × 2 字节如果按 float16 算大概 512 MB 级别。这还没有算上视觉特征占用的部分。llm_demo 在实现 KV Cache 时有一个值得注意的设计它把 KV 缓冲区按最大上下文长度一次性预分配而不是每次增长。这样做有两个好处一个是避免多次 malloc 造成的内存碎片另一个是 decode 阶段不用频繁做拼接操作。因为注意力计算要访问全部的 KV如果 KV 分散在多段内存里访存效率会下降。但代价也很明确预分配大小必须设置得足够大否则长对话会越界如果设置得过大内存会浪费。所以 llm_demo 在 config 里通常有一个max_context_length参数这个值的设定需要结合模型能力和设备内存做权衡。4.3 视觉特征在KV Cache里的特殊处理多模态模型引入了一个新问题视觉特征经过 prefill 之后它们的 K、V 也要进入缓存。这部分缓存的长度是固定的跟图像 patch 数量有关不会随着 decode 过程增长。llm_demo 在初始化 KV Cache 时会为视觉特征预留单独的静态区域避免它们和动态增长的文本 KV 混在一起造成管理混乱。从代码结构上看KV Cache 内部通常维护一个past_token_num计数器它表示已经缓存了多少个位置。prefill 之后这个计数器的值会等于 prompt token 数加上视觉特征数。之后每 decode 一步计数器加 1。如果输入的图像尺寸不同视觉特征数也会不同这对 KV Cache 的预分配大小计算是一种考验源码里一般按最大可能图像尺寸预留。实际操作中如果模型允许动态图像分辨率编码后的视觉 token 数变化很大。我在理解长文图像比如整页PDF截图时遇到过内存不够的情况就是因为视觉 token 数量远超设计预期。llm_demo 通常不支持这种动态扩展所以选模型导出尺寸时要先想清楚。5. 算子性能和量化是Omni落地的真正瓶颈5.1 注意力算子的融合别让访存打败算力MNN 在端侧做 LLM 推理时最大的瓶颈并不是计算量而是访存带宽和算子启动开销。纯文本模型已经存在这个问题多模态模型把问题放大了视觉编码器是典型的 CNN 加 Transformer 混合结构文本语言模型是标准的 decoder-only Transformer两者对算子融合的要求不一样。源码里能看到 MNN 的 backend 对 attention 做了专门优化。普通的 attention 实现需要三次矩阵乘加一次 Softmax中间结果要多次读写全局内存导致带宽占用极大。llm_demo 中如果检测到设备支持特定硬件指令集会走 Flash Attention 风格的融合实现把 QK^T、mask、Softmax、V 的乘加合在一个 kernel 里完成中间结果保持在寄存器或共享内存中。这个融合模式对多模态模型尤其重要因为视觉部分 attention 的序列长度可能到一两千如果不做融合光 QK^T 的中间矩阵就要占用很大的临时内存还会拖慢整体推理速度。5.2 低比特量化对多模态各层的影响量化方案是 llm_demo 里少不了的环节。常见的设置是语言模型主干用 4bit 量化视觉编码器用 8bit 或者保持 float16。为什么会有这种差异我自己的体会是视觉编码器对量化误差更敏感图像特征一旦被过量压缩位置、颜色、文字细节会大量丢失导致模型“看得见但看不懂”。而语言模型主干因为训练时有很强的语言先验4bit 量化后质量损失相对可控。llm_demo 的量化实现更像是“分模块量化”。代码里会为每个子模型提供一个单独的量化配置语言模型部分走llm_export里的 weight-only 量化视觉和音频编码器则保持原始精度加载。这样做也方便在调试时对比量化前后的效果差异。我在实测一个 7B Omni 模型时发现如果视觉编码器用 4bit图像描述任务的准确率会明显下降尤其对 OCR 相关的任务影响最严重文本识别错误率比 8bit 高了一倍多。所以如果你做的是图像问答、文档理解类场景视觉部分建议保住 8bit 甚至 float16。5.3 实测不同平台的耗时分布参考不同芯片的推理耗时分布差异很大。以我实测过的某端侧 SoC 上跑一个 7B Omni 模型为例prefill 阶段大约是 1.5 到 3 秒其中视觉编码占了 40% 左右文本 prefill 占 50%剩下的开销在 mask 构造和算子启动。decode 阶段每 token 大约 40 到 80 毫秒在这个阶段视觉编码已经不参与了主要时间都花在语言模型主干的矩阵乘上。阶段主要耗时来源占比参考图像预处理解码、resize、归一化5%~10%视觉编码ViT 前向30%~40%文本 prefill首轮注意力计算20%~30%decode单 token 自回归30%~40%值得说的是很多人只盯着 decode 的延迟忽视了 prefill 和视觉编码在多模态场景里的占比。如果用户每次提问都附带一张新图片那么图像预处理加视觉编码的耗时是躲不掉的这直接决定了“从按下回车到看到第一个字”的等待时间。产品侧如果要求低首字延迟得从视觉编码器的轻量化入手。6. 一套可以抄的端侧部署方案与踩坑记录6.1 针对Omni模型的参数配置建议部署 Omni 模型时我认为有五个参数最值得关心。线程数不是越大越好。MNN 在端侧一般推荐设置成 4 到 6 个线程具体看芯片的核心数和频率。我试过把线程数拉到 8结果因为内存带宽争抢decode 速度反而下降了 10% 以上。内存上限要卡死。通过setMemoryLimit之类的接口限制内存使用防止模型加载时把系统内存打满导致进程被杀。对于 7B 级模型保留 1GB 给运行时和图像处理是必要的。图像输入分辨率要按模型导出时的配置走不要随意放大。我发现很多人喜欢把图像输入设成 1024×1024觉得分辨率高识别准确率就高但视觉编码器如果训练时只见过 448×448强行上高分辨率反而会造成特征分布偏移。上下文长度设置要和量化位宽一起考虑。4bit 模型能承担更长的上下文因为每层的 KV Cache 占用小省下来的内存可以放更多历史 token。采样参数中的 repeat 惩罚建议开启多模态模型在视觉背景下特别容易钻入同一个语义区域反复输出。6.2 我在移植过程中踩过的几个坑第一个坑是动态 shape 引发的模型加载失败。直接用 llm_export 导出的模型如果输入维度被固定死换一个不同尺寸的图像就会报 shape 不匹配。解决方式是在导出时把视觉分支的输入维度设成动态或者干脆把图像垫到某个固定尺寸再统一处理。第二个坑是特殊 token id 对不上。我在拿到一个第三方转换模型后发现生成结果始终包含大量unk排查了半天问题出在 tokenizer 的 vocab 文件里缺少 image_start 和 image_end 的 id导致 prompt 拼接时插入的不是预期的特殊令牌。llm_demo 在这类问题上没有做太强的校验转换模型之后一定要手动检查这几个 id。第三个坑是 prefill 阶段的内存峰值。有些底端设备跑纯文本模型没问题一上 Omni 模型就 OOM原因往往不是模型本身而是视觉特征在 prefill 时产生的临时注意力矩阵。视觉 token 数如果达到上千prefill 的中间计算结果会瞬间吃掉大量内存解决方式是手动把视觉特征分段送入语言模型而不是一次性全量拼接。6.3 建议的源码分析方法如果你想自己把 llm_demo 彻底啃下来我的方法可以供参考不要从 llm.cpp 一路顺着读那样容易陷进参数解析和日志打印里。先跑一次完整的对话记录日志输出找到prefill和decode的边界然后在 prefill 函数入口打断点观察输入序列的 token ids 和视觉特征的 shape再从上下行代码里找attention_mask的生成逻辑确认图像位置是不是真的没有遮挡最后拿一张固定图片反复测试 KV Cache 的past_token_num变化把这个数字和理论计算对一遍。源码分析最忌讳的是只看结构不跑数据。llm_demo 的好处是它本身就带交互模式你随时可以把中间的隐藏状态 dump 出来和 PyTorch 的实现做比对。我就是靠这种对比在视觉位置编码上找到了一个拼接顺序错误差了几行代码图像描述效果却天差地别。工具链上llm_export 脚本里的参数大多可以直接复用但要对模型卡片里的“原始预训练上下文长度”保持清醒。端侧部署通常跑不到那个长度原因是内存不允许强迫自己理解每个参数的实际含义比盲目加大上下文长度更有价值。如果后面有时间我还想再写一篇关于视觉编码器单独前向的性能剖析把 ViT 和投影层各自的耗时数据补上方便大家在做裁剪时找准方向。