构建面向 Agentic 推理的下一代 LLM 推理引擎

InfoQ 中文 2026-11-03T23:40:16+08:00

随着 Agent 类负载快速增长，为聊天场景设计的推理引擎在吞吐目标、缓存复用和硬件适配上逐渐失效。TokenSpeed 独立贡献者尚旭春在 AICon 2026 深圳站分享了该项目的设计思路：从 Agent 负载特征出发，重构调度、缓存与算子层，并给出在 Blackwell 平台上的阶段性性能数据。

为什么 Agent 场景需要重新设计推理引擎

Chat 与 Agent 是两类负载

传统 Chat 负载的优化目标相对单一：用户发一次请求，模型生成一段完整回复。系统只需尽量把 batch 打满、最大化吞吐，KV Cache 加载一次后即释放，复用率天然偏低。

Agent 场景则不同。以 coding agent 为例，模型要持续思考、调用工具、拿到结果再继续思考，整个过程由多轮短请求循环构成。第一次请求可能携带几十 K 的前缀上下文，但后续每轮输出往往只有几百个 token。推理引擎因此必须重新审视缓存复用策略、请求调度机制和整体优化目标。

延迟约束下的吞吐

Agent 场景对延迟更敏感。工具调用链中一旦推理环节明显滞后，用户会直接感到“不跟手”。因此真正要优化的是延迟约束下的吞吐。TokenSpeed 内部设定的指标是：每用户 TPS 不低于 70，在此前提下尽量提升单 GPU 吞吐。

以 TensorRT-LLM 为基准，TokenSpeed 在 batch size 为 1 的最小延迟场景下端到端延迟快 9%；在每用户 100 TPS 的典型工作点上，单 GPU 吞吐高 11%。

KV Cache 复用率决定推理成本

KV Cache 命中率是成本的关键变量。DeepSeek 的缓存复用率据称在 99% 以上，而业界多数系统目前只能达到 70% 到 80%。在 Agent 场景中，同一段几十 K 的上下文如果能被反复命中，每次命中都能省去 prefill 计算，系统整体成本会显著下降。

测试中，上下文长度从 128K 扩展到 1M 时，TokenSpeed 的解码吞吐率仅下降 16%。

线性注意力与 MOE 带来的新命题

越来越多模型开始采用线性注意力，例如千问的 Gated DeltaNet 和 Kimi K3 使用的 KDA。这类模型不再维持 token 与 KV 的一一对应关系，而是通过一个持续迭代的状态来表示上下文，其中包括 SSM 状态和卷积状态。如何在推理引擎中高效更新、保存和恢复这些状态，是新的系统设计问题。

MOE 自 DeepSeek 发布后已成为主流选择，但优化难度并不低。以 Kimi K3 为例，模型有 896 个专家，虽然采用 Latent MOE 技术大幅降低了成本，并把激活专家比例提升到 16%，但路由和通信复杂度明显增加。

硬件特性与 Runtime-Kernel 协同

Blackwell 这一代硬件引入了 Tensor Memory、新一代 Tensor Core 等特性，AMD 的 CDNA 架构也有其独特能力。推理引擎如果无法充分利用这些硬件特性，Blackwell 相对 Hopper 的性能优势就无法真正释放。TokenSpeed 因此在 Kernel 层针对 Blackwell 做了深度优化。

调度器与 Kernel 也需要协同设计：调度设计得再好，如果底层 Kernel 存在瓶颈或耦合过深，系统整体仍达不到理想状态；反过来，Kernel 优化到极致，如果调度器产生大量 bubble，吞吐也打不满。TokenSpeed 在架构上把调度、执行和底层算子做了系统性解耦与重构。

Runtime 架构：控制、执行、算子三层解耦

TokenSpeed 的整体架构分为三层：控制面、Runtime 和 Kernel。控制面完全用 C++ 编写，Runtime 使用 Python 实现，Kernel 则独立为一个子包。这种分层设计的核心目的是保持模块边界清晰：

