特色项目
摘要：TokenSpeed 推理引擎在 GPU 上运行 Qwen3.5-397B-A17B 模型时达到了创纪录的 580 tps。这种针对智能体（Agentic）负载的极致性能归功于系统性地消除内存拷贝、先进的算子融合以及完全重叠的 CPU-GPU 执行——始终保持 GPU 处于满载状态。在功能方面，TokenSpeed 还支持混合前缀缓存（Prefix Caching）和统一的 Prefill-Decode 状态传输，以应对复杂的智能体服务场景。
1. 介绍
Qwen 开源模型代表了一个能力强大的大语言模型家族，旨在实现广泛的易用性和灵活的部署。它们拥有涵盖不同参数规模的完整开源版本矩阵，满足从资源受限的边缘设备到复杂云环境的多样化需求。这些模型在海量高质量语料库上训练，在自然语言理解、高级逻辑推理、全栈编码和超长上下文处理方面展现出卓越的水平。此外，凭借对自主智能体规划、多步任务执行和工具调用的内置支持，Qwen 开源系列使全球开发者和研究人员能够高效地构建、定制和部署强大的 AI 应用。
Qwen3.5 模型作为 Qwen 开源系列的旗舰产品，通过采用混合注意力机制进一步突破了边界，该机制基于门控增量网络（Gated Delta Network, GDN），将标准全注意力层与线性注意力层交替组合。与传统的纯 Transformer 架构不同，这种混合设计在保持强大建模能力的同时，显著降低了长序列推理的计算复杂度。
TokenSpeed 是由 LightSeek 基金会根据 MIT 许可证发布的一款高性能开源 LLM 推理引擎，专为智能体负载打造。它的目标是在保持 vLLM 开发者友好性的同时，提供媲美 TensorRT-LLM 的“光速”性能。它采用原生的 SPMD 架构和静态编译从零构建，显著加速了复杂多步智能体任务的执行，使开发者能够高效部署超快、生产级的 AI 应用。
本文将介绍 TokenSpeed 推理框架中 Qwen3.5 模型的设计、实现和优化，涵盖运行时架构设计（PD 分离、前缀缓存、调度器）、关键性能优化及性能基准测试。
2. 运行时设计与特性
Qwen3.5 使用混合架构：大多数层为 GDN（具有层级 conv_state 和 temporal_state 的线性注意力），每隔 N 层插入一个具有常规 KV 缓存的标准全注意力层。TokenSpeed 提供了对 GDN 在前缀缓存、调度和 Prefill-Decode 分离方面的完整支持，实现了对整个混合栈的高效服务。
2.1 GDN/Mamba 前缀缓存
前缀缓存对于智能体负载至关重要，因为多轮工具调用序列频繁共享长上下文和对话历史。TokenSpeed 的前缀缓存分为两层。C++ 负责逻辑缓存：基数树匹配、页面 ID、驱逐和 Mamba 插槽生命周期。Python 负责物理张量：GPU KV 页面、Mamba conv_state / ssm_state、流排序、写时复制（CoW）、归零和快照拷贝。
对于常规 KV 缓存，前缀命中意味着重用缓存的页面 ID。对于 Mamba，这还不够。一个可重用的前缀必须携带同一前缀边界处的循环状态。TokenSpeed 通过将 MambaSlot 附加到表示缓存 KV 前缀的同一个基数树节点上来解决这个问题。
插槽生命周期
每个活动的 Mamba 请求可能持有两种插槽类型：
working插槽：当前前向传播步骤使用的可变状态。checkpoint插槽：后续可发布到前缀树的快照目标。
调度器在 C++ 中分配这些插槽，但由 Python 写入实际的张量内容。
只有在满足两个条件时，检查点插槽才变为可重用：Python 已将其填充为干净状态，且 C++ 已将其附加到块对齐的基数树节点上。
前缀匹配与写时复制
当后续请求匹配到树时，HybridPrefixCache 首先执行常规 KV 前缀匹配，然后找到最近的 Mamba 检查点节点。如果存在该节点，调度器返回 mamba_cow_src_index。
然后，Python 在运行前向传播之前将该缓存的检查点拷贝到请求的私有 working 插槽中。缓存树插槽不会被修改；只有请求的 working 插槽会发生改变。
保持检查点清洁
主要的正确性风险在于重用插槽中的陈旧数据：MambaChunkAllocator 分配整数插槽 ID 但不清除 GPU 内存。TokenSpeed 通过运行时规则防止状态陈旧。
具体而言，新分配的 working 插槽通过两种方式保证安全：要么接收来自已知干净检查点的写时复制拷贝，要么 Python 在使用前显式将其归零。检查点仅在对齐边界处发布，因此树从不将任意中间状态作为可重用的前缀状态进行广播。
重叠调度下的分块 Prefill
分块 Prefill 在重叠模式下引入了一个细微之处：CPU 可能在提交前一个块的输出之前调度下一个块。由于 CUDA 流排序，检查点仍然是安全的。
前一个块的 Mamba 前向传播在 execution_stream 上写入检查点。在下一个循环迭代开始时，默认流等待 execution_stream。只有在此之后，C++ 才将前一个块插入树中并分离其检查点插槽。下一个块随后会获得一个新的检查点插槽。
关键的不变量是：
C++ 可以在重叠调度期间发布检查点插槽 ID，但任何随后的 GPU 使用者都在前一个块的快照写入之后排序，且该发布的插槽不再被重用作为下一个检查点的目的地。
Decode 重叠
Decode 存在不同的风险：下一个 decode 可能在 CPU 提交前一个结果之前修改同一个 working 插槽。TokenSpeed 通过在分发下一个 decode 之前对块对齐的 decode 状态进行快照来处理这个问题。
这在 working 插槽推进之前保留了干净的状态。
总结
TokenSpeed 的 Mamba 前缀缓存是安全的，因为它将 Mamba 状态视为树拥有的检查点，而不是请求的附带效应。C++ 控制插槽何时成为前缀树的一部分。Python 控制张量内容何时被拷贝、归零和快照。它们共同维护一个核心不变量：从前缀树可达的每个 Mamba 插槽都包含该树节点所代表前缀的干净、对齐的状态。
2.2 调度器
混合架构对调度器提出了独特要求：它必须同时将 KV Cache（全注意力层）和 Mamba State（线性注意力层）作为两个独立的资源池进行管理。
Mamba 状态与混合模型管理
TokenSpeed 的调度器实现了以下关键机制：
- 双资源池管理：每个请求同时持有 KV Cache 块索引和 Mamba Pool 插槽索引（
mamba_pool_indices），调度器管理两者的分配和释放。 - 状态生命周期
- 请求到达时：分配 mamba_pool 插槽
- Prefill 期间：填充初始状态（或从前缀缓存加载）
- Decode 期间：每一步原地更新状态
- 完成或抢占时：释放插槽
- 推测解码支持：调度器维护中间状态缓存（
spec_cache），存储每个推测步骤的 Conv/SSM 状态快照，以便在验证失败时进行回滚。 - 层级路由：
HybridLinearAttnBackend根据layer_id将前向调用路由到适当的后端（全注意力或线性注意力），并为每种后端类型提供单独的元数据初始化。
2.3 GDN PD
2.3.1 挑战
对于混合模型，Mamba 层维护着超出常规键值对的状态张量。这些状态必须随 KV 缓存一起从 prefill 节点传输到 decode 节点，这需要在全注意力层和 Mamba 层之间进行正确的层级对齐。
2.3.2 我们的方案
我们为 PD 分离引入了端到端 Mamba 缓存支持，包括：
1. 统一状态传输：两个世界，一条链路
核心见解是，尽管 Mamba 状态具有不同的语义，但它们可以使用与 KV 缓存相同的 RDMA 机制进行传输——只要系统知道如何寻址它们。
我们在每个节点上设计了一个双张量池：一个池存储卷积状态（因果卷积的短期记忆），另一个存储循环 SSM 状态（长期压缩历史）。两者都作为连续的 GPU 内存预分配，每个请求每层拥有且仅拥有一个插槽。在注册时，prefill 和 decode 节点交换缓冲区描述符——基地址、每插槽大小以及从每个物理缓冲区到其相应全局层 ID 的映射。
当传输开始时，系统将每个请求的插槽索引映射为物理字节偏移量，将连续插槽分组为分散-聚集块，并作为批量 RDMA 写入发出。从网络的角度来看，Mamba 状态只是另一组内存区域——无需序列化，无需中间暂存。关键区别在于寻址：KV 缓存由页表索引，而 Mamba 状态由调度器分配的平面插槽 ID 索引。
2. 跨层调度：统一心跳
最微妙的部分是何时传输每一层的状态。
在层级传输模式下，prefill 节点不会等待整个前向传播完成才开始数据移动。相反，它在每个层组完成后立即开始发送数据——实现计算与通信的重叠。但对于混合模型，这意味着传输线程必须跨越注意力层和 Mamba 层跟踪进度，就像它们是一个连续的流水线一样。
我们引入了一个统一步骤计数器，在每一层的前向传播后递增一次——无论类型如何。传输线程监视该计数器，并为每个层窗口发送属于该窗口的任何数据：全注意力层的 KV 页面，Mamba 层的状态插槽。模型的层类型模式对传输逻辑变得不可见——它只是简单地询问“哪些缓冲区映射到层 4 到 7？”，并在计数器达到 7 后将它们全部发送。
在 decode 端，该机制的镜像是一个层完成屏障（layer-done barrier）：模型前向传播可以在第 15 层状态到达之前开始执行第 0 层。每一层的计算调用状态池，仅在该特定层尚未加载时才会阻塞。这允许 decode 将网络接收与早期层执行重叠，从而在有用的工作中隐藏传输延迟。
3. PD 感知 Token 生命周期：三阶段握手
最后一部分将状态传输与 Token 生成生命周期连接起来。在分布式系统中，prefill 节点不仅产生状态，还产生第一个输出 Token。Decode 节点在开始生成之前两者都需要。
我们设计了一个三阶段握手：
- 传输完成：最后一层组的所有 KV 页面和 Mamba 状态已发送。但传输线程尚未宣布成功——它在屏障处等待，等待前向传播结束。
- 生成 Token：Prefill 前向传播完成并发出第一个输出 Token。事件循环记录该 Token 并通知等待的传输线程。
- 状态交付：传输线程通过侧信道向 decode 端点发送一条轻量级状态消息（携带引导 Token）。只有当 decode 节点同时收到批量状态数据和此 Token 时，它才会向其调度器发出“远程 Prefill 完成”事件。
该协议确保了一个不变量：decode 节点从不会在状态不完整的情况下开始生成，也绝不会浪费一步重新推导第一个 Token。Mamba 状态、KV 缓存和引导 Token 作为逻辑原子单元到达——即使它们通过不同的路径并在不同的时间传输。
3. 性能优化
3.1 Mamba 状态更新优化
通过索引间接消除 Mamba 状态拷贝
在使用 Mamba 式线性注意力的推测解码中，目标验证阶段传统上带有隐藏的内存成本。在草稿模型生成推测 Token 后，基础模型运行前向传播进行验证。由于每个草稿 Token 将 Mamba 状态推进一步，引擎需要为每个推测位置保留中间状态，然后根据接受的 Token 数量恢复正确的状态。
之前的流水线使用专门的中间状态缓存来处理此问题：内核在验证期间将每步 Mamba 状态写入侧缓冲区，并在验证后使用 fused_mamba_state_scatter_with_mask 内核将接受位置的状态拷贝回调度器拥有的 working 插槽。Scatter 本身是一个跨 num_layers × state_dim 的全张量拷贝——代价并不低，且在每个解码步骤执行。
核心思想：移动指针，而非数据
我们不再在单独缓存中缓冲中间状态并在之后 Scatter，而是让内核将每一步的输出直接写入专用的物理行，然后简单地记住哪一行持有规范状态。
状态缓冲区扩展了一个在调度器分配的基础插槽后附加的草稿区域；每个请求拥有由其 req_pool_index 索引的草稿行私有切片。一个轻量级表 current_input_indices 记录了每个请求当前持有其规范 Mamba 状态的物理行。
目标验证期间：
- 输入重定向：内核从
current_input_indices记录的行中读取其初始状态（这可能是 working 插槽、CoW 分叉插槽或来自前一步的草稿行）。此处不进行数据移动，仅进行索引查找。 - 输出路由：一个每请求的
output_state_indices张量告诉内核确切写入每一步输出的位置：插槽 0 是 working 行，插槽 1..N 是请求私有的草稿行。内核直接写入这些预分配的位置，完全消除了中间缓存。 - 验证后记账：一旦获知接受长度，我们只需更新
current_input_indices [req]指向对应于最后一个接受 Token 的草稿行。这是一个 O(1) 整数写入，而不是 O(L·D) 张量拷贝。
3.2 运行时优化
3.2.1 重叠就是一切
遵循现代推理引擎的通用做法，TokenSpeed 使用 CUDA 多流并行来重叠非顺序操作。通过在多个流中并发执行独立负载，TokenSpeed 有效减少了调度开销并改善了端到端延迟。
共享专家与路由专家重叠
Qwen3.5 MoE 层包含共享专家和路由专家。共享专家处理所有 Token，而路由专家仅处理 TopK 选定的 Token。两者天然可并行，并通过 StreamFork 类进行流分叉和同步：
- 主流执行 TopK 路由、专家调度和 MoE GEMM
- 辅助流并发执行共享专家前向传播（gate_up → SiLU → down）和 Sigmoid 门控
- 两个流在组合结果前通过事件同步
这种重叠隐藏了共享专家计算延迟，降低了生产部署中单个 MoE 层的时间。
GDN 输入投影双流优化
GatedDeltaNet 层的输入投影包含两个独立的线性层（in_proj_qkvz 和 in_proj_ba），也在流中并行执行。
此优化仅在 CUDA Graph 捕获期间激活，较小的 in_proj_ba 投影被完全隐藏在备用流上较大的 in_proj_qkvz 背后。
3.2.2 融合越多，延迟越低
Gemma AllReduce 融合
GemmaRMSNorm 使用 x * (1 + weight) 代替标准 RMSNorm 的 x * weight，这此前阻碍了使用 TRT-LLM 的融合 AllReduce + Residual + RMSNorm 内核。
TokenSpeed 预计算 gemma_weight = weight + 1.0 并将其作为 gamma 参数传递给标准融合内核——以启用 GemmaRMSNorm 通信融合。融合后，每层的 AllReduce + 残差相加 + RMSNorm 从三个单独的内核启动合并为一个。
此融合覆盖所有 Qwen3.5 解码器层，并自动在 SM90+ 单节点 TP 部署上启用。
注意力中的融合 QK-RMSNorm + 部分 RoPE + 门控拆分
在原始注意力路径中，QKV GEMM 投影后，5 个单独的内核被顺序启动以归一化、旋转和拆分 Q/K/gate 向量：
| 步骤 | 操作 | 从 HBM 读取 | 写入 HBM |
|---|---|---|---|
| 1 | Q RMSNorm | q | q_normed |
| 2 | K RMSNorm | k | k_normed |
| 3 | Q RoPE | q_normed | q_rotated |
| 4 | K RoPE | k_normed | k_rotated |
| 5 | Gate 拆分 + 连续拷贝 | q_gate | gate |
每个中间张量（q_normed, k_normed 等）被写入全局内存只是为了被下一个内核立即读取——纯属带宽浪费。fused_qk_rmsnorm_rope_gate 用单个 Triton 内核替换了所有 5 次启动。所有中间值都保留在寄存器中。
MoE 共享专家中的融合 Gate-Sigmoid-Mul-Add
在 MoE 块中，共享专家输出在与路由专家输出合并前被门控。原始代码路径为此概念上的单一表达式启动了 5 个单独内核：
| 步骤 | 内核 | 说明 |
|---|---|---|
| 1 | 逐元素乘法 | h[i] * w[i] — 逐元素乘积 |
| 2 | 归约 | 部分乘积求和 → 每个 Token 1 个标量 |
| 3 | Sigmoid | σ(gate_val) — 对标量进行逐元素运算 |
| 4 | 乘法 | σ(gate_val) * shared_output — 广播标量 × 全向量 |
| 5 | 加法 | final_hidden_states += scaled |
主要低效之处：gate_val 是每个 Token 的标量，但未融合路径在启动之间将逐元素乘积和归约后的标量物化到 HBM。中间的 scaled 张量（完整的 [num_tokens, hidden_dim]）也被写入并立即重新读取。fused_gate_sigmoid_mul_add 在一个 Triton 内核中原地计算完整表达式 final += σ(x·w) * shared。点积归约、Sigmoid、广播乘法和累加都在每个 Token 的单个线程块内完成——中间值从不离开寄存器。
3.2.3 千次同步之死
TokenSpeed 的解码循环将核心前向传播（目标模型、采样器和草稿模型）捕获到单个 CUDA 图中。一旦捕获，数千个 GPU 内核通过一次启动重放，完全消除了逐个内核的调度开销。
但 CUDA 图在设计上是静态的。在图重放之间，运行时仍必须在主机上执行动态工作：准备输入、解析调度索引、在推测验证后更新 Mamba 状态指针以及协调传输状态。这些图之间的“间隙”正是 CPU 开销隐藏的地方——粗心的 .item() 或不必要的 D2H 拷贝可能会阻塞整个流水线。
TokenSpeed 将这种跨图 CPU 开销视为首要优化目标：即使在图之外，也要让主机远离关键路径。
消除设备到主机的往返
最阴险的同步模式是“无辜查询”——从 GPU 读取单个标量以在主机上做出分支决策。TokenSpeed 将这些替换为初始化时已知的预计算最坏情况界限，或在 H2D 传输前捕获 CPU 端最大值，以便 GPU 张量及其界限可以同时可用。对于推测解码状态管理，边界检测和插槽选择使用 GPU 端标记值——下游内核通过边界检查跳过无效条目，而不是进行 CPU 端过滤。整个决策树保持在设备上。
编译融合索引算术
混合模型中的运行时调度涉及繁重的索引操作：计算插槽映射、草稿 Token 布局和验证后的指针更新。在即时（Eager）PyTorch 中，每一步都成为一个单独的内核启动，中间值写入 HBM。TokenSpeed 使用 torch.compile 注释这些例程，允许 Inductor 将 10–14 次单独启动融合为一个或两个逐元素内核，所有值都流经寄存器。GPU 保持繁忙，CPU 提交一次启动而不是十四次。
一切皆异步
H2D 传输自始至终使用带有非阻塞拷贝的固定内存（Pinned memory）。传输系统轮询固定主机计数器而不是调用 synchronize()，层级加载使用基于事件的屏障，仅唤醒需要数据的特定层。CPU 在当前批次仍在传输时准备下一个批次。
累计效果：TokenSpeed 的解码循环保持接近零的 CPU 开销——主机线程将时间花在提交工作上，而不是等待结果。
3.3 FA4 支持
Flash Attention 4 (FA4) 是针对 NVIDIA Blackwell 架构的下一代注意力内核。Qwen3.5 默认使用 head_dim=256，对注意力计算后端提出了重大需求——并非所有内核都能开箱即用地高效支持此配置。
对 head_dim=256 的 FA4 支持已被贡献并合入上游社区仓库。在 TokenSpeed 中，对 Qwen3.5 的原生 FA4 支持目前处于活跃开发中，并将在即将发布的版本中提供，进一步释放 Blackwell GPU 用于 Qwen3.5 推理的全部计算潜力。
4. 基准测试
以 Qwen3.5-397B-A17B 为例，我们对 NVIDIA Blackwell GPU 上的 Qwen3.5 模型进行了系统性的性能评估。感谢 EvalScope 团队提供基准测试工具；下述所有性能结果均使用 EvalScope Benchmark 获得。
测试环境：所有基准测试均使用 TokenSpeed 最新 Docker 镜像（lightseekorg/tokenspeed-runner:latest），基于近期版本。基准测试脚本和复现说明可在 TokenSpeed 的 GitHub 仓库中找到。
4.1 基础基准测试
我们使用固定的输入/输出长度并测量解码吞吐量（输出 Token/秒）。我们在不同的并行配置（TP/EP）下评估了不同 Batch Size 的性能。使用了两种主要的测试配置：
- 配置 1：Attn TP + MoE TP
- 配置 2：Attn TP + MoE EP
我们在启用了 MTP 和禁用 MTP 的 B200 上对 Qwen3.5-397B-A17B-NVFP4 的解码吞吐量进行了基准测试。
在 Attn TP8 + MoE TP8 / Attn TP8 + MoE EP8 的所有输入/输出长度配置中，MTP 在 bs=1（延迟是主要瓶颈）时带来了 +100%~+159% 的吞吐量增益。在更高的并发度下，增益与输出长度强相关：长输出负载（例如，输出长度 >4096）在 bs=32/64 时保持了 +38%~+90% 的显著加速，而短输出负载（例如，bs=64 时的 1024 个 Token）增益减小到接近零或略有负值，因为当解码已受吞吐量限制时，推测开销开始超过接受利益。
4.2 智能体负载基准测试
智能体应用的快速激增（包含工具调用历史和多轮对话上下文）从根本上重塑了生产负载的特性。为了反映真实的智能体行为，我们使用模拟真实智能体调用模式（50K 第一轮上下文，后续每轮附加 800 个 Token，总共 10-15 轮）的智能体负载测试套件。
在 B200 使用 NVFP4 的情况下，TokenSpeed 为 Qwen3.5-397B-A17B 在智能体负载下提供了卓越的单用户吞吐量。所有四种并行配置——TP4、TP4EP4、TP8 和 TP8EP8——在 bs=1 时均维持 500+ tok/s，其中 TP8 达到了约 580 tok/s 的峰值。
在 concurrent=16 时，TP4 系列扩展到约 2K tok/min/GPU 系统吞吐量，而 TP8 系列达到约 1K tok/min/GPU。在相同的 GPU 数量下，纯 TP 和 TP+EP 配置表现出相当的吞吐量-延迟权衡，为用户提供了部署灵活性而不牺牲性能。值得注意的是，多轮智能体负载实现了超过 90% 的平均 KV 缓存命中率，显著降低了 Prefill 开销，并对整体吞吐量增益做出了贡献。
4.3 长上下文基准测试（高达 1M）
长上下文处理是智能体负载的另一个关键挑战。虽然前缀缓存可以在多轮对话中命中大量重复前缀并显著降低 Prefill 开销，但 Decode 阶段在每一步仍然必须读取并处理整个历史 KV，缓存无法绕过这一点——上下文越长，每一步的 Decode 内存访问成本就越高。
基于 NIAH（大海捞针）1M 样本，我们切分了四种提示长度——128K / 256K / 512K / 1M——进行评估。在 Qwen3.5-397B-A17B 上，解码吞吐量在 128K 内保持在约 530 tok/s/user，256K 时为约 495，1M 时为约 445（在 TP8 上测量），使得从 128K 到 1M 的端到端衰减仅为约 16%——长上下文吞吐量衰减得到了很好的控制。
5. 结论
通过上述优化和架构设计，TokenSpeed 为 Qwen3.5 模型提供了出色的性能——特别是在智能体负载中，实现了超低延迟生成和高推理吞吐量。TokenSpeed 将继续突破 Qwen 推理优化的边界，在栈的每一层追求更极致的性能。
我们诚邀您关注 TokenSpeed 项目，亲自体验光速推理吞吐量。GitHub 上提供了完整的安装指南，使得在支持的硬件上部署和基准测试变得简单明了。我们也热忱欢迎社区提交以性能为导向的 Pull Request——每一次贡献都有助于 Qwen 模型系列运行得更快、更智能。
6. 致谢
这项工作通过开源生态系统中的紧密合作得以实现。我们要感谢阿里云通义团队、NVIDIA DevTech、Mooncake 团队和 LightSeek 基金会的工程协作与实现支持。我们还要感谢 NVIDIA 和 Verda 提供 Blackwell GPU 基础设施和计算支持。