智能体推理引擎TokenSpeed：重新定义LLM推理效率，优化从内核到调度全栈

大型语言模型推理的性能瓶颈，早已不再局限于“矩阵乘法运算的速度够不够快”这个单一维度。

随着 Claude Code、Codex、Cursor 这类智能体编码系统从演示级产品蜕变为真正的生产力工具，推理系统所应对的流量特征也发生了根本性转变：上下文长度动辄超过 50K tokens，对话轮次跨越数十轮，工具调用、代码检索、补丁生成以及反思修正等操作不断交织在一起。

TokenSpeed: A Speed-of- Light LLM Inference Engine for Agentic Workloads
博客地址：https://lightseek.org/blog/lightseek-tokenspeed.html
代码仓库：https://github.com/lightseekorg/tokenspeed
全文约 8000 字，阅读时长约 35 分钟，播客版本约 22 分钟

LightSeek 官方博客指出，为了承载这种增长趋势，业界正在建设需要数十吉瓦电力供应的大型数据中心，并投入数千亿美元级别的资本。

在这种规模下，每块 GPU 的吞吐量哪怕只提升几个百分点，都可能转化为极为可观的算力节省。TokenSpeed 的雄心正在于此：它并非仅仅优化某个 kernel，而是从建模、调度、KV 资源管理、安全复用、内核系统到请求入口， 全面重构一套面向 agentic workloads 的 LLM 推理引擎。

这组对比图揭示了 TokenSpeed 在 MoE 大模型推理中的显著性能优势：横轴代表用户请求的 Token 速率（体现并发压力），纵轴表示 GPU 的处理效率（单位：千级 Token/分钟），覆盖了 Attn TP4/TP8 与 MoE TP/EP 四种并行配置。在所有场景下，TokenSpeed 的表现均优于 TensorRT-LLM，尤其在 MoE 专家并行（EP）与高并发场景中差距更为明显，这说明 TokenSpeed 针对 MoE 的专家路由与跨设备通信进行了深度优化，有效降低了调度开销。同时，Attn TP4 的整体效率显著高于 TP8，反映出过高的张量并行度会引入明显的通信瓶颈，从而抵消算力收益；而 TokenSpeed 在 TP8 配置下依然能维持性能优势，体现了其并行优化的鲁棒性。总体来看，TokenSpeed 实现了更优的“负载-效率”权衡，在高用户并发下能保持更高的 GPU 利用率，更适配智能体工作负载对低延迟和高吞吐的需求。这张图展示了 Kimi K2.5 MoE 大模型在 NVIDIA B200 GPU 上，针对智能体（Agentic）工作负载的推理性能对比。横轴为用户侧单请求 Token 速率（Token/Sec/User，代表服务延迟与并发压力），纵轴为 GPU 侧 Token 处理效率（Token/Min/GPU，代表硬件利用率），核心对比了自研推理系统（TS）与 NVIDIA TensorRT-LLM（TRTLLM）在张量并行（TP）、专家并行（EP）、数据并行（DP）等不同策略下的表现。从趋势上看，所有配置都呈现出典型的“负载-效率”权衡：用户请求速率越高（服务延迟越低），GPU 因调度与通信开销导致的处理效率就越低。关键结论包括：其一，自研 TS 框架整体显著优于 TRTLLM，特别是 TP4 配置在低负载下 GPU 效率可达 1500k Token/Min/GPU，远超 TRTLLM 同配置的 1100k，这体现了针对 MoE 模型与智能体负载的深度优化优势；其二，MoE 专家并行（EP）始终优于 MoE 张量并行（TP），所有配置中 EP 曲线均在 TP 曲线上方，说明 EP 更适配 MoE 的专家路由特性，减少了跨设备通信开销；其三，并行度并非越高越好，TP8 效率低于 TP4，DP8 表现最差，这说明过高并行度的通信延迟抵消了算力收益，在智能体负载下低并行度加 EP 的组合更具优势。更多信息见：https://github.com/lightseekorg/tokenspeed/tree/main/test/agentic_benchmark

这两张图表对比了 TensorRT-LLM MLA 与 TokenSpeed MLA（开源 OSS/二进制版）在大模型 MLA（Multi-Head Latent Attention）机制下的 Prefill 与 Decode 阶段延迟性能。Prefill 阶段（上图）覆盖了 5 种业务场景，TokenSpeed 二进制版（Binary）的延迟始终最优，在 use case 3 中较 TensorRT-LLM 降低了约 20%，展现了定制化算子优化的优势；TokenSpeed 开源版（OSS）在部分场景中延迟略高于 TensorRT-LLM，这可能意味着其开源版本的优化尚未完全释放。Decode 阶段（下图）聚焦于不同批大小（4/8/16），TokenSpeed 的延迟优势被进一步放大：当批大小为 16 时，TensorRT-LLM 的延迟达到 338.8us，而 TokenSpeed 仅为 146.3us，降幅超过 56%，这说明其针对 MLA 解码的核心瓶颈（如注意力计算、归约操作）进行了深度优化，尤其适配高并发场景下的低延迟需求。整体来看，TokenSpeed 的 MLA 实现，特别是二进制优化版，在 Prefill 与 Decode 阶段都展现出了显著的性能优势，为大模型推理提供了更高效的底层实现方案。

unsetunset本文目录unsetunset