三者各司其职，避免了过去推理引擎中调度逻辑、执行逻辑和算子实现相互渗透所导致的维护困境。对关注 AI 编程与氛围编程的开发者而言，Agent 推理的延迟表现和缓存命中率，直接决定了工具链的使用体验与调用成本。

控制面、Runtime、Kernel：三层各用各的语言

控制面最终选择 C++，有实际的历史原因。团队成员多从 C++ 时代过来，对 Rust 的熟悉程度有限；而 C++ 在性能和表达力上已经够用，于是确定用它来写控制面。Runtime 保留 Python，是因为深度学习生态成熟，PyTorch、CUDA 等库集成度高，模型适配层的开发效率明显高于其他语言。Kernel 子包独立出来后，NVIDIA 和 AMD 的专家可以针对各自硬件做深度优化，不必理解调度器或 Runtime 的内部实现。

这种分层还带来一个额外收益：并行编译可以自动化。过去开发推理引擎，模型并行切分的代码往往散落在各个模型适配模块里，和注意力实现、MOE 实现、通信逻辑搅在一起。支持的模型一多，这些散落的并行代码就越来越难维护，接入一个新模型可能牵动多个模块。TokenSpeed 的做法是把并行切分逻辑从模型适配代码中完全抽出来，做成独立的 Placement 系统。模型开发者只需声明模块边界、输入输出类型，以及这个模块在并行维度上的属性——是复制（replicate）、切片（shard）还是部分和（partial），后续的切分和通信由系统自动处理。

具体来说，每个模块可以挂到不同的并行组上：注意力模块挂 Attention Tensor Parallel 组，MOE 模块挂 MOE 并行组，各自独立选择并行策略。基于这些 Placement 声明，类型编译系统会遍历整个解码层的执行计划，比较相邻模块的 Placement 是否一致。发现不一致，就自动插入对应的通信原语：从 partial 到 replicate 需要 all-reduce，也就是把各卡的部分结果求和后再广播；从 shard 到 replicate 需要 all-gather，把切片收集成完整张量；从 shard 到 partial 需要 reduce-scatter，先求和再切分。整个编译流程输出统一的 compiled_decoder_layer 执行计划，运行时分三步接入：prefill 阶段先做前置通信对齐，再执行模型算子，最后做通信输出。自动并行编译不仅降低了模型适配的维护成本，还因为类型系统提前给出了信息，让相邻算子的 fusion 成为可能，减少 kernel launch 开销，性能上也有正收益。

请求生命周期交给状态机

控制面的核心是一个覆盖单个请求全生命周期的状态机。一个推理请求从进入到完成，会经历几个关键状态：submitted、prefilling、decoding、retracting/timeout、finished。每个状态本身是一个资源容器，持有该状态所需的 KV 页引用、SSM/MLA 的 slot 列表等资源。状态迁移时，这些资源的持有权同步传给下一个状态。整个控制面就是一个持续流转的状态机系统。

状态迁移由事件驱动。每个事件承载对应的 C++ operator，资源所有权通过 C++ 的 unique_ptr move 语义传递——同一时刻只有一个持有者，转移后原持有者自动失效。这样做的好处是把运行时的资源管理问题提前到了编译期：代码能编译通过，说明编译器已经按类型系统检查过资源转移路径，资源泄漏和分配错误在编译阶段就会暴露。资源本身用 RAII 管理，即构造时获取、析构时自动释放；无论是 Cache Block 的引用计数，还是 Device Node 或 Host Side 的引用计数，在数据搬移或操作完成之前都不会被提前释放。

控制面状态机与执行面完全解耦。调度器只关心何时把请求从 prefilling 推进到 decoding、何时触发 retracting，不需要知道底层 Kernel 如何执行具体计算。调度策略因此可以独立迭代，不必因为算子调整而频繁修改。Runtime 层接收控制面给出的执行计划，按计划中的顺序和资源配置执行计算，完成后把结果返回控制面，驱动状态机继续流转。

