TokenSpeed调度器深度剖析有限状态机如何从编译期保证KV资源安全【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeedTokenSpeed 是一款主打光速speed-of-light性能的大模型推理引擎而它的核心引擎之一 tokenspeed-scheduler 用纯 C 实现了调度器并通过**有限状态机FSM**把每条请求的 KV Cache 页面归属写进类型系统——非法状态切换在编译期就被排除异常路径在运行期直接抛错。本文将带你理解这套设计如何让 GPU 显存分配永远对号入座。一、为什么 LLM 推理引擎最怕KV 页面失控在 LLM 推理中KV Cache 是显存里最昂贵也最容易被踩的资源Prefill 分块长提示词按 chunk 分片计算页面归属必须跨多轮保持Retract 重调度容量不足时调度器会驱逐在途请求回收其页面给新请求PD 分离Prefill 节点的页面要跨越节点传输给 Decode 节点。任何一个环节忘记释放、重复释放或写错页面都会造成显存泄漏或数据污染。传统写法依赖人肉约定 运行时断言而 TokenSpeed 选择把约定下沉到C 类型系统。二、每个请求一个状态机9 个状态覆盖完整生命周期TokenSpeed 调度器把每个请求抽象为一个状态变量定义在 states.h状态含义是否持有 KV 页面BootstrappingPD 分离下等待对端建链❌Submitted已提交、尚未分配页面❌Prefilling本地分块 prefill 进行中✅RemotePrefilling对端 P 节点正在 prefill✅PrefillAwaitingResult等待末块结果落地✅PrefillDoneprefill 完成等待首个 decode✅Decoding逐 token 生成中✅Retracted已被驱逐等待重新准入❌页面已归还Finished结束❌其中RemotePrefilling、PrefillAwaitingResult等仅在 PD 分离部署下出现见 pd_states.h融合引擎永远不会进入这些状态——状态集合本身就编码了部署形态的约束。三、资源束ForwardResources页面整包移交不存在第三种路径编译期安全的第二个支柱是资源束设计。forward_states.h 中的ForwardResources把 KV 页面相关的四样东西捆成一个 move-only 整体block_tables该请求占用的 KV 页表req_pool_index请求池槽位cache_progress前缀缓存推进进度results_in_flight已发出、结果未回的前向次数。源码注释写得很直白状态转移要么把整包资源移交给唯一后继状态要么把页面归还缓存协调器后让空束消亡——不存在第三条路径也不存在任何字段会在转移途中被遗忘。这正是从编译期保证安全的第一层含义持有页面的状态必然携带资源束字段级别的遗漏在结构上就不可能发生。四、重载 兜底抛错非法转移被结构性拒绝TokenSpeed 的事件SchedulePrefillEvent、ScheduleDecodeEvent、RetractEvent等是实现了operator()访问器重载的结构体见 forward_events.h每个事件只对合法来源状态定义重载例如ScheduleDecodeEvent只接受PrefillDone、PrefillAwaitingResult和Decoding三种状态。未定义的组合则落入统一兜底 InvalidTransitionHandler直接抛出带事件名和状态名的std::logic_error。这套白名单重载 兜底异常带来三重收益签名即契约事件的返回类型声明了合法后继状态的variant新增/删减状态时编译器会立即暴露所有需要适配的转移零运行时判断合法路径没有if (state ...)分支std::visit直接分派错误即崩溃而非静默非法转移如Submitted上触发 decode立即抛错不会带着脏状态继续跑下去。五、驱逐Retract时刻KV 资源安全的关键战场容量不足时调度器会驱逐受害请求这是最容易出 bug 的时刻。TokenSpeed 的处理方式是让RetractEvent走一次普通状态转移forward_events.h页面随转移当场归还缓存协调器且在同一轮计划构建中把释放的容量直接授予被阻塞的准入请求受害者在结果未回results_in_flight 0或 PD 传输仍钉住页面时不可被驱逐——因为 in-flight 计数就住在资源束里随页面一起走任何页面还在、计数丢了的组合在类型上构造不出来被驱逐请求进入Retracted后带上retraction_epoch时间戳重准入顺序直接由状态本身推导无需额外维护队列。完整的准入、驱逐与恢复协议可以在官方设计文档 docs/design/scheduler.md 中逐节对照阅读例如被驱逐者可能正是阻塞者本人这类边界情况都有明确的规则陈述。六、从编译期到测试台保证如何被验证这套 C 调度器通过 nanobind 暴露给 Python构建配置见 pyproject.toml并配有两层测试Python 层test_fsm_and_scheduling.py 覆盖Submitted → Prefilling → Decoding → Finished全链路及分块 prefill 的多轮计划C 层tests/cpp/ 下的 test_kv_cache_lifecycle.cpp、test_scheduler_plan.cpp 等直接对状态转移与页面生命周期做单元断言。调度器省下的显存管理开销最终体现为推理吞吐与延迟——TokenSpeed 官方基准中 Kimi K2.5 等模型的调度效率表现如下图所示图同前此处不再重复。七、小结把约定变成类型TokenSpeed 调度器给出的方法论值得所有管理稀缺资源显存、连接、锁的系统借鉴状态即资源持有 KV 页面的状态强制携带资源束字段遗漏在构造上不可能重载即白名单合法转移写成operator()重载编译器替你审查状态图兜底即护栏非法转移统一抛错杜绝静默腐化驱逐走转移高危操作retract复用同一套转移机制页面归还与容量授予在同一轮完成。对新手而言这套设计的启示是当一段业务规则反复依赖记住别漏掉某一步时就该考虑把它编码进类型系统——让错误代码无法编译比让错误代码无法运行更难被绕过。【免费下载链接】tokenspeedTokenSpeed is a speed-of-light LLM inference engine.项目地址: https://gitcode.com/gh_mirrors/to/tokenspeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考