一、快速上手与最小使用路径
二、为什么智能体负载需要新的推理引擎
三、架构总览：Python 执行平面与 C++控制平面的分工
四、入口层：轻量 CLI 背后的异步请求主干
五、请求主路径：从 tokenize 到调度器再回到用户
六、C++状态机调度：TokenSpeed 最核心的系统设计
七、KV 缓存：从张量缓存变成资源所有权系统
八、MLA KV 布局：为长上下文与 FP8 缓存做结构化折中
九、Attention 后端：按模型架构和硬件自动落位
十、Kernel 系统：把异构硬件优化变成可注册、可解释、可覆盖的机制
十一、MLA 性能优化：TokenSpeed 押注的 Blackwell 低延迟战场
十二、性能预览：以 TPS/User 和 TPM/GPU 构造真实 Pareto 边界
十三、PD 分离与服务拓扑：下一阶段的系统战场
十四、协作生态与开源位置：站在 TensorRT-LLM 肩膀上继续推进
十五、当前边界：预览版意味着什么
结语：TokenSpeed 的本质是把推理性能问题系统化

unsetunset一、快速上手与最小使用路径unsetunset

快速上手：预览版部署与体验

TokenSpeed 目前仍处于预览阶段。其 README 文档明确指出，此版本的核心目标在于复现 Kimi K2.5 on B200 以及 TokenSpeed MLA[1] on B200 的性能数据。项目正处于快速迭代期，官方不建议直接将其用于生产环境。如需快速体验，推荐通过 runner 容器来搭建开发环境。

首先，拉取官方镜像并启动容器：

docker pull lightseekorg/tokenspeed-runner:latest  

docker run -itd   
  --shm-size 32g   
  --gpus all   
  -v /raid/cache:/home/runner/.cache   
  --ipc=host   
  --network=host   
  --pid=host   
  --privileged   
  --name tokenspeed   
lightseekorg/tokenspeed-runner:latest   
/bin/bash

进入容器后，克隆项目代码并安装依赖：

git clone https://github.com/lightseekorg/tokenspeed.git  
cd tokenspeed  

export PIP_BREAK_SYSTEM_PACKAGES=1  
pip install -e "./python" --no-build-isolation  
pip install -e tokenspeed-kernel/python/ --no-build-isolation  
pip install -e tokenspeed-scheduler/

环境配置完成后，可以通过以下命令查看帮助或启动服务：

tokenspeed env  
tokenspeed serve --help

一个最简的服务启动命令如下：

tokenspeed serve openai/gpt-oss-20b   
  --host 0.0.0.0   
  --port 8000   
  --tensor-parallel-size 1

若想启动更接近官方性能预览的 Kimi K2.5 配置，可参考以下参数骨架：

tokenspeed serve nvidia/Kimi-K2.5-NVFP4   
  --served-model-name kimi-k2.5   
  --host 0.0.0.0   
  --port 8000   
  --trust-remote-code   
  --max-model-len 262144   
  --kv-cache-dtype fp8   
  --quantization nvfp4   
  --tensor-parallel-size 4   
  --enable-expert-parallel   
  --chunked-prefill-size 8192   
  --max-num-seqs 256   
  --attention-backend trtllm_mla   
  --moe-backend flashinfer_trtllm   
  --reasoning-parser kimi_k2   
  --tool-call-parser kimi_k2

服务启动后，兼容 OpenAI API 的客户端可以这样访问：

from openai import OpenAI  

client = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1")  
response = client.chat.completions.create(  
model="kimi-k2.5",  
messages=[{"role": "user", "content": "Write a concise deployment checklist."}],  
max_tokens=256,  
)  
print(response.choices[0].message.content)

更多参数细节可参考：Getting Started[2]、Launching a Server[3]、Server Parameters[4]。

二、为何智能体负载需要全新的推理引擎？

传统的 LLM serving 框架大多围绕“批量吞吐”与“通用聊天”场景进行优化。然而，coding agent 的流量模式更像是一场复杂的城市交通调度：某些请求携带超长上下文，某些请求仅需极短的 decode 步骤，某些请求重复使用相同的系统提示词和仓库上下文，还有的请求在工具调用后立即进入下一轮交互。官方博客将这种场景定义为 agentic-inference regime。

TokenSpeed 的设计初衷，就是从第一性原理出发，专门服务于此类负载。其核心特性包括：

基于编译器支撑的 modeling mechanism，实现并行 placement 与 collective 的自动生成；
高性能的 C++ scheduler；
由类型系统约束的安全 KV 资源复用机制；
可插拔、分层的 kernel 系统，支持异构加速器；
集成 SMG[5] 的低开销 CPU 请求入口。

这表明，TokenSpeed 并非“又一个 vLLM 替代品”，也不是“TensorRT-LLM 的轻量封装”。它试图回答一个更底层的问题：当智能体应用将 token 生产规模推向数据中心级别时，推理系统如何同时保障 GPU 吞吐、用户侧 TPS、长上下文复用以及工程的可迭代性？

三、架构总览：Python 执行平面与 C++ 控制平面的分工

TokenSpeed 的核心设计理念可以概括为：Python 负责易用性与快速迭代，C++ 负责调度控制与资源正确性，kernel 子系统则专注于硬件相关的性能优化。

README 中列出的核心组件包括：

Modeling layer：采用 local-SPMD 设计，通过模块边界的 I/O placement annotation，让轻量级静态编译器自动生成 collective communication。
Scheduler：由 C++ 控制平面与 Python 执行平面共同构成。请求的生命周期、KV cache 所有权以及 overlap timing 均由有限状态机来表达。
Kernels：一个分层、可插拔、集中注册的 kernel 系统。
Entrypoint：集成了 SMG 的 AsyncLLM，旨在降低 CPU 侧的请求处理开销。

其中，scheduler 的分层设计尤为关键：

控制平面使用 C++ 编写，并与类型系统紧密结合，力求在编译期就约束好 KV cache 的状态转移和资源使用。
执行平面则采用 Python 编写，方便研究人员和工程师快速加入新模型、新特性或进行新实验。

这种分工类似于“飞机驾驶舱与发动机舱”：控制逻辑必须稳定可靠、易于验证，而执行逻辑则需具备高度的灵活性，以便快速调整。