工程上的收益很直接：并发请求调度中的竞态和资源泄漏明显减少。过去 Python 实现里常见的资源管理 bug，比如请求被取消后 KV 页没有及时回收，或者异常路径下引用计数没有正确递减，在 C++ 实现中靠编译期检查和 RAII 得到了系统性解决。这也是控制面完全用 C++ 重写的核心动机之一。

对做 AI 编程工具和氛围编程的团队来说，这种“人只声明意图、工具链负责切分与资源安全”的思路值得借鉴：把容易出错的细节交给编译器和运行时，上层代码才更好写、也更好改。

KV Cache 管理：从基数树走向平面存储

在整个 TokenSpeed 架构中，KV Cache 管理是团队投入精力最多的模块之一。原因在于，Agent 场景对前缀复用的要求远高于传统对话场景：同一个 Agent 请求在一轮完整工作流中会反复携带相同的前缀上下文，如果每次循环都重新做一遍 prefill 计算，推理成本会迅速膨胀。

第一版实现借鉴了 SGLang 的思路，用 Radix Tree（基数树，一种按 token 序列层级组织的树形结构）来管理前缀缓存。它的做法是维护一棵从根到叶的树，每个节点对应一段 token 序列，所有请求的前缀匹配都在这棵树上完成。这套设计在纯 attention 场景下运行良好，但在适配 Kimi K3 时，我们发现了问题：Radix Tree 的语义与部分新型注意力模式存在冲突。Mamba 类模型和 Sliding Window Attention（滑动窗口注意力）天然破坏了 Radix Tree 依赖的层级匹配假设，它们的缓存访问模式并不是从前缀到后缀的简单线性延伸。

这一经验促使我们重构了整个 KV Cache 系统。在支持 Kimi K3 的过程中，我们把原有的 Radix Tree 全部迁移到 Flat KV 设计。新的缓存系统不再关心缓存属于哪种注意力类型，只要求上层把不同类型的缓存按统一格式上报，声明各自的类型和占用空间；底层则把所有缓存统一打包成相同尺寸的 block，做平面化存储管理。

K3 的第一版采用逻辑 block 与物理 block 相等的方式，以 1536 个 token 为一个 block 进行管理。但这个粒度在长上下文场景下，前缀复用效率和页内碎片问题都比较突出。后续我们借鉴了 vLLM 的做法，引入两层 block 设计：第一层是逻辑块，大小统一为 128；第二层是物理块，大小为 1536。128 的逻辑块负责前缀复用和更细粒度的缓存匹配，物理层保持不变。这套双层设计在后来的测试中取得了不错的优化效果。

缓存抢占、恢复与 PD 分离

缓存抢占与恢复是另一项关键能力。TokenSpeed 目前完全接入 Mooncake 的 KV Store 系统：当显存不足时，请求的整个 KV 状态会被卸载到 KV Store 中持久化；等该请求重新被调度回来，再通过 lazy（延迟加载）化的 reload 机制重新载入并继续计算。

在 PD 分离（Prefill 与 Decode 分离）方面，我们最初实现的是逐层传输。问题在于 RDMA 带宽利用率非常低——每一层的传输都要付出独立的网络往返开销。后续优化允许用户指定把若干层打包成一个 block 传输，也就是一种可定制的分层传输策略。用户因此可以根据具体网络拓扑和模型规模自行调节传输粒度。

Kernel 层：Blackwell 上的 SOTA 与 MLA 的三个瓶颈

Kernel 层的优化主要由团队中的 NVIDIA 和 AMD 专家推动。目前 TokenSpeed 的 Kernel 在 Blackwell 上已达到 state-of-the-art 性能水平，并被 vLLM 社区集成为 Blackwell 上的默认 Kernel 包。

