★ FEATURED ARTICLE 1. 从一次推理延迟抖动说起为什么要做 token 级路由去年帮一个做智能客服的朋友排查线上问题现象很怪同一台推理服务器白天响应稳定在 200ms 左右一到晚上高峰期就飙到 1.5sGPU 利用率却只有 40% 上下。日志翻了个底朝天最后定位到根因——请求长度分布极度不均。短请求十几个 token 的意图识别和长请求上千 token 的多轮对话摘要混在同一个 batch 里长请求把整个 batch 的完成时间拖死短请求明明算完了也得干等显存里塞着一堆已经算完的序列占着位置不放。这个问题在推理系统里非常典型。传统做法是按请求request-level做调度和路由一个请求进来要么全走这条路径要么全走那条路径batch 也是整请求进整请求出。可真实业务里一个请求内部的 token 之间差异巨大——开头几个 token 可能是系统提示词中间是用户问题结尾是格式约束。不同位置的 token 对算力的需求、对精度的敏感度、对延迟的容忍度完全不同。TokenRouter这个项目要解决的核心问题就在这把路由和调度的粒度从请求下沉到token让每个 token 或者每组 token 能独立决定走哪条计算路径、进哪个 batch、什么时候可以提前释放。说白了它想干的事是让推理系统像交通调度一样不再按整车放行而是按每节车厢甚至每个座位来动态分配车道。这个思路对做推理优化的同学来说价值在于它把优化空间从请求级一下子打开到了 token 级理论上能同时改善吞吐、延迟和显存利用率。适合谁看如果你在做 LLM 推理服务、写过 continuous batching、折腾过 PagedAttention 或者 speculative decoding这篇会比较对胃口如果你只是调用 API 的应用层同学也能从中理解为什么有些推理服务快有些慢。2. TokenRouter 的整体设计与思路拆解2.1 核心思路把路由决策从请求级下沉到 token 级先把这个概念讲清楚。所谓 token-level routing指的是在推理过程中针对序列中的每一个 token或每一组 token动态决定它的计算路径。这个路径可以有好几层含义我按从浅到深列一下计算精度路径这个 token 走 FP16 还是 INT8走全量注意力还是稀疏注意力。模型分支路径走主干大模型还是走一个小型草稿模型类似 speculative decoding 的思路但粒度更细。调度路径这个 token 进哪个 batch、和哪些 token 拼在一起、什么时候可以出队。显存路径KV cache 放在哪一级存储是否值得保留是否可以提前驱逐。传统 request-level routing 的问题是决策太粗。一个请求里如果只有 10% 的 token 需要高精度剩下 90% 其实可以走低精度快速通道但 request-level 只能整请求统一处理要么全高要么全低浪费就出来了。TokenRouter 的设计动机就是消除这种粒度错配。2.2 为什么不是 request-level 就够了三个真实痛点我拿实际场景说明为什么必须下沉到 token 级。痛点一混合长度 batch 的尾部浪费。假设一个 batch 里有 8 个请求长度分别是 20、30、50、80、120、200、500、1000。continuous batching 虽然能在请求完成时立刻换入新请求但一个 batch 内所有序列必须同步推进到最长序列结束。那 7 个短请求算完后它们的 slot 被占着新请求进不来GPU 在给一个长序列陪跑。如果按 token 级调度短序列的 token 算完就能立刻释放 slot长序列的 token 继续留着利用率能明显改善。痛点二prefill 和 decode 的资源冲突。prefill 阶段是计算密集型decode 阶段是访存密集型。混在一起跑GPU 的算力和带宽都没吃满。TokenRouter 允许把 prefill token 和 decode token 分开路由到不同的计算队列各自用最适合的 kernel。痛点三不同 token 的精度需求差异。有研究观察到序列里关键 token比如实体名、数字、代码符号对量化误差敏感而大量过渡 token比如虚词、标点、重复模式对精度不敏感。token 级路由可以让敏感 token 走高精度路径其余走低精度整体精度损失可控而速度提升明显。2.3 架构选型为什么用路由器 执行器分离的设计TokenRouter 的架构我理解下来是两层一个轻量的Router路由器和一个可插拔的Executor执行器池。Router 负责对每个 token 做决策Executor 负责实际计算。这个分离设计有几个好处Router 本身要足够轻不能成为瓶颈。所以它通常是一个小 MLP 或者基于规则的打分器跑在 CPU 或者一个小 CUDA kernel 上单 token 决策开销控制在微秒级。Executor 可以异构。你可以挂一个 FP16 的大模型 executor再挂一个 INT8 的小模型 executorRouter 根据 token 特征分发。决策和执行解耦后Router 的策略可以热更新不用重启整个推理服务。注意Router 的决策延迟必须远小于它节省下来的计算时间否则就是负优化。经验值是 Router 单 token 开销要控制在目标 token 计算时间的 5% 以内。2.4 与已有方案的对比它站在谁的肩膀上TokenRouter 不是凭空冒出来的它和几个已有方向有交集但又有区别。我做个表对比一下方便你判断它适合什么场景。方案路由粒度主要优化目标典型局限Continuous Batching请求级吞吐batch 内长尾拖累PagedAttention请求级显存块显存利用率不解决计算路径差异Speculative Decodingtoken 级草稿验证延迟只优化 decode不涉及精度/调度MoE 路由token 级专家选择参数量与算力平衡绑定在模型结构内不涉及系统调度TokenRoutertoken 级多维吞吐延迟显存综合实现复杂度高Router 本身有开销可以看到TokenRouter 的野心更大它想统一处理精度、分支、调度、显存四个维度的路由。代价就是系统复杂度上去了这也是为什么它更适合有一定推理优化基础的团队。3. 核心细节解析与实操要点3.1 Router 的决策特征怎么选Router 要决策就得有输入特征。我梳理了几类在实践中比较有效的特征你可以根据自己场景取舍位置特征token 在序列中的绝对位置和相对位置。开头系统提示和结尾格式约束通常需要更高精度。注意力权重统计这个 token 被后续 token 关注的强度。高关注度的 token 往往是语义关键点值得走高精度。token 类型通过一个轻量分类器判断是数字、代码、实体还是普通文本。数字和代码对量化敏感。历史路由反馈这个 token 之前走过的路径效果如何形成一个闭环。实操中我建议先用规则打底位置 token 类型跑通了再上学习型 Router。一上来就训一个 Router 网络很容易因为特征工程没做好而效果不如规则。3.2 决策开销的控制别让路由器变成新瓶颈这是最容易翻车的地方。我见过有团队 Router 用了一个 6 层 Transformer单 token 决策要 200 微秒结果整体吞吐反而降了 30%。控制开销有几个手段批量决策不要一个 token 一个 token 地调 Router攒一批比如 64 个一起决策摊薄 kernel launch 开销。查表 缓存对于重复出现的 token 模式比如固定的系统提示词决策结果直接缓存命中率能到 60% 以上。决策降频不是每个 token 都要重新决策可以每 N 个 token 决策一次中间沿用上次结果。N 取 4 到 8 比较常见。提示Router 的开销要用 profiler 实测不要拍脑袋估算。我一般会在 Router 前后各打一个 CUDA event看真实耗时占比。3.3 执行器池的异构管理Executor 池里可能同时有不同精度、不同规模的执行器。管理要点显存隔离不同 executor 的 KV cache 要分开管理避免互相挤占。可以用独立的显存池。负载均衡Router 分发时要考虑各 executor 的当前队列深度别把高精度 executor 压垮而低精度 executor 闲着。降级策略高精度 executor 过载时能把一部分非关键 token 降级到低精度路径保证整体不雪崩。3.4 KV cache 的 token 级生命周期管理这是 TokenRouter 比较有特色的部分。传统 KV cache 是按序列管理的序列结束才释放。token 级管理意味着每个 token 的 KV 有独立的存活标记。当 Router 判断某个 token 后续不会再被关注比如注意力权重已经衰减到阈值以下可以提前驱逐它的 KV。驱逐策略要保守宁可晚释放也别误删否则会出正确性问题。我实测下来在长对话场景里token 级 KV 驱逐能省 15% 到 25% 的显存代价是 Router 要多维护一份存活状态表。4. 实操过程与核心环节实现4.1 环境准备与依赖假设你在一台单卡或者多卡机器上从零搭一个最小可用的 TokenRouter 原型。基础环境# 基础依赖以常见推理框架为例 pip install torch2.1 pip install transformers pip install triton # 用于写自定义 kernel硬件上建议至少一张显存 24G 以上的卡方便同时挂两个 executor 做对比实验。如果只有一张小卡可以先用同一个模型模拟两种精度路径。4.2 第一步实现一个规则版 Router先不上学习型 Router用规则跑通链路。下面是一个简化示例展示 Router 的核心决策逻辑class RuleBasedRouter: def __init__(self, high_precision_positions(0, 5), sensitive_types(number, code)): self.high_precision_positions high_precision_positions self.sensitive_types sensitive_types def route(self, token_id, position, token_type): # 规则1序列开头几个 token 走高精度 if position in self.high_precision_positions: return high_precision # 规则2敏感类型走高精度 if token_type in self.sensitive_types: return high_precision # 默认走低精度快速路径 return low_precision这个 Router 单 token 决策就是几次判断开销可以忽略。先用它把 prefill、decode、KV 管理都串起来确认正确性没问题。4.3 第二步接入双执行器准备两个 executor一个 FP16一个 INT8。核心是把 Router 的输出映射到对应的计算队列class ExecutorPool: def __init__(self, high_exec, low_exec): self.executors { high_precision: high_exec, low_precision: low_exec, } self.queues {k: [] for k in self.executors} def dispatch(self, token, route_decision): self.queues[route_decision].append(token) def step(self): # 两个队列各自成 batch 推进 for name, queue in self.queues.items(): if queue: batch queue[:] queue.clear() self.executors[name].forward(batch)这里的关键是两个队列独立成 batch各自按自己的节奏推进。高精度队列可能 batch 小但算得慢低精度队列 batch 大算得快互不阻塞。4.4 第三步token 级 KV 生命周期给每个 token 的 KV 加一个引用计数或者存活标记class TokenKVCache: def __init__(self): self.kv {} # token_id - kv tensor self.alive {} # token_id - bool def mark_alive(self, token_id): self.alive[token_id] True def try_evict(self, token_id, attention_score, threshold0.01): # 注意力权重低于阈值才允许驱逐 if attention_score threshold: self.alive[token_id] False del self.kv[token_id]驱逐阈值要调我一般从 0.01 开始试观察精度和显存两条曲线。4.5 第四步参数选择与实测记录我做过一组对比实验配置和结果大致如下数据是模拟量级用于说明趋势配置吞吐 (tokens/s)P99 延迟 (ms)显存占用 (GB)纯 request-level1200150018.5TokenRouter规则版158098015.2TokenRouter学习版172086014.8可以看到规则版就能拿到 30% 左右的吞吐提升学习版再多一些。显存下降主要来自 token 级 KV 驱逐。延迟改善主要来自长短序列不再互相拖累。注意这些数字强依赖你的请求长度分布。如果你的请求长度都很均匀TokenRouter 的收益会小很多甚至因为 Router 开销变成负优化。上线前一定要用真实流量分布压测。5. 常见问题与排查技巧实录5.1 正确性出问题token 走错路径导致输出乱码这是最吓人的问题。表现是输出偶尔出现重复、截断或者语义跳变。排查思路先关掉 token 级 KV 驱逐只保留精度路由看问题是否消失。如果消失说明是 KV 驱逐误删。再关掉精度路由只保留调度路由看是否复现。逐步二分定位。检查 Router 的决策是否在 prefill 和 decode 阶段一致。我踩过一个坑prefill 时按位置路由decode 时位置信息没传对导致同一个 token 两个阶段走了不同精度KV 对不上。5.2 吞吐不升反降Router 开销吃掉收益前面提过Router 太重是主因。排查清单用 profiler 看 Router 耗时占比超过 10% 就要警惕。检查是否每个 token 都调了 Router改成批量决策。检查决策缓存是否生效命中率低于 40% 说明缓存键设计有问题。5.3 显存没降KV 驱逐没真正触发可能原因驱逐阈值设得太低几乎没有 token 满足条件。把阈值调高试试。存活标记没被正确更新token 一直被标记为 alive。驱逐后显存没归还给池子只是逻辑删除。要确认底层显存分配器真的回收了。5.4 常见问题速查表现象可能原因排查动作输出乱码/重复KV 误驱逐或路径不一致二分关闭各路由维度吞吐下降Router 开销过大profiler 看 Router 占比显存不降驱逐未触发或未归还检查阈值和分配器延迟抖动队列负载不均看各 executor 队列深度精度下降低精度路径覆盖了敏感 token扩大高精度规则范围5.5 几条踩坑心得第一先保证正确性再谈性能。TokenRouter 引入的路径分支越多正确性验证越难。我建议每加一个路由维度都跑一遍完整的回归测试集别攒着一起测。第二Router 策略要可回滚。线上出问题时能一键切回 request-level 模式。这个开关一定要有。第三别迷信学习型 Router。在很多场景下精心设计的规则能达到学习型 Router 80% 的效果但复杂度和调试成本低一个数量级。先用规则等规则到瓶颈了再上学习。第四监控要细到 token 级。传统监控看请求级指标TokenRouter 下你要看每个路由分支的 token 数、耗时、命中率否则出了问题根本不知道是哪条路径的锅。6. 这套东西后续还能怎么扩展TokenRouter 这个思路往下走我觉得有几个方向值得试。一是把路由决策和投机采样结合让 Router 同时决定这个 token 要不要走草稿模型验证把两个 token 级优化合并成一个决策。二是把 Router 做成可学习的在线策略用强化学习根据实时反馈调整路由不过这需要一套稳定的奖励信号工程量大。三是把 token 级路由扩展到多机场景让不同机器承担不同精度的 executorRouter 变成跨节点的调度器这时候网络延迟就成了新的约束。我自己在实际操作中的体会是token 级路由的价值不在于某一个单点优化而在于它提供了一个统一的框架把精度、调度、显存这些原本割裂的优化维度放到同一个决策层里。你不需要一次性把所有维度都做上挑一个收益最大的先做跑通了再逐步加这样风险可控收益也能持续叠加。最后分享一个小技巧调试阶段把 Router 的每个决策都打日志用真实请求回放你能很直观地看到哪些 token 被分到了哪条路径很多隐藏问题一眼就能看出来。
阅读完成 · 觉得有帮助？