四、入口层：轻量 CLI 背后的异步请求主干

TokenSpeed 向外界提供了三种接入方式：tokenspeed serve 命令行、兼容 OpenAI 的 HTTP 服务器，以及 Python 层面的 Engine.generate 接口。CLI 部分设计得极为精简，主要职责是命令的分发调度，以及将重量级模块的导入操作推迟到真正需要时才执行。

# 来源：python/tokenspeed/cli.py  
def main() -> None:  
parser = argparse.ArgumentParser(  
prog="tokenspeed",  
description="TokenSpeed is a speed-of-light LLM inference engine.",  
)  

subparsers = parser.add_subparsers(dest="command")  

serve_parser = subparsers.add_parser(  
"serve",  
help="Launch the TokenSpeed inference server.",  
)  
if len(sys.argv) >= 2 and sys.argv[1] == "serve":  
from tokenspeed.runtime.utils.server_args import ServerArgs  

ServerArgs.add_cli_args(serve_parser)  
serve_parser.set_defaults(func=_serve)  

args, extra_args = parser.parse_known_args()  

if args.command is None:  
parser.print_help()  
sys.exit(1)  

args.func(args)

真正的核心入口实际上是 Engine 类。通过代码注释可以清晰地看到，整个系统由 TokenizerManager、Scheduler 子进程以及输出处理链路三大部分构成，各进程之间依靠 ZMQ 库进行通信。

# 来源：python/tokenspeed/runtime/entrypoints/engine.py  
class Engine(EngineBase):  
"""  
The entry point to the inference engine.  

    - The engine consists of three components:  
        1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.  
        2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.  
        3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.  

Note:  
    1. The HTTP server, Engine, and TokenizerManager both run in the main process.  
    2. Inter-process communication is done through ICP (each process uses a different port) via the ZMQ library.  
"""

在 Engine.generate 方法中，用户的输入首先被封装为一个 GenerateReqInput 对象，随后根据是否启用流式返回，分别进入两条不同的处理路径。

# 来源：python/tokenspeed/runtime/entrypoints/engine.py  
obj = GenerateReqInput(  
text=prompt,  
input_ids=input_ids,  
sampling_params=sampling_params,  
image_data=image_data,  
audio_data=audio_data,  
video_data=video_data,  
return_logprob=return_logprob,  
logprob_start_len=logprob_start_len,  
top_logprobs_num=top_logprobs_num,  
token_ids_logprob=token_ids_logprob,  
custom_logit_processor=custom_logit_processor,  
return_hidden_states=return_hidden_states,  
stream=stream,  
bootstrap_host=bootstrap_host,  
bootstrap_port=bootstrap_port,  
bootstrap_room=bootstrap_room,  
)  
if stream:  
return self.llm.generate_stream(obj)  
else:  
return self.llm.generate(obj)

同步 API 并未额外实现一套独立的同步 IPC 客户端，而是借助 `LLM` 类桥接到后端的 asyncio 事件循环：