MLA（Multi-head Latent Attention，多头潜在注意力）是国内开源模型广泛采用的注意力机制。它把 KV 压缩成单一的 latent 状态，从而节省写入成本，但计算模式对硬件并不天然友好，主要有三个瓶颈。

第一是算力浪费。在 decode 阶段，query 与 latent key 的第一次矩阵乘法中 query 维度为 1，整个矩阵乘法实际退化成了向量乘法，Blackwell 上的 Tensor Core 计算资源因此被严重闲置。

第二是带宽。Agent 场景下前缀长度常达到几十 K 甚至几百 K，即便使用 INT8 量化，数据加载带宽仍是限制因素。

第三是形状不匹配。MLA 的 head 数与训练长度组合，往往让矩阵乘法呈现瘦高形态，对 SM（流式多处理器）利用率和 kernel launch 开销都产生负面影响。

![]

在 Blackwell 上重写 Kernel

围绕上述问题，TokenSpeed 在 Blackwell 上重新设计了 Kernel，主要做了三件事。

第一，用上 Blackwell 新引入的片上存储 Tensor Memory（张量内存）。团队为它单独写了分配器，让 Tensor Core 直接读写这块内存，跳过经寄存器搬运中间数据的开销。

第二，在矩阵乘法层面引入折叠因子（fold factor）。decode 阶段每个 query 只有一个 token，维度为 1，Tensor Core 的算力被大量浪费。当 head 数小于 128 时，可以把 query token 折叠进 head 维度，折叠后有效 head 维度乘以折叠因子，矩阵维度变高，Tensor Core 就能跑得更满。

第三，设计 Split KV 机制，把计算拆成两阶段：先算 Split KV 的部分和，再做归约。Split 参数和 workspace 大小在运行时自动计算并缓存；跨 SM 协同执行降低了 MMA barrier 的开销，数据写回也做了带宽友好处理。几项优化叠加后，TokenSpeed 的 Kernel 在 Blackwell 上达到领先性能。

投机解码：三种算法已支持，实测最高 2.33 倍

投机解码是当前提升推理吞吐的关键手段，EAGLE、DeepSeek 等算法都带来过明显收益。TokenSpeed 仓库目前支持三种投机解码算法：MTP、EAGLE3 和 DeepSeek；团队也在用 TokenSpeed 支持 DeepSeek 推理引擎中提出的投机解码算法，后续会为开源模型专门训练一批对应的投机模型权重，放进仓库供社区使用。

测试数据如下：

线性注意力模型：搬指针，不搬数据

对线性注意力模型，TokenSpeed 也做了针对性优化。这类模型的 MLA 状态会随 token 生成不断向前更新，到达某个截止点时必须写回，而写回本身有开销。TokenSpeed 的做法是不搬数据、只动指针：每一步的状态直接写入该请求的私有物理行，用一个索引标记状态位置；需要截止时只更新索引指向，不搬整个状态块，把开销压到最低。

prefill 场景同样特殊。Agent 负载的前缀命中率通常在 90% 以上，此时只需追加几百个 token 的计算，请求的 shape 是变长的，Q 和 KV 长度也不相等。TokenSpeed 的 MLA Kernel 支持无 padding 的变长批处理，默认实现基于 QDLS，性能表现良好，同时也提供预编译 binary 通过 TMA 接入。

跨平台 Kernel 设计

推理硬件越来越多样，跨平台能力成了推理引擎绕不开的问题。NVIDIA、AMD、华为、平头哥等厂商的加速器在硬件架构和编程模型上各不相同，如果 Kernel 代码直接绑死在某套硬件接口上，每支持一种新硬件都得大规模重构。

TokenSpeed 设计 Kernel 层时参考了 Linux 内核的抽象思路，建立 register（注册）与 selector（选择）机制。

