十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

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

LiteRT-LM MTP 投机解码源码解析:草稿-验证两阶段机制完整指南

LiteRT-LM MTP 投机解码源码解析:草稿-验证两阶段机制完整指南 LiteRT-LM MTP 投机解码源码解析草稿-验证两阶段机制完整指南【免费下载链接】LiteRT-LMLiteRT-LM is Googles production-ready, high-performance, open-source inference framework for deploying Large Language Models on edge devices.项目地址: https://gitcode.com/GitHub_Trending/li/LiteRT-LMLiteRT-LM 是 Google 开源的端侧大语言模型LLM高性能推理框架其内置的MTPMulti-Token Prediction多 Token 预测投机解码机制通过草稿-验证两阶段流水线显著提升了 Decode 阶段的出字速度。本文面向新手带你快速读懂 LlmLiteRtMtpDrafter 的核心源码轻量草稿模型如何起草 Token、大模型如何并行验证、以及奖励 Tokenbonus token是如何被接受的。为什么需要投机解码端侧推理的速度瓶颈 大模型生成文本时Decode 阶段每个 Token 都要完整跑一遍主模型。Token 数量少时计算量小但内存带宽开销大——每次生成 1 个 Token 的耗时接近生成整句的耗时。MTP 投机解码的思路是先用一个超轻量的小模型Drafter快速猜出接下来 G 个 Token再把这 G 个 Token 连同当前 Token 一起一次性送进主模型并行验证 G1 个位置主模型输出的前 G 个 Token 与草稿逐个比对连续正确的部分全部保留第一个不一致的位置由主模型的正确答案顶替即奖励 Token。这样一次 Decode 调用最多可产出G1个 Token而大模型只运行了一次速度收益十分可观。整体架构草稿模型、验证模型与双采样器核心类 LlmLiteRtMtpDrafter 同时持有两个角色组件说明关键代码mtp_drafter_model_轻量 MTP 草稿模型输出logits与projected_activationsllm_litert_mtp_drafter.ccbase_model_verify_signature_主模型专门使用名为verify的签名批量验证llm_litert_mtp_drafter.ccdrafter_sampler_草稿采样器贪心Top-K1batch 大小为 1CreateGreedySamplerverifier_sampler_验证采样器贪心batch 大小为 G1llm_litert_mtp_drafter.cc两个采样器都配置为Top-K1、p0、temperature1.0的纯贪心采样器——因为投机解码要求草稿与验证结果可确定性比对随机采样会导致比对失效。草稿步数 G 不是写死的而是从主模型verify签名的input_pos输入维度自动推导num_draft_steps input_pos_dims[0] - 1见 llm_litert_mtp_drafter.cc。第一阶段草稿——用轻量模型快速写 G 个 Token草稿入口是 RunDraftingLoop它是一个循环 G 次的流水线拼接入场激活值把上一步的 Embedding 向量与上一步输出的projected_activations拼接得到[1, 1, hidden_size * 2]的输入例如 1536×23072 维见 ConcatenateEmbeddingsAndActivations。首次草稿由外部传入激活值之后的步骤直接复用草稿模型自己的输出实现自回归接力。异步运行草稿模型mtp_drafter_model_.RunAsync(...)以异步方式执行充分利用 GPU 流水线。采样出 1 个 Token对输出 logits 应用约束掩码如有后贪心采样得到sampled_draft_id追加进drafted_tokens。记录约束状态若启用了约束解码同步推进Constraint::State为验证阶段逐位置施加掩码做准备。整个过程见 llm_litert_mtp_drafter.cc 第 590-644 行。第二阶段验证——主模型一次性并行核对 G1 个 Token草稿完成后PrepareVerifierInputBuffers 负责组装验证输入input_pos填入连续的position .. position G将当前 Token G 个草稿 Token查表得到 G1 个 Embedding填充全局/滑动窗口注意力掩码验证阶段使用RingBufferAttentionMaskMode::kVerify模式。随后 RunVerification 调用主模型的verify签名一次执行返回G1 个位置的采样结果。若存在约束还会对每个位置批量应用对应的LogitMaskApplyMasksToLogitsSequence保证验证结果同样满足格式约束。接受逻辑连续比对 奖励 Token ✨整个机制最巧妙的部分在 Draft 的收尾逐个比对从左到右比较主模型验证输出与草稿 Token统计连续相等的num_correct_tokens第一个不等处立即停止奖励 Token取主模型在第一个不一致位置或末尾输出的 Token 作为bonus_token它一定是主模型说了算的正确结果最终输出前num_correct_tokens个草稿 Token bonus_token一次 Decode 就产出了 G1 个 Token 中的合格部分状态接续记录last_verified_token_id_idx_下一轮草稿直接从最后一个被接受的 Token 处接力保证上下文无缝衔接。无论草稿命中多少每轮至少产出 1 个 Token奖励 Token且质量与纯主模型解码完全一致——这是投机解码无损加速的精髓。执行器集成如何开启 MTP 投机解码MTP 与主执行器的对接位于 llm_litert_compiled_model_executor.cc配置项enable_speculative_decoding定义在 llm_executor_settings.h默认关闭开启后采用懒加载首次 Decode 时才调用EnsureMtpDrafterLoaded()编译草稿模型不额外消耗启动时间草稿模型在包内以kTfLiteMtpDrafter枚举值 13类型存放见 model_resources.hPrefill 之后的第一步 Decode会额外做一次常规解码把激活值喂给 Drafter 完成点火之后进入纯草稿-验证循环。构建目标runtime_executor_llm_litert_mtp_drafter定义在 CMakeLists.txt单元测试见 llm_litert_mtp_drafter_test.cc。与约束解码的联动Tool Call 场景下的 MTPLiteRT-LM 的投机解码并非孤立存在——草稿与验证两个阶段都会感知Constraint约束解码器底层基于 llguidance 语法引擎。草稿阶段逐步推进语法状态验证阶段则为 G1 个位置批量计算并应用掩码确保即使输出 JSON 工具调用这类强格式内容加速也不破坏约束。这一联动让 MTP 与工具调用流程协同工作可观测性一眼看懂草稿命中率 LlmLiteRtMtpDrafter 内置了两个计数器num_drafted_tokens_累计起草数与num_verified_tokens_累计被接受数对象析构时打印成功率见 llm_litert_mtp_drafter.cc。这个Success rate正是评估你的模型/场景是否适合启用投机解码的关键指标命中率越高端到端加速越明显。关键文件速查表文件作用llm_litert_mtp_drafter.hDrafter 类声明、缓冲区布局注释llm_litert_mtp_drafter.cc草稿/验证/接受全流程实现llm_litert_compiled_model_executor.ccDecode 入口与 MTP 分支llm_executor_settings.henable_speculative_decoding开关llm_litert_mtp_drafter_test.cc单元测试与假模型构造小结LiteRT-LM 的 MTP 投机解码把小模型猜、大模型查的协作做到了工程化——贪心采样保证可确定性比对异步执行隐藏延迟约束状态全程同步配合内置命中率统计让端侧大模型在几乎不增加内存的前提下实现无损提速。读懂这 800 多行的 Drafter 源码你就掌握了投机解码在真实生产框架中的完整落地方式。【免费下载链接】LiteRT-LMLiteRT-LM is Googles production-ready, high-performance, open-source inference framework for deploying Large Language Models on edge devices.项目地址: https://gitcode.com/GitHub_Trending/li/LiteRT-LM创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表