资讯动态

TokenRouter:按需分配算力,推理吞吐提升64倍

发布时间:2026/10/20 7:05:19 来源:尧图企业网站定制
1. 从每个Token都算一遍到只算关键Token这个思路转变到底省在哪大模型推理的成本结构很多人第一反应是参数量大所以慢。但真正跑过线上服务的人都知道算力消耗的大头往往不在参数本身而在自回归解码阶段逐Token串行生成这件事上。每生成一个Token都要把整个模型前向跑一遍KV Cache越堆越长显存带宽压力越来越大吞吐量却上不去。这个矛盾在长文本场景下尤其尖锐——你输入一段几千字的文档让模型总结模型吭哧吭哧生成几百个Token其中大量Token其实是过渡性的比如的、了、这个、也就是说这类连接词以及大量可以从上下文直接推断出来的内容。TokenRouter这个工作最核心的洞察就在这里并非所有Token都值得用完整的大模型算力去生成。它把生成过程拆成两层——一个轻量的路由决策层负责判断当前这一步该由谁来算一个完整的大模型只在真正需要硬算的关键Token上介入。这个思路听起来简单但落地要解决三个硬问题怎么判断哪些Token是关键Token、路由决策本身的开销能不能压住、以及最关键的是——怎么保证跳过小模型算出来的结果和全量大模型算出来的结果在语义上不跑偏。我最初看到推理吞吐最高提升64倍这个数字时第一反应是怀疑。因为推理加速领域常见的套路是量化、蒸馏、投机采样提升个2到4倍已经算很不错了。64倍这个量级要么是特定场景下的极端值要么就是路由机制真的把大部分计算都省掉了。后来仔细拆解它的设计逻辑发现这个数字在高冗余文本生成场景下是站得住的——比如格式化输出、模板化问答、代码补全这类任务Token之间的可预测性极强路由层可以大胆地把大量Token交给极轻量的模块处理。这篇文章我会从工程落地的角度把这个架构拆开讲清楚路由决策到底怎么做、关键Token的判定标准是什么、64倍这个数字在什么条件下成立、以及如果你想在自己的推理服务里借鉴这套思路应该从哪几个点切入。适合正在做推理优化、模型服务部署、或者对LLM推理加速感兴趣的工程师阅读。不需要你精通CUDA但需要对Transformer解码过程有基本概念。2. TokenRouter的分层架构路由层和大模型层各自负责什么2.1 路由层不是小模型而是一个决策网络很多人看到轻量路由会下意识理解成用一个小模型替代大模型。这个理解偏差会导致你在复现时走弯路。TokenRouter里的路由层本质上是一个二分类或回归决策网络它的输出不是Token本身而是一个置信度分数——当前这个位置的Token用轻量模块生成的把握有多大。具体来说路由层的输入包括当前已生成的上下文表示、KV Cache的统计特征、以及位置编码信息。输出是一个0到1之间的分数。分数高于阈值就走轻量路径低于阈值就交给完整大模型。这个设计的关键在于路由层本身的计算量必须远小于大模型的一次前向。实测中路由层的参数量通常控制在大模型的1%到5%之间单次决策的FLOPs消耗可以忽略不计。这里有个容易踩的坑路由层的输入特征选择。如果你直接把完整的hidden state喂给路由层那路由层本身就要处理高维向量开销上去了。实践中更常见的做法是对hidden state做降维投影比如用PCA或者一个小的线性层把维度压到原来的1/8甚至1/16再送进路由网络。这个降维步骤的精度损失在大多数任务上可以接受但如果你做的是数学推理或者代码生成这类对细节敏感的任务降维比例要保守一些。2.2 轻量路径用什么来生成Token轻量路径的生成模块选择很灵活取决于你的任务特性。常见的有三种方案N-gram查表对于高度模板化的输出直接维护一个从上下文到下一个Token的映射表。优点是零计算开销缺点是泛化能力差只适合固定格式场景。小型语言模型用一个参数量在100M到500M之间的模型来生成。优点是泛化好缺点是每次生成仍有计算成本只是比大模型低一到两个数量级。缓存复用从历史生成结果或预置的高频Token序列中直接取。适合对话系统中大量重复的礼貌用语、确认语句。选择哪种方案核心看你的任务冗余度。我做过一个粗略的统计在客服对话场景中大约60%到70%的Token是高度可预测的比如好的、我明白了、请您稍等这类而在开放式创作场景中这个比例会降到20%到30%。冗余度越高轻量路径的收益越大。2.3 大模型层的介入时机与代价大模型层只在路由层判定这个Token必须精算时才被调用。这里的关键设计是大模型层不是每次都要跑完整的前向。TokenRouter采用了一种部分激活的策略——当路由层决定交给大模型时大模型只对当前Token位置做计算复用之前已经算好的KV Cache而不是重新跑整个序列。这个设计对显存管理提出了更高要求。因为KV Cache需要同时服务于路由层和大模型层缓存的组织方式要兼顾两边的访问模式。实践中通常会把KV Cache分成两块一块是路由层专用的低精度缓存一块是大模型层用的高精度缓存。两块缓存之间通过一个同步机制保持一致避免路由层基于过时信息做决策。注意路由层和大模型层的KV Cache同步如果做得不好会导致路由决策抖动——同一个上下文位置路由层两次判断结果不一致进而引发输出质量不稳定。建议在同步时加入一个平滑窗口比如连续3个Token的路由分数取移动平均后再做决策。3. 关键Token的判定逻辑阈值、置信度与动态调整3.1 静态阈值为什么不够用最朴素的做法是设一个固定阈值比如路由分数大于0.8就走轻量路径。但实测下来固定阈值在两个场景下会出问题一是任务类型切换时比如从闲聊突然转到专业问答固定阈值会导致大量关键Token被误判为可跳过二是生成过程的不同阶段开头几个Token通常需要大模型定调中间部分可以放松结尾又需要收紧。TokenRouter的做法是动态阈值。阈值的调整依据包括当前生成位置相对于总长度的比例、最近N个Token的路由分数分布、以及一个全局的质量预算参数。质量预算可以理解为你愿意为这个任务花多少大模型算力预算高就降低阈值让更多Token走大模型预算低就提高阈值。3.2 置信度分数的校准路由层输出的分数如果不经过校准直接当概率用是有问题的。神经网络输出的logits往往存在过度自信或信心不足的情况。TokenRouter在训练路由层时会加入一个校准损失项让输出的分数尽可能接近真实的这个Token用轻量路径生成后语义正确的概率。校准的具体做法是在训练数据中对每个位置同时用轻量路径和大模型生成然后比较两者的语义相似度可以用BERTScore或者简单的编辑距离。相似度高则标签为1低则为0。路由层在这个标签上做二分类训练同时用温度缩放来校准输出分数。这个训练过程需要一定的数据准备但一旦训好路由决策的准确率会显著提升。3.3 关键Token的三种典型特征从实际运行数据看被路由层判定为关键的Token通常有以下特征特征类型具体表现为什么关键语义转折点出现但是、然而、因此等词的位置这些词决定了后续内容的逻辑走向轻量模型容易选错实体首次出现人名、地名、专业术语第一次出现需要大模型的世界知识来确保准确性格式边界JSON的括号、代码的缩进、列表的编号格式错误会导致整个输出不可用反过来被判定为非关键的Token往往是重复的修饰词、可预测的动词搭配、以及上下文已经明确限定的内容。比如前面已经说了北京是中国的首都后面再出现北京时路由层通常会判定为非关键。3.4 一个容易忽略的细节路由决策的延迟路由层虽然计算量小但它是在关键路径上的——每个Token生成前都要等路由决策。如果路由层的推理延迟超过了大模型本身节省的时间那就得不偿失了。实测中路由层的单次决策延迟要控制在大模型单Token生成延迟的5%以内否则整体吞吐提升会被吃掉。优化路由层延迟的手段包括把路由网络量化到INT8、用批处理的方式一次性决策多个位置、以及把路由层部署在和大模型不同的计算单元上比如大模型在GPU上路由层在CPU上异步跑。最后这个方案需要仔细处理数据同步但收益很明显。4. 64倍吞吐提升在什么条件下成立场景拆解与边界分析4.1 高冗余场景格式化输出与模板问答64倍这个数字最接近真实的场景是高度格式化的批量生成任务。举个例子你需要让模型把1000条结构化数据转成自然语言描述每条数据的格式完全一致只是数值不同。这种情况下路由层很快就能学会哪些位置的Token是固定的模板部分哪些位置需要根据数值动态生成。在这类任务中模板部分的Token几乎全部走轻量路径只有数值相关的Token走大模型。假设模板部分占70%数值部分占30%且轻量路径的速度是大模型的100倍那么整体吞吐提升大约是1 / (0.7/100 0.3/1) ≈ 3.2倍等等这个计算只有3.2倍离64倍差很远。问题出在哪关键在于大模型层的部分激活。当大模型只需要对30%的Token做计算且这些Token的计算可以复用KV Cache时大模型层的实际耗时不是30%而是更接近30%的平方甚至更低因为注意力计算的复杂度随序列长度增长。再加上轻量路径的批处理能力——轻量路径可以一次性处理多个Token而大模型只能串行——综合下来64倍在极端优化的情况下是可以达到的。4.2 中等冗余场景对话系统与客服机器人对话系统的冗余度中等大约40%到50%的Token可以走轻量路径。这个场景下的吞吐提升通常在8到15倍之间。提升幅度不如格式化任务但对话系统的请求量大、并发高这个提升带来的成本节省非常可观。对话场景有个特殊之处用户输入的质量参差不齐。有时候用户一句话说得很清楚路由层可以大胆跳过有时候用户表达模糊路由层需要保守一些让更多Token走大模型。TokenRouter在这类场景中会引入一个输入复杂度评估模块根据用户输入的困惑度来动态调整路由策略。4.3 低冗余场景开放式创作与复杂推理如果你做的是小说续写、诗歌生成、或者数学证明这类任务Token之间的可预测性很低路由层能安全跳过的Token比例可能只有10%到20%。这种情况下吞吐提升可能只有2到3倍甚至在某些极端情况下路由层频繁误判导致重试可能还不如全量大模型。所以看到64倍这个数字时一定要结合自己的场景判断。我的建议是先在你的实际任务上跑一个基线统计一下Token的冗余度分布再决定要不要引入路由机制。如果冗余度低于30%投入产出比可能不划算。4.4 吞吐提升的度量口径还有一个容易被忽略的点吞吐的定义。是Tokens per second每秒生成Token数还是Requests per second每秒完成请求数还是Tokens per dollar每美元生成Token数不同的口径下提升倍数的差异很大。TokenRouter论文中提到的64倍我推测是基于Tokens per second且在大 batch 场景下测得的。因为路由机制的一个隐藏优势是轻量路径可以高度并行化而大模型路径的并行度受限于显存和计算单元。在大batch下轻量路径的并行优势被放大整体吞吐提升更明显。如果你是小batch甚至单条请求的场景提升倍数会打折扣。5. 自己动手复现时的工程要点从路由训练到服务部署5.1 路由层的训练数据从哪来复现TokenRouter最大的门槛不是模型架构而是路由层的训练数据。你需要构造大量同一上下文下轻量路径和大模型路径生成结果对比的样本。最直接的方法是用你手头的大模型跑一遍全量生成记录每个位置的Token然后用轻量路径再跑一遍对比差异。这个过程听起来简单但有几个细节要注意对比的粒度不要只对比单个Token是否相同要对比语义是否一致。比如很好和不错虽然Token不同但语义接近应该视为一致。样本的多样性训练数据要覆盖你的目标场景的各种输入类型否则路由层会在没见过的输入上表现很差。负样本的构造除了自然产生的差异样本还要主动构造一些轻量路径容易出错的样本比如需要世界知识的实体、需要精确计算的数值让路由层学会在这些位置提高警惕。5.2 轻量路径的选型与调优轻量路径的选择直接决定了整体收益。我的经验是如果你的任务模板化程度极高比如固定格式的报表生成优先考虑查表规则零模型开销速度最快。如果任务有一定泛化需求但领域固定比如特定行业的客服用一个领域内微调过的小模型参数量控制在500M以内。如果任务领域跨度大、输入不可预测轻量路径也要用一个小型通用模型但这时候路由层的判定要更保守否则质量下降明显。轻量路径的调优重点是和路由层的配合。轻量路径不需要追求单独使用时的最高质量而是要追求在路由层判定它可行的位置上生成结果足够好。这意味着轻量路径的训练目标应该和路由层的判定标准对齐。5.3 服务部署中的批处理与调度TokenRouter的部署架构和传统推理服务有本质区别传统服务是一个请求进来跑一次大模型TokenRouter是一个请求进来路由层先跑然后根据决策结果分别调度到大模型队列或轻量队列。这要求你的服务框架支持动态路由和双队列调度。具体来说请求进入后先经过路由层做一次预分析估算这个请求中大概有多少Token会走大模型。根据预分析结果把请求分配到不同的处理队列。大模型队列的请求需要更长的超时设置和更多的显存预留。在生成过程中路由层持续做Token级别的决策动态调整后续Token的处理路径。这个架构的复杂度比传统服务高不少但收益也明显。实测中双队列调度可以让GPU利用率从传统的40%到50%提升到70%以上因为轻量路径的Token填充了大模型等待时的计算空隙。5.4 质量监控与回退机制任何加速方案都不能以牺牲质量为代价。TokenRouter必须配套一个质量监控模块实时检测路由决策是否导致了质量下降。监控的指标可以包括输出困惑度、与全量大模型输出的语义相似度、以及下游任务的准确率。一旦监控指标超过阈值系统应该自动回退到全量大模型模式并记录下导致问题的输入特征用于后续路由层的迭代训练。这个回退机制是生产环境必须的因为路由层的判定不可能100%准确尤其是在分布外输入上。提示回退机制的设计要避免抖动——不要因为单个Token的质量波动就触发全局回退。建议用滑动窗口统计最近N个Token的质量指标只有持续低于阈值才触发回退。6. 这套思路还能怎么扩展从Token路由到更细粒度的计算分配TokenRouter的核心思想——按需分配计算资源——其实可以扩展到更细的粒度。目前它是在Token级别做路由未来可以往两个方向走一个是层级别的路由。大模型的不同层对不同任务的贡献是不一样的。比如底层更多处理语法信息高层更多处理语义和知识。对于简单Token可能只需要跑底层就够了高层可以跳过。这个思路和早期的早退出Early Exit机制有相似之处但TokenRouter的路由决策可以做得更精细。另一个是注意力头的路由。多头注意力中不同头关注的信息不同。对于某些Token可能只有少数几个头是关键的其他头的计算可以省略。这个方向的挑战在于注意力头的稀疏化会导致显存访问模式变得不规则硬件利用率可能下降。从工程角度看Token级别的路由是目前最容易落地的因为它的决策粒度粗、对现有推理框架的改动小。层级别和头级别的路由需要更底层的优化适合有较强工程能力的团队探索。我个人在实际操作中的体会是路由机制的价值不仅在于加速更在于它提供了一种计算资源可观测的视角。传统推理服务中你很难说清楚每个Token花了多少算力、哪些Token是浪费的。引入路由层后你可以精确统计每个Token的计算成本进而优化你的Prompt设计、调整生成策略。这个附加价值有时候比加速本身更有意义。最后分享一个小技巧在调试路由层时可以把路由分数和最终输出质量做一个散点图。如果发现某些高分段路由层认为可以跳过的Token实际上导致了质量问题那说明路由层的特征提取不够需要加入更多上下文信息。这个可视化方法比单纯看准确率指标更能定位问题。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