模型适配代码完全不感知底层跑的是什么硬件，它只调用公共的 AI API 或 prefix API。这些 API 经过仔细定义，确保后续不轻易变动，其描述类型涵盖张量格式、执行约束等必要参数。厂商实现各自的 Kernel 时，通过 register 机制把内核注册进系统，声明支持的操作类型、capability、signature、优先级以及特定的 Tensor Core 特性。运行时由 select kernel 根据平台能力和数据格式动态过滤，挑出最合适、最稳定的 Kernel 完成前向计算。

这套设计给不同阶段的开发者提供了分层接入路径：新模型接入时，先用 reference 代码（例如 PyTorch 实现）把模型整体跑通，建立正确性基线；再把 reference 代码替换为 specialized 定制 Kernel，逐步提性能。优先级分层机制允许 reference 实现、performance 实现和 specialized plugin 实现同时存在，系统按注册时的优先级自动选择。TokenSpeed 的 Kernel 包有独立的 AMD Build，AMD 团队正是基于这套系统独立完成了 AMD GPU 的适配。

AMD 平台上的验证：多数场景优于 TensorRT

在 AMD MI355X 机器上，开发者利用 CDNA4 架构提供的 DSL 能力，可以精细控制流水线编排和内存布局。注意力部分采用 persistent prefill kernel（常驻式预填充内核，减少重复启动与调度开销），实现跨 CDA 的均衡调度。

测试结果如下：在 15 种典型的 prefill 形状中，有 14 种快于 TensorRT，领先幅度为 1.4 到 2.3 倍；在 MOE 的小 batch decode 场景中，速度比 TensorRT 快 1.7 到 2.1 倍。端到端测试在所有测量点上都取得了性能优势。

工程实践与性能收益

经过几个月迭代，TokenSpeed 在 Blackwell 平台上取得了多项可量化的成果。发布版本以 TensorRT-LLM 作为对比基准；TensorRT-LLM 是 NVIDIA 官方推理引擎，在 Blackwell 上本身已具备相当高的优化水平。

测试环境为单机 Blackwell、非 PD 分离架构（即预填充与解码不拆开部署），并使用 EvalScope 这一键式 benchmark 工具生成流量。结果显示，在“延迟约束下最大化吞吐”这一核心维度上，TokenSpeed 在所有测试场景中都不弱于 TensorRT-LLM：batch size 等于 1 的最小延迟场景快 9%；在每用户 100 TPS 的工作点上，单 GPU 吞吐提升 11%。在千问 3.5 397B 模型上，B200 八卡、batch size 等于 1 的配置跑到 580 TPS，在当时属于比较突出的水平。

MLA Kernel 的优化效果同样明显。在五种典型的 coding agent prefill 负载上，TokenSpeed 全部优于 TensorRT-LLM；在投机解码（先由小模型快速草拟 token、再由大模型校验）的 batch size 等于 1 场景下，延迟相对 TensorRT-LLM 几乎减半。

跨模型支持方面，千问 3.5 的上下文从 128K 扩展到 1M 时，解码吞吐仅下降 16%，扩展性较好。在 Kimi 2.5 上，每用户 TPS 达到 423。Kimi K3 当时处于 Day-0 支持状态，相关优化仍在持续迭代。

四条设计原则

回顾整个项目，可以提炼出几条核心设计原则：

  1. 优化目标要换。 在 Agent 场景下，推理引擎不应再单纯追求吞吐，而应转向“延迟约束下的吞吐最大化”，并以 Pareto 前沿（不同延迟档位下可达到的最优吞吐集合）作为评价引擎性能的核心框架。
  2. 分层要清晰。 控制面只做调度决策，不操作数据流；Runtime 只负责执行；Kernel 只负责计算。
  3. 并行切分要解耦。 把并行切分逻辑从模型适配链路中抽离出来，通过 Placement 系统和自动并行编译，实现声明式的模型部署。
  4. Kernel 要吃透新硬件。 既要充分发掘新硬件特性，也要通过抽象层为多硬件生态提供友好的接入路径。

TokenSpeed 开源数月，仍是一个年轻项目，正在快速迭代，也期待社区的反馈与贡献。

关键要点

图片

查看原文