```python
# 来源：python/tokenspeed/runtime/engine/llm.py  
def run(self, coro) -> Any:  
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()  

def generate(self, obj: GenerateReqInput) -> dict:  
async def _one() -> dict:  
gen = self.async_llm.generate_request(obj)  
return await gen.__anext__()  

return self.run(_one())  

def generate_stream(self, obj: GenerateReqInput) -> Iterator[dict]:  
q: queue.Queue[Any] = queue.Queue()  

async def _drain() -> None:  
pending_exc: BaseException | None = None  
try:  
async for item in self.async_llm.generate_request(obj):  
q.put(item)  
except BaseException as exc:  
pending_exc = exc  
finally:  
if pending_exc is not None:  
q.put(pending_exc)  
q.put(_STREAM_END)  

asyncio.run_coroutine_threadsafe(_drain(), self._loop)  

while True:  
item = q.get()  
if item is _STREAM_END:  
return  
if isinstance(item, BaseException):  
raise item  
yield item

这种设计的意义在于：业务端能够使用同步接口，而内部依然维持 AsyncLLM 这一单一主干，避免了因同步/异步双路径而引发的状态分裂问题。

五、请求主路径：从 tokenize 到调度器再回到用户

AsyncLLM 作为主进程的异步前端，承担着请求接入、请求状态管理、scheduler IPC 通信以及输出分发循环等职责。生成路径的核心集中在 generate_request：

# 来源：python/tokenspeed/runtime/engine/async_llm.py  
async def generate_request(  
self,  
obj: GenerateReqInput | EmbeddingReqInput,  
):  
created_time = time.time()  

self.auto_create_handle_loop()  
self.input_processor.validate_request(obj)  
obj.normalize_batch_and_arguments()  

async with self.model_update_lock.reader_lock:  
is_single = obj.is_single  
if is_single:  
tokenized_obj = await self._tokenize_one_request(obj)  
self._send_one_request(obj, tokenized_obj, created_time)  
async for response in self._wait_one_response(obj):  
yield response  
else:  
async for response in self._handle_batch_request(obj, created_time):  
yield response  

async def _tokenize_one_request(  
self,  
obj: GenerateReqInput | EmbeddingReqInput,  
) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput:  
return await self.input_processor.tokenize_one_request(obj)

请注意这里的 model_update_lock.reader_lock。它表明 TokenSpeed 将在线权重更新与请求生成纳入了同一套并发控制机制，从而确保推理过程中模型权重不会被不安全地切换。

向 scheduler 发送请求的逻辑如下：

# 来源：python/tokenspeed/runtime/engine/async_llm.py  
def _send_one_request(  
self,  
obj: GenerateReqInput | EmbeddingReqInput,  
tokenized_obj: TokenizedGenerateReqInput | TokenizedEmbeddingReqInput,  
created_time: float | None = None,  
):  
state = ReqState(  
RequestOutputCollector(),  
False,  
asyncio.Event(),  
obj,  
created_time=created_time,  
tokenized_time=tokenized_obj.created_time,  
)  
self.rid_to_state[obj.rid] = state  
self.engine_core_client.send_to_scheduler.send_pyobj(tokenized_obj)

### 整体链路概览

整个处理流程可归纳为：

HTTP/OpenAI/Python API  
→ Engine.generate / async_generate  
→ AsyncLLM.generate_request  
→ InputProcessor.tokenize_one_request  
→ EngineCoreClient.send_to_scheduler  
→ C++ Scheduler 生成 ExecutionPlan  
→ Python 执行 forward/cache op  
→ 输出回到 AsyncLLM.handle_loop  
→ OutputProcessor 唤醒对应 ReqState  
→ stream 或 final response 返回用户  

## 六、C++ 状态机调度：TokenSpeed 最核心的系统设计

TokenSpeed 官方博客特别指出：scheduler 将控制平面与执行平面解耦。控制平面由 C++ 有限状态机实现，并尽可能利用类型系统来确保 KV cache 状态转移、资源复用以及生命周期管理的正确性。

在代码中，请求状态并非简单的 waiting/running/finished，而是被拆分为一组精细的状态：

// 来源：tokenspeed-scheduler/csrc/fsm/states.h  
// Put resources into each particular state  
namespace tokenspeed::fsm {  

using State = std::variant<Bootstrapping, Submitted, Prefetching, PrefetchDone, Aborting, Prefilling, PrefillDone,  
Decoding, Draining, WritingBack, Retracting, Retracted, Finished>;  

}  // namespace tokenspeed::fsm  

这里的 `std::variant` 非常关键。它意味着资源不再挂载于请求对象上的若干可空字段，而是与状态紧密绑定。

- 例如，处于 `Decoding` 状态的请求拥有 decode 所需的 KV allocator 和 req pool index；
- 处于 `WritingBack` 状态的请求则持有 device/host node ref，用于确保异步写回期间对应页不会被提前释放。

调度器每轮会生成一个 `ExecutionPlan`：

// 来源：tokenspeed-scheduler/csrc/scheduler/execution_plan.h  
class ExecutionPlan {  
public:  
template <typename OperationType>  
ExecutionPlan& With(OperationType operation) {  
operations_.emplace_back(operation);  
return *this;  
}  

template <typename OperationType>  
ExecutionPlan& With(std::vector<OperationType> ops) {  
for (auto& op : ops) {  
operations_.emplace_back(std::move(op));  
}  
return *this;  
}  

const std::vector<Operation>& Operations() const { return operations_; }  

private:  
std::vector<Operation> operations_;  
};  

`NextExecutionPlan` 负责清理已完成请求、跳过正在 cache 操作中的请求，并从可调度候选中生成 forward/cache operation：

// 来源：tokenspeed-scheduler/csrc/scheduler/scheduler.cpp  
ExecutionPlan Scheduler::NextExecutionPlan() {  
ExecutionPlan plan;  

std::vector<WriteBackOperation> write_back_ops;  
write_back_ops = std::move(newWriteBackOperation(requests_));  

std::erase_if(requests_, [](const auto& req) { return req.second->template Is<fsm::Finished>(); });  

std::vector<Request*> candidates;  
for (auto& [id, req] : requests_) {  
if (!req->Is<fsm::Draining>() && !req->Is<fsm::Prefetching>() && !req->Is<fsm::Retracting>() &&  
!req->Is<fsm::WritingBack>()) {  
candidates.push_back(req.get());  
}  
}  

auto [fwd_ops, cache_ops] = newForwardOperation(candidates);  
plan.With(FlatForwardOperation{std::move(fwd_ops)});  

if (!write_back_ops.empty()) {  
plan.With(CacheOperation{FlatWriteBackOperation{write_back_ops}});  
}  
if (auto* lb = std::get_if<std::vector<LoadBackOperation>>(&cache_ops)) {  
if (!lb->empty()) {  
plan.With(CacheOperation{FlatLoadBackOperation{*lb}});  
}  
}  
return plan;  
}

这就是 TokenSpeed 与众多以 Python 为核心的调度器之间的关键差异：它并非依赖规约与运行时检查来保障资源正确性，而是**将“请求当前所处的阶段、所持有的资源、以及能够响应哪些事件”直接内化为控制系统的固有属性。**

## 七、KV 缓存：从张量缓存进化为资源所有权体系

> 在包含长上下文的智能体应用场景里，KV 缓存具有极高的价值。**系统提示词、代码仓库上下文、工具说明以及历史执行轨迹往往会被反复使用。因此，TokenSpeed 将 KV 缓存视为一种按页管理的资源，并引入了 device、host、L3 storage 等多个层级。**

在调度器一侧，滚动哈希（rolling hash）被用于实现页面级别的前缀匹配：

// 来源：tokenspeed-scheduler/csrc/scheduler/scheduler.cpp  
std::vector<std::string> Scheduler::CalcRollingHash(const std::vector<std::int32_t>& input_tokens, bool apply_match) {  
        conststd::int32_t page_size = config_.page_size;  
        conststd::size_t num_pages = input_tokens.size() / page_size;  
        std::vector<std::span<conststd::int32_t>> token_pages;  
        token_pages.reserve(num_pages);  
        for (std::size_t i = 0; i < num_pages; ++i) {  
            token_pages.emplace_back(input_tokens.data() + i * page_size, page_size);  
        }  
        if (!apply_match) {  
            return ComputePagedHashes(token_pages, "");  
        }  
        MatchResult result = kv_prefix_cache_.Match(token_pages);  
        conststd::int32_t host_matched = result.host.DepthInPage();  
        if (host_matched >= static_cast<std::int32_t>(num_pages)) {  
            return {};  
        }  
        constauto& hashes = result.host.last_node->PageHashes();  
        std::string prior = hashes.empty() ? std::string{} : hashes.back();  

        return ComputePagedHashes(  
            std::vector<std::span<conststd::int32_t>>(token_pages.begin() + host_matched, token_pages.end()), prior);  
}

请求处理完毕后，KV cache 并不会立即被释放，而是有可能进入 host 写回流程。`FinishEvent` 会将已完成的 device KV 页插入到 prefix cache 中；若需要持久化到 host 端，则会转入 `Draining` 状态。

```cpp
// 来源：tokenspeed-scheduler/csrc/fsm/forward_events.cpp
template <typename ForwardStateT>
std::variant<Draining, Finished> FinishEvent::apply(ForwardStateT&& state) {
auto full_paged_tokens = state.GetFullPagedTokens(true);
std::vector<std::int32_t> prefix_pages = DevicePagesFromRoot(state.GetDeviceNode());
std::int32_t alloc_count =
static_cast<std::int32_t>(full_paged_tokens.size()) - static_cast<std::int32_t>(prefix_pages.size());

auto local_allocator = std::move(state).TakeLocalKVAllocator();
OwnedPages alloc_pages = local_allocator->TakeFirst(alloc_count);

kv_prefix_cache_->Insert<ResourceType::Device>(full_paged_tokens, prefix_pages, std::move(alloc_pages),
page_hashes_);

MatchResult match = kv_prefix_cache_->Match(full_paged_tokens);
if (!disable_l2_cache_ && (match.device.DepthInPage() > match.host.DepthInPage())) {
std::vector<TreeNode*> write_diff = match.NodesWithout<ResourceType::Host>();
std::int32_t host_pages_num = 0;
for (TreeNode* node : write_diff) {
host_pages_num += node->Device().NumPages();
}
std::unique_ptr<HostNodeRef> temp_lock = std::make_unique<HostNodeRef>(match.host.last_node);
if (!kv_prefix_cache_->EnsureCapacityByEvict<ResourceType::Host>(host_pages_num)) {
return Finished{};
}
kv_prefix_cache_->AllocateResourceOfType<ResourceType::Host>(write_diff);

return Draining{BuildWriteBackPairs(write_diff), std::move(device_node_ref), std::move(host_node_ref)};
}
return Finished{};
}

这段代码揭示了 TokenSpeed 为何要将生命周期拆解得如此细致：从业务层面看，请求已经终结；但从资源层面看，KV 页可能仍处于异步传输过程中。若缺少 Draining/WritingBack 这类状态，cache 复用极易引发竞态条件。

八、MLA KV 布局：面向长上下文与 FP8 缓存的结构化权衡

TokenSpeed 的 MLA 实现是官方博客重点阐述的核心优化之一。该项目独立提供 tokenspeed-mla 软件包，专门针对 NVIDIA Blackwell SM100/SM103 架构，支持 MLA 的 prefill 与 decode 阶段，以及 MLA K/V 打包 + FP8 量化功能。

在运行时，MLA KV Pool 会根据所选的量化策略来分配不同的内存布局：

# 来源：python/tokenspeed/runtime/layers/attention/kv_cache/mla.py  
with memory_saver_adapter.region():  
# The padded page 0 is used for writing dummy outputs from padded tokens.  
if self.quant_method == "per_token_head":  
self.kv_buffer = [  
(  
torch.zeros(  
(self.size + self.page_size, 1, kv_lora_rank),  
dtype=self.store_dtype,  
device=device,  
),  
torch.zeros(  
(self.size + self.page_size, 1, 1),  
dtype=torch.float32,  
device=device,  
),  
torch.zeros(  
(self.size + self.page_size, 1, qk_rope_head_dim),  
dtype=self.model_dtype,  
device=device,  
),  
)  
for _ in range(layer_num)  
]  
else:  
self.kv_buffer = [  
torch.zeros(  
(self.size + self.page_size, 1, self.kv_cache_dim),  
dtype=self.store_dtype,  
device=device,  
)  
for _ in range(layer_num)  
]

当启用 per_token_head 量化后，KV 缓存会被拆解为三个独立组件：

k_lora_cache：采用 FP8 格式存储，有效降低显存占用；
k_scale_cache：以 FP32 格式保存缩放因子，确保量化恢复时的精度不受损；
k_rope_cache：使用模型原始数据类型，保障 RoPE 相关部分的数值稳定性。

在向 MLA KV buffer 写入数据时，TokenSpeed 会根据不同的量化方式选择对应的处理路径：

# 来源：python/tokenspeed/runtime/layers/attention/kv_cache/mla.py  
def set_mla_kv_buffer(  
self,  
layer: PagedAttention,  
loc: torch.Tensor,  
cache_k_nope: torch.Tensor,  
cache_k_rope: torch.Tensor,  
):  
layer_id = layer.layer_id  
if self.quant_method == "per_token_head":  
k_lora = cache_k_nope.float()  
k_rope = cache_k_rope.float()  
scale = k_lora.abs().amax(dim=-1, keepdim=True).clamp(1e-26) / 448.0  
k_lora = (k_lora / scale).to(torch.float8_e4m3fn)  
k_rope = (k_rope / scale).to(self.model_dtype)  
self.kv_buffer[layer_id][0][loc] = k_lora.view(self.store_dtype)  
self.kv_buffer[layer_id][1][loc] = scale  
self.kv_buffer[layer_id][2][loc] = k_rope  
else:  
if cache_k_nope.dtype != self.dtype:  
cache_k_nope = cache_k_nope.to(self.dtype)  
cache_k_rope = cache_k_rope.to(self.dtype)  

set_mla_kv_buffer_triton(  
self.kv_buffer[layer_id],  
loc,  
cache_k_nope,  
cache_k_rope,  
enable_pdl=pdl_enabled(),  
)

由此可见，TokenSpeed 中的 KV cache 绝非一个简单的张量，而是一种针对模型架构、量化策略、Kernel 访问方式以及跨层调度进行协同设计的内存布局。

九、Attention 后端：依据模型架构与硬件自动选择最优落位

TokenSpeed 的注意力注册机制会根据模型架构与硬件平台，自动选取默认的 backend 实现。以 MLA 架构为例：在 Blackwell 平台上默认采用 trtllm_mla，而在 Hopper 上则优先选择 flashmla。

# 来源：python/tokenspeed/runtime/layers/attention/registry.py  
def _get_default_backend_name(arch: AttentionArch) -> str:  
platform = current_platform()  
if arch == AttentionArch.MLA:  
if platform.is_blackwell:  
return "trtllm_mla"  
if platform.is_hopper:  
return "flashmla"  
return "trtllm_mla"  
else:  
return "mha"  

def _get_backend_cls(name: str, arch: AttentionArch) -> type[AttentionBackend]:  
if name is None:  
candidates = [_get_default_backend_name(arch)]  
for candidate in candidates:  
entry = _BACKEND_REGISTRY.get(candidate)  
if entry is not None and arch in entry[0]:  
return entry[1]  
raise ValueError(  
f"No backend supports arch {arch}. Available: {list(_BACKEND_REGISTRY)}"  
)

在构建注意力组件时，系统会首先通过 profile 机制评估可用的 KV 缓存页数，随后再创建对应的 backend 与内存池：

# 来源：python/tokenspeed/runtime/layers/attention/registry.py  
max_total_num_pages = profile_max_num_pages(  
config,  
gpu_id,  
server_args.mapping.world_size,  
server_args.gpu_memory_utilization,  
server_args.block_size,  
num_layers,  
gpu_memory,  
world_group=server_args.mapping.world_group,  
draft_attn_config=draft_attn_config if draft_attn_config else None,  
draft_num_attention_layers=(  
draft_model_config.num_attention_layers if draft_attn_config else None  
),  
cache_cell_size=profile_cache_cell_size,  
)  
max_num_tokens = _resolve_max_num_tokens(  
max_total_num_pages,  
server_args.block_size,  
server_args.max_total_tokens,  
)

这一设计揭示了 TokenSpeed 的核心取舍逻辑：并非要求用户手动理解所有底层细节，而是将硬件特性、模型结构、显存预算以及 backend 选择等复杂因素，统一集中在 runtime 的组件构建阶段进行处理。

十、Kernel 系统：将异构硬件优化转化为可注册、可解释、可覆盖的机制

官方博客着重指出，TokenSpeed 的 kernel 层已从核心引擎中解耦，成为独立的模块化子系统。

tokenspeed-kernel 的 README 文档清晰地展示了其分层架构：public API → select_kernel → KernelRegistry → 各 op family 下包含的 Triton、Gluon、CuTe DSL、vendor wrapper 以及 reference 实现。

kernel 的注册结构如下：

# 来源：tokenspeed-kernel/python/tokenspeed_kernel/registry.py  
@dataclass(frozen=True)  
class KernelSpec:  
"""Complete specification of a registered kernel."""  

# Identity  
name: str  
family: str  
mode: str  
features: frozenset[str] = frozenset()  
solution: str = ""  

# Capabilities  
capability: CapabilityRequirement = field(default_factory=CapabilityRequirement)  
dtypes: frozenset[torch.dtype] = frozenset()  
traits: dict[str, frozenset[Any]] = field(default_factory=dict)  

# Selection metadata  
priority: int = int(Priority.PERFORMANT) + 2  
tags: frozenset[str] = frozenset()

选择逻辑会综合考量配置文件中的 override 设置、上下文中的 override 参数、环境变量、平台能力、数据类型（dtype）、特征（traits）以及优化目标（objective）：

```python
# 来源：tokenspeed-kernel/python/tokenspeed_kernel/selection.py  
def select_kernel(  
family: str,  
mode: str,  
dtype: torch.dtype,  
    *,  
features: frozenset[str] | None = None,  
platform: PlatformInfo | None = None,  
objective: SelectionObjective = SelectionObjective.DEFAULT,  
traits: dict[str, Any] | None = None,  
override: str | None = None,  
expected_kernel_name: str | None = None,  
) -> SelectedKernel:  
platform = platform or current_platform()  

config_entry = _get_config_override(family, mode)  
if config_entry is not None:  
if override is None:  
if config_entry.name:  
override = config_entry.name  
elif config_entry.solution:  
override = config_entry.solution  

env_key = f"TOKENSPEED_KERNEL_OVERRIDE_{family.upper()}_{mode.upper()}"  
env_override = os.environ.get(env_key)  
if env_override:  
override = env_override  
registry = KernelRegistry.get()  

cache_key = _make_cache_key(  
family, mode, dtype, platform.arch, objective, features, traits  
)  
if override is None:  
cached = registry.cache_get(cache_key)  
if cached is not None:  
return cached  

if override:  
return _resolve_override(registry, family, mode, dtype, override, platform)  

candidates = registry.get_for_operator(  
family, mode, features=features, platform=platform, dtype=dtype  
)  

if traits:  
candidates = _filter_by_traits(candidates, traits)  

scored = _rank_by_objective(candidates, objective, platform, traits)  
winner = scored[0][0]  

impl = registry.get_impl(winner.name)  
result = SelectedKernel(name=winner.name, impl=impl)  
registry.cache_put(cache_key, result)  
return result

这套系统的价值在于：当运行环境从 NVIDIA Blackwell 切换到 Hopper、AMD，或者 shape/dtype 发生变化时，kernel 选择不再依赖散落在各处的 if-else 分支。通过 registry、traits、priority、objective 和 override 机制，可以实现统一的管理和调度。

十一、MLA 性能优化：TokenSpeed 在 Blackwell 低延迟战场上的关键布局

官方博客将 TokenSpeed MLA 列为核心优化方向之一。其目标并非简单地“支持 attention”，而是针对 coding agent 这类负载进行极致调优：

prefill 面向长 prefix KV cache；
decode 面向 batch size 为 4、8、16 等 speculative decoding 常见场景；
q_seqlen 通常较短；
num_heads 在某些场景下较小，导致 Tensor Core 利用率不足。

TokenSpeed MLA decode 的关键优化策略是：将 query-sequence 轴折叠进 head 轴，使 q_seqlen 和 num_heads 共同填满 BMM1 的 M tile。README 中也明确指出，当 num_heads < 128 且 num_heads * q_seqlen <= 128 时，会启用 fold_sq 机制，以改善小 head 数下的 M 维度利用率。

官方博客给出的性能预览显示，在典型 decode 工作负载上，这种优化结合其他调优手段，可相对 TensorRT-LLM MLA 实现接近一半的延迟降低；而 binary prefill kernel 则利用 NVIDIA internal knobs 对 softmax 实现进行了精细调整，在五类 coding-agent prefill 场景中超越了 TensorRT-LLM MLA。需要注意的是，这些数据来自官方性能预览，且当前仓库仍处于 preview 状态。

MLA Decode 性能基准测试示例

执行 MLA 解码（Decode）阶段的基准测试，可以使用以下命令：

python ./tokenspeed-mla/python/tokenspeed_mla/mla_decode_fp8.py   
  --batch_size 4   
  --softmax_scale 0.07216882   
  --page_size 64   
  --seq_len_k 81920   
  --in_dtype Float8E4M3FN   
  --out_dtype Float8E4M3FN   
  --seq_len_q 4   
  --warmup_iterations 1   
  --iterations 10   
  --num_heads 16   
  --skip_ref_check

Prefill 性能基准测试示例

执行 Prefill 阶段的基准测试，可以参考以下命令：

python ./tokenspeed-mla/python/tokenspeed_mla/fmha.py   
  --is_causal   
  --bottom_right_align   
  --in_dtype Float8E4M3FN   
  --out_dtype Float8E4M3FN   
  --q_shape 1,8192,128,192   
  --k_shape 1,8192,128,192   
  --warmup_iterations 10   
  --iterations 10   
  --skip_ref_check

值得注意的是，官方博客披露，TokenSpeed MLA 已被 vLLM 项目采纳[6]。这表明该优化方案并非 TokenSpeed 的独家技术，它有望成为更广泛推理生态中的底层能力。

十二、性能预览：以 TPS/User 和 TPM/GPU 构建真实的 Pareto 边界

TokenSpeed 官方博客并未止步于提供单一吞吐量数值，而是采用了更贴近编码智能体（coding agent）场景的指标来评估系统性能。下图展示了 Kimi K2.5 MoE 大模型在 B200 GPU 上，针对智能体工作负载的推理性能对比。

横轴（X 轴）：TPS/User，即每位用户每秒生成的 Token 数，作为衡量延迟体验的指标。
纵轴（Y 轴）：TPM/GPU，即每个 GPU 每分钟生成的 Token 数，作为衡量吞吐效率的指标。
曲线生成：通过扫描并发度（sweeping concurrency）绘制出 Pareto 曲线。

该图表清晰地展示了所有配置下典型的“负载-效率”权衡关系：用户请求速率越高，GPU 因调度和通信开销导致的处理效率就越低。自研推理框架 TS 的整体表现显著优于 TRTLLM。在低负载场景下，其 TP4 配置的 GPU 效率可达 1500k Token/Min/GPU，远超 TRTLLM 同配置的 1100k，这充分体现了针对 MoE 模型与智能体负载的深度优化优势。同时，专家并行（EP）普遍优于张量并行（TP），过高的并行度会因通信延迟抵消算力收益，这表明低并行度结合 EP 的方案更适配智能体工作负载的性能需求。

这种评估方式比单纯追求“最大吞吐量”更适合智能体推理场景。因为编码智能体并非离线批处理系统，用户体验高度依赖于生成速度。官方设定的目标是：在保证每用户 TPS 不低于某个下限的前提下，最大化每 GPU 的 TPM。常见的下限为 70 TPS，在某些场景下要求达到 200 TPS 甚至更高。

官方博客指出，在 Kimi K2.5、NVIDIA B200 以及非 PD 分离（non-PD disaggregation）的部署配置中，对于编码智能体常见的 70 TPS/User 以上的区间，最佳配置是 Attention TP4 + MoE TP4。在整个 Pareto 前沿（Pareto frontier）上，TokenSpeed 的表现均优于 TensorRT-LLM：在 batch size 为 1 的最低延迟场景下，性能提升约 9%；在 100 TPS/User 附近，吞吐量提升约 11%。

这些数字的重要性不在于“9% 或 11% 的提升听起来是否足够大”，而在于数据中心规模效应带来的巨大影响。正如官方博客开篇所述，当 Token 生成规模已经大到需要数十吉瓦级数据中心来支撑时，每 GPU 效率的微小提升，都会被整个生产集群（production fleet）放大，产生显著的效益。

十三、PD 分离与服务拓扑：下一阶段的系统战场

官方博客明确指出，当前的性能预览聚焦于单机、非分离（non-disaggregated）部署场景。PD 分离（PD disaggregation）功能正在清理和优化中，后续会单独发布文章进行详细介绍。不过，在仓库的代码中已经能看到相关雏形，例如 MiniLoadBalancer 负责在 Prefill 服务器和 Decode 服务器之间进行配对选择：

# 来源：python/tokenspeed/runtime/pd/mini_lb.py  
def select_pair_round_robin(self):  
    assert len(self.prefill_configs) > 0, "No prefill servers available"  
    assert len(self.decode_servers) > 0, "No decode servers available"  
    prefill_config = self.prefill_configs[self.prefill_index]  
    self.prefill_index = (self.prefill_index + 1) % len(self.prefill_configs)  
    decode_server = self.decode_servers[self.decode_index]  
    self.decode_index = (self.decode_index + 1) % len(self.decode_servers)  
    return prefill_config.url, prefill_config.bootstrap_port, decode_server

Prefill 和 Decode 阶段的资源需求完全不同：Prefill 阶段处理长上下文的大块计算，而 Decode 阶段则是严格延迟敏感的小步迭代。PD 分离的潜在价值在于，可以让不同的硬件池、不同的并发策略以及不同的缓存传输机制来服务于不同的阶段，而不是将所有请求都塞进同一个调度模型。

十四、协作生态与开源位置：站在 TensorRT-LLM 肩膀上继续推进

官方博客的致谢部分至关重要，因为它揭示了 TokenSpeed 在生态中的位置。TokenSpeed 由 LightSeek Foundation 牵头，并与 NVIDIA DevTech、AMD Triton、Qwen Inference、Together AI、Mooncake、LongCat、FluentLLM、EvalScope[7]、NVIDIA Dynamo 等多个团队协作开发。

官方博客明确感谢了 TensorRT-LLM 的维护者，称 TensorRT-LLM 设定了他们衡量自身性能的标杆，并且 TokenSpeed 的许多优化也受到了 TensorRT-LLM 的启发，包括 one-CUDA-graph optimization 和 forward pass optimizations。同时，它也向 Triton、FluentLLM、vLLM、EvalScope、FlashInfer、SGLang 等开源推理社区表达了感谢。

这段背景让 TokenSpeed 的价值更加清晰：它并非凭空构建的封闭系统，而是在现有推理生态高度成熟的基础上，针对 agentic workloads 进行全新的系统级权衡。一方面对标 TensorRT-LLM 在 Blackwell 上的极致性能，另一方面也努力保留 vLLM 那样的易用性与开源可迭代性。

十五、当前边界：预览版意味着什么

TokenSpeed 目前仍处于 preview release 阶段。README 和官方博客均强调：该项目始于 2026 年 3 月中旬，engine 和 kernels 仍在积极开发中，生产环境的 hardening 计划将在后续推进，未来数周内还会有大量 PR 合并。

README 中列出的 ongoing work 包括：

模型覆盖：Qwen 3.6、DeepSeek V4、MiniMax M2.7；
runtime 特性：PD、EPLB、KV store、Mamba cache、VLM、metrics；
平台优化：Hopper、MI350 及相关 runtime 改进。

因此，TokenSpeed 当前更适合被视为一条技术路线和性能工程样板，而非“可立即替换线上推理集群”的成熟产品。其最值得研究之处，在于它如何将智能体推理中的典型痛点系统化： 长上下文、短 decode、高并发、KV 复用、低 CPU overhead、异构 kernel 选择，以及控制平面的安全性。

结语：TokenSpeed 的本质是把推理性能问题系统化

如果只看 MLA kernel，TokenSpeed 像一个追逐 Blackwell 峰值的底层优化项目；如果只看 tokenspeed serve，它又像是一个 OpenAI 兼容的推理服务。但将代码、README 和官方博客串联起来看，它真正重要的地方在于：将 LLM 推理从“单点算子优化”推进到了“系统级协同优化”。

Python 前端负责易用接口与异步请求治理；
C++ scheduler 通过状态机将生命周期和资源所有权显式化；
KV cache 从临时张量演变为由 page、prefix、host/device、writeback、retract 共同参与的资源系统；
kernel registry 将硬件差异、dtype、shape traits 和性能目标纳入可解释的选择过程；
MLA 内核优化则针对 coding agent 的长 prefix、短 decode、小 head 利用率问题，提供底层加速。

智能体时代的推理流量，将不再是简单的“给一个 prompt，生成一段答案”。它会是成千上万条长上下文、短步长、多轮交互、频繁复用的请求流。谁能像调度高速公路上的车辆一样管理这些请求，像仓储物流一样复用 KV cache，像切换不同路况下的发动机模式一样选择 kernel，谁就更接近所谓的 speed-of-light inference。

TokenSpeed 目前仍很年轻，且明确处于预览阶段。但它给出的技术方向足够清晰：未来 LLM 推理系统的竞争，不止是单个 kernel 谁更快，而是谁能在建模、调度、缓存、内核、入口和服务拓扑之间，建立一套可验证、可扩展、可持续优化的系统。

参考资料[1] TokenSpeed MLA: https://github.com/lightseekorg/tokenspeed/tree/main/tokenspeed-mla

[2] Getting Started: https://lightseek.org/tokenspeed/guides/getting-started

[3] Launching a Server: https://lightseek.org/tokenspeed/guides/launching

[4] Server Parameters: https://lightseek.org/tokenspeed/configuration/server

[5] SMG: https://pytorch.org/blog/lightseek-smg/

[6] TokenSpeed MLA 已被 vLLM 采用: https://github.com/vllm-project/vllm/pull/41778

[7] EvalScope: https://github.com/modelscope/evalscope

关注“鲸栖”小程序，掌握最新AI资讯

本文来自网络搜集，不代表鲸林向海立场，如有侵权，联系删除。转载请注明出处：https://www.itsolotime.com/archives/33572