TokenSpeed — 代码解读报告 #

Date: 2026-05-13

URL: https://github.com/lightseekorg/tokenspeed

Version: v0.1.0 (preview)

Domain: framework | Language: Python 82%, C++ 14%, CUDA/CuTe DSL 4% | LOC: ~205K

License: MIT | Org: LightSeek Foundation

Blog: https://lightseek.org/blog/lightseek-tokenspeed.html

---

What It Does #

一句话总结: 面向 agentic workloads 的 speed-of-light LLM 推理引擎，通过编译期 SPMD 并行化、FSM 类型安全调度器、分层插件化内核系统和 Blackwell MLA 内核，达到 TensorRT-LLM 级性能和 vLLM 级可用性。

类比: 如果 vLLM 是一台配置灵活的日本车（开箱即用、社区活跃），TensorRT-LLM 是一台高性能但难以维护的赛车（极致性能、深度定制），TokenSpeed 试图成为一台「带自动挡的赛车」——底层用 C++ FSM + Blackwell 内核取得极致性能，上层用 Python + Placement 注解保持模型开发效率。调度器的类型安全 FSM 设计则像 Rust 的所有权系统——在编译期而非运行时防止 KV cache 误用。

目标用户: 需要为 coding agent (Cursor/Claude Code/Codex) 等 agentic workloads 部署高性能推理的基础设施团队。

生态关系:

• 竞争: TensorRT-LLM (性能标杆)、vLLM (可用性标杆)、SGLang
• 上游血统: FluentLLM (meituan-longcat/SGLang-FluentLLM) 提供了部分 runtime 基础
• 合作: NVIDIA DevTech, AMD Triton, Qwen Inference, Together AI, Mooncake, NVIDIA Dynamo
• 被采用: TokenSpeed MLA 已被 vLLM 采纳 (PR #41778)

---

Architecture Overview #

┌──────────────────────────────────────────────────────────────────────┐
│                          User / Agent                                 │
│   ts serve <model>  →  SMG HTTP Gateway  →  gRPC Engine              │
└──────────────────────────────┬───────────────────────────────────────┘
                               │
                    ┌──────────▼──────────┐
                    │     AsyncLLM        │  Python (engine/)
                    │  request intake +   │  ZMQ ↔ scheduler IPC
                    │  output dispatch    │
                    └──────────┬──────────┘
                               │
            ┌──────────────────┼──────────────────┐
            │                  │                  │
   ┌────────▼────────┐ ┌──────▼──────┐ ┌─────────▼─────────┐
   │   C++ Scheduler  │ │  Execution  │ │  Model Loader     │
   │  (tokenspeed-    │ │  Plane      │ │  (HF/ModelScope)  │
   │   scheduler)     │ │  (Python)   │ └───────────────────┘
   │                  │ │             │
   │  FSM 13 states   │ │ CUDA Graphs │
   │  RAII KV safety  │ │ Spec Decode │
   │  Radix prefix    │ │ DP Control  │
   └────────┬─────────┘ └──────┬──────┘
            │                  │
            └────────┬─────────┘
                     │
        ┌────────────▼────────────┐
        │   Compiled Model Layer   │  models/base/
        │  Placement annotations   │  compiler.py → CommOps
        │  Static SPMD collectives │  execution.py → StepRunner
        └────────────┬────────────┘
                     │
        ┌────────────▼────────────┐
        │   Kernel Registry        │  tokenspeed-kernel/
        │  select_kernel(family,   │  5-band priority
        │    mode, dtype, platform)│  plugin override
        └────────────┬────────────┘
                     │
    ┌────────────────┼────────────────┐
    │                │                │
┌───▼───┐    ┌──────▼──────┐   ┌─────▼─────┐
│ MLA   │    │ Attention   │   │ GEMM/MoE  │
│(CuTe  │    │ (FA3/FA4/   │   │(DeepGEMM/ │
│ DSL   │    │  FlashInfer/│   │ TRT-LLM/  │
│Blackw.)│    │  Triton)    │   │ DeepEP)   │
└────────┘    └─────────────┘   └───────────┘

---

Core Concepts & Data Flow #

1. 请求生命周期 — FSM 驱动 #

TokenSpeed 用 C++ std::variant<13 states> 建模请求生命周期，所有资源通过 RAII 绑定到状态：

Submitted → Prefilling → PrefillDone → Decoding → Draining → WritingBack → Finished
                                          ↓ (memory pressure)
                                     Retracting → Retracted → (re-admit) → Decoding

关键设计: 状态转移通过 move 语义消费前一状态的资源（unique_ptr、unique_ptr），使得 KV cache 双重释放或遗漏释放在编译期就是错误。这是 C++ 的零成本抽象——运行时没有引用计数或锁。

2. 调度循环 #

Scheduler::NextExecutionPlan()
  → collect candidates (skip Draining/Prefetching/WritingBack/Retracting)
  → sort by priority: Prefilling(0) < Submitted(1) < Decoding(2) < Retracted(3)
  → greedy batch: fill max_scheduled_tokens + max_batch_size
  → prefill-first rule: if any prefill scheduled, skip ALL decode
  → pressure valve: if nothing scheduled but decode candidates exist → retract longest
  → return ExecutionPlan(FlatForwardOps, CacheOps)

Python execution plane:
  → run forward pass (CUDA graphs / eager)
  → Scheduler::Advance(ExecutionEvent)
  → repeat

3. Placement 编译器 — 静态 SPMD #

模型作者只写 Placement 注解，编译器自动插入通信操作：

# 模型作者写的:
class DeepSeekV4DecoderLayer:
    def attn_spec(self):
        return ModuleSpec(kind=ATTN, input=Replicate(ATTN_TP), output=Partial(ATTN_TP))
    def mlp_spec(self):
        return ModuleSpec(kind=MOE, input=Replicate(MOE_TP_EP), output=Partial(MOE_TP_EP))

# 编译器自动生成的:
# hidden(Partial@ATTN_TP) → AllReduce → hidden(Replicate@ATTN_TP)
# hidden(Partial@MOE_TP_EP) → ReduceScatter → hidden(Shard@MOE_TP_EP)

支持融合优化：FusedReduceNormOp（AllReduce + LayerNorm 合并）、ResidualAllGatherOp（残差连接与 AllGather 重叠）。

4. 内核选择流水线 #

select_kernel("attention", "decode", dtype=fp8, platform=SM100, features={"paged","mla"})
  → filter by CapabilityRequirement (arch, features, vendor)
  → filter by dtype, features, traits
  → sort by priority band
  → return SelectedKernel (cached for future calls)

5 个优先级带保证了可预测的选择行为：REFERENCE(0) < PORTABLE(4-7) < PERFORMANT(8-11) < SPECIALIZED(12-15) < PLUGIN(16-19)。外部插件始终可以覆盖内置内核。

---

Key Design Decisions #

1. C++ 控制面 + Python 执行面 #

选择: 调度器逻辑（FSM、KV 分配、prefix matching）用 C++ 实现，模型前向和内核调用用 Python。

好处: 调度决策在微秒级完成，不受 GIL 影响；FSM 类型安全在编译期保证。Python 面保持了内核/模型迭代效率。

代价: nanobind FFI 边界增加维护复杂度；C++ 测试需要独立的 GoogleTest 套件。

对比: vLLM 调度器全 Python；TensorRT-LLM 全 C++ 但可用性差。

2. Placement 注解 + 静态编译 vs DTensor / 手写通信 #

选择: 轻量级 Placement dataclass + compile_decoder_layer 静态分析，不用 torch.DTensor。

好处: 推理时零开销（DTensor 设计给训练，有大量 dispatch 开销）；模型作者不需要手写 AllReduce/AllGather。

代价: 灵活性不如动态 DTensor（不能在运行时改变并行策略）。

独特性: 三个独立的 ParallelGroup（ATTN_TP, DENSE_TP, MOE_TP_EP）支持不对称 TP——例如 Attention TP4 + MoE TP4EP 不同的并行度。

3. Retraction 作为 FSM 一等公民 #

选择: 内存压力下不是简单丢弃请求，而是 Decoding → Retracting → Retracted，KV 写回 host 后重新调度。

好处: 避免重新预填充已处理的长上下文（agent 场景中 context 常 >50K tokens）。

代价: FSM 状态数从 ~8 增加到 13；retraction 需要额外的 host 内存。

对比: vLLM 的 preemption 策略类似但在 Python 层实现，没有编译期安全保证。

4. MLA 的 q-head 折叠 #

选择: Decode 时将 q_seqlen 折叠到 head 维度（H_eff = num_heads × F），利用 BMM1 M tile。

好处: Agent 场景的 decode 通常 batch 很小（speculative decoding 下 q_seqlen=4-16），标准 MLA decode 的 Tensor Core 利用率很低。折叠后接近最优。

代价: 只适用于 q_seqlen 能被 F 整除的情况（runtime 选择最大合法 F）。

效果: 配合 speculative decoding 的 typical decode 工作负载，延迟几乎减半（相比 TRT-LLM MLA）。

5. 内核注册表的 5-band 优先级 #

选择: 固定 5 个优先级带而非连续数值。

好处: 外部插件作者不需要审计每一个内置注册就能确定该用什么优先级——只需选 PLUGIN 带就能保证覆盖。

代价: 同一带内只有 4 个相对位置（band+0 到 band+3）。

设计哲学: 类似 CSS 的 specificity 层级——提供「局部推理」能力，不需要全局知识就能预测选择结果。

---

核心三问 #

• Q1: 解决什么问题？ 现有推理引擎在 agentic workloads（长 context、高 turn 数、speculative decode）下的 performance gap——vLLM 可用性好但性能不够，TensorRT-LLM 性能好但可用性差，缺少为 agent 场景专门优化的引擎。

• Q2: 什么方法？ (1) C++ FSM + RAII 类型安全调度，保证 KV cache 安全的同时降低调度延迟；(2) Placement 静态编译器，让模型作者用注解而非手写通信；(3) 分层内核注册表 + Blackwell MLA 特化内核，在保持可扩展性的同时取得极致性能；(4) Retraction + Radix Prefix Sharing，为长 context agent 场景减少重复计算。

• Q3: 效果如何？ Kimi K2.5 on B200: Pareto 全面优于 TRT-LLM（min-latency 快 ~9%，100 TPS/User 吞吐高 ~11%）。MLA decode 在 spec decode 典型工作负载下延迟近乎减半。Preview 阶段，生产加固中。

---

逻辑故事还原 #

时代定位 #

2026年中，coding agents（Claude Code、Cursor、Codex）成为 AI 基础设施的主要工作负载。这类 workload 的特征——长上下文（>50K tokens）、高多轮、严格 TPS 要求——暴露了现有推理引擎的弱点：vLLM 的 Python 调度器在大 batch 下成为瓶颈，TRT-LLM 的全 C++ 设计让模型适配周期过长。市场需要一个为 agent 场景量身定制的引擎。

背景 #

LightSeek Foundation 联合 NVIDIA DevTech、AMD Triton、Qwen、Together AI 等团队，从 FluentLLM (SGLang fork) 出发，2026 年 3 月中旬开始重写调度器和内核层。核心目标：Blackwell 上达到 speed-of-light，同时保持 Python 模型开发体验。

约束推导 #

为何不直接改 vLLM？ 因为 vLLM 的调度器在 Python 中，KV cache 管理通过 runtime check 而非类型系统保证安全。在 1000+ 并发请求的 agent 场景下，Python 调度本身成为瓶颈，且 runtime check 难以覆盖所有并发角落情况。

为何不直接用 TRT-LLM？ 因为 TRT-LLM 的模型定义、权重加载、通信逻辑全在 C++/CUDA 中，添加新模型需要数周 C++ 开发。Agent 场景下新模型上线速度至关重要（DeepSeek V4、Qwen 3.5 等快速迭代）。

为何不用 DTensor 做并行？ 因为 DTensor 设计给训练场景，推理时的 dispatch 开销不可接受。但完全手写通信又太繁琐。折中是 Placement 注解 + 静态编译：开发期像 DTensor 一样简单，运行时零开销。

破局 #

TokenSpeed 的核心 insight 是：推理引擎的关键不是单一维度的极致，而是控制面安全性、执行面性能、开发面效率三者的平衡。C++ FSM 解决安全性（编译期防止 KV 误用），Blackwell MLA + 内核注册表解决性能（hardware-aware 选择最优内核），Placement 编译器解决效率（注解代替手写通信）。

核心技术壁垒 #

1. FSM + RAII KV 安全: std::variant<13 states> 配合 unique_ptr move 语义，是推理引擎中首次将「KV cache 资源安全」提升到编译期保证。这比 vLLM 的 Python BlockManager 在正确性上有本质优势。
1. MLA q-head 折叠: 将 q_seqlen 折叠到 head 维度来填充 Tensor Core tile，是对 agent 场景小 batch decode 的精准优化。这个 insight 需要深入理解 Blackwell UTCMMA 的 tile scheduling。
1. 不对称 TP 编译器: 三个独立的 ParallelGroup（ATTN_TP, DENSE_TP, MOE_TP_EP）允许每个组件使用不同的并行度。例如 Kimi K2.5 的最佳配置是 Attention TP4 + MoE TP4——编译器自动在组件边界插入 resharding。

设计绑定批判 #

• 绑定 Blackwell: MLA 内核目前仅支持 SM100/SM103；Hopper/MI350 优化标注为 "ongoing work"
• 绑定 FluentLLM 传承: Runtime 代码中 FluentLLM 的影子随处可见（engine 结构、entrypoint 设计），可能限制了更激进的架构重构
• 绑定 preview 状态: 77 commits、~3 月开发、PD disaggregation 未合并——这是概念验证而非生产系统，API 稳定性无保证

---

Key Findings #

• C++ std::variant-based FSM 将 KV cache 资源安全从 runtime check 提升到 compile-time guarantee——13 个状态中每一个都通过 RAII 绑定其拥有的页面、tree node、allocator，状态转移通过 move 语义消费资源
• Placement 编译器 + 三个独立 ParallelGroup 实现了不对称 TP，每个组件（attention、dense MLP、MoE）可以使用不同的并行度，编译器在边界自动插入 resharding
• MLA decode 的 q-head 折叠（fold_sq_factor）在 agent 典型 decode 工作负载（batch=4-16, long prefix）下延迟接近减半
• 内核注册表的 5-band 优先级设计提供了「局部推理」——插件作者只需知道 PLUGIN 带就能保证覆盖，不需审计所有内置注册
• Retraction 作为 FSM 一等公民（而非 crash/retry）避免了 agent 长 context 的昂贵重新预填充
• 调度策略: greedy batch + prefill-first + 最长请求 retraction 作为 pressure valve，L3 prefetch 定义了但未接入主循环（structured but not auto-emitted）

Limitations #

• Preview 状态: 3 个月开发、77 commits、多个核心 PR 未合并（PD、EPLB、KV store、Mamba cache、VLM、metrics）
• 模型覆盖有限: 目前主要支持 DeepSeek V3/V4、Qwen 3/3.5、Llama 系列和 Kimi K2.5；非 MoE/MLA 模型的优化深度不明
• Blackwell 中心: MLA 内核仅 SM100/SM103；Hopper 和 MI350 标注为 ongoing
• FluentLLM 传承债务: Engine/entrypoint 层保留了 FluentLLM 的结构，可能限制了更深层的架构优化
• L3 prefetch 未激活: 调度器定义了 host→device prefetch 逻辑，但 schedulePrefetch 未从 NextExecutionPlan 调用
• 缺少公开 benchmark 数字: 性能对比仅以 PNG 图表形式发布，无可复现的 CSV 数据

Infrastructure Impact #

• Kernel: MLA decode q-head 折叠是 Blackwell 上 agent workload 的重要优化方向，已被 vLLM 采纳；CuTe DSL 内核展示了 DSL-based kernel authoring 的可行性
• Framework: Placement 编译器的静态 SPMD 设计是 DTensor (训练) 在推理场景的轻量替代，可能影响其他推理引擎的并行化设计
• Framework: C++ FSM + RAII 类型安全调度为推理引擎的正确性设立了新标杆——从 runtime check 到 compile-time guarantee
• Algorithm: Retraction 作为 first-class FSM state 展示了如何在 memory pressure 下避免 agent 长 context 的重新计算
• Hardware: 不对称 TP（ATTN_TP ≠ MoE_TP）的编译器支持回应了 MoE 模型中 attention 和 expert 对并行度需求不同的现实

---

Deep Analysis (code) #

1. Project Identity #

• Name: TokenSpeed
• One-liner: Speed-of-light LLM inference engine for agentic workloads
• Domain: framework (inference serving)
• Owner: LightSeek Foundation
• Scale: ~205K LOC (530 Python, 153 C++/CUDA files)
• License: MIT
• GitHub: https://github.com/lightseekorg/tokenspeed
• Dev start: mid-March 2026; 77 commits as of 2026-05-12

2. Architecture & Module Map #

Package	Language	LOC	Responsibility
tokenspeed (python/)	Python	~100K	主 runtime: AsyncLLM、engine、models、layers、PD、spec decode
tokenspeed-kernel	Python + TVM-FFI	~45K	内核注册、选择、10 family × N solutions、thirdparty wrappers
tokenspeed-mla	Python + CuTe DSL	~8K	Blackwell MLA prefill/decode 内核
tokenspeed-scheduler	C++ + nanobind	~36K	FSM 调度器、KV allocators、radix tree、prefix cache

Key Design Decisions (expanded) #

1. radix tree + hybrid prefix cache 同时管理 KV 页面和 Mamba 状态——支持 DeepSeek V4 这样的 attention+Mamba 混合架构，无需两套独立的 cache 管理。
1. PagedCacheGroup 支持多组 paged cache——DeepSeek V4 的 sliding-window layers 和 full-history layers 使用不同的 cache group，scheduler 分别管理其页面生命周期。
1. 内核 thirdparty 隔离 所有第三方内核（FlashAttention、DeepGEMM、DeepEP、TRT-LLM）放在 thirdparty/ 下，通过 _triton.py 统一导入 Triton——防止 import side effect 污染，也方便 vendor 切换。
1. ts serve = SMG gateway + gRPC engine 一键启动 HTTP 网关和推理引擎，SMG 处理 structured generation、token streaming、OpenAI-compatible API。

3. Entry Points #

Entry: tokenspeed serve <model> (cli/__main__.py → serve_smg.py)
  Input: model path/name, --tensor-parallel-size, --attn-tp-size, --moe-tp-size, ...
  Output: OpenAI-compatible HTTP API (via SMG gateway)
  Side effects: spawns gRPC engine subprocess, loads model, starts serving

Entry: Engine (entrypoints/engine.py)
  Input: ServerArgs, PortArgs
  Output: AsyncLLM instance (programmatic API)
  Side effects: creates scheduler IPC (ZMQ), loads model

4. Critical Path: Request → Token #

[HTTP Request]
  → SMG Gateway → gRPC → AsyncLLM.add_request()
    → EngineCoreClient (ZMQ IPC) → scheduler subprocess
      → C++ Scheduler::SubmitRequests([RequestSpec])
      → Scheduler::NextExecutionPlan()                    ~µs (C++)
        → sort candidates by FSM state priority
        → greedy batch fill (tokens + batch size)
        → produce FlatForwardOperation
      → Python: ModelRunner.forward(execution_plan)       ~ms-s (GPU)
        → compiled decoder layer:
          → for each step: pre_comms → runner → post_comms
          → runner selects kernel via registry
        → CUDA graph capture / replay
      → Scheduler::Advance(ExecutionEvent)                ~µs
        → FSM state transitions (move semantics)
        → KV cache page allocation/release
      → repeat until Finished
    → OutputProcessor → Detokenizer → HTTP stream

瓶颈: GPU forward pass（prefill 可达数秒，decode ~ms）。C++ 调度器的 NextExecutionPlan 在微秒级完成，不是瓶颈。真正的优化空间在内核效率和 batch 策略。

5. Concurrency Model #

• 主进程: AsyncLLM 管理请求生命周期，asyncio event loop
• 调度器进程: 通过 ZMQ IPC 通信，C++ scheduler 单线程（无锁）
• GPU 执行: ModelRunner 在调度器进程中执行 forward pass
• Data Parallel: DataParallelController 管理多 DP 副本，跨 replica 分发请求
• 已知约束: Scheduler 单线程设计——在极高 QPS 下 ZMQ IPC 可能成为瓶颈（但 C++ 调度本身极快）

6. Supported Models #

Model	Architecture	Key Feature
DeepSeek V3	MoE + MLA	标准 MoE 路径
DeepSeek V4	MoE + MLA + sliding window + Mamba	多 cache group, hybrid prefix cache
Kimi K2.5	MoE + MLA	speed-of-light 优化目标
Qwen 3 / 3.5 / 3.5 MoE	Dense / MoE	Qwen 团队联合优化
Llama / EAGLE-3	GQA	spec decode drafter 支持
GPT-OSS	-	AMD optimized path
MiniMax M2	-
LongCat Flash	DiT	视频生成

7. Comparison with Alternatives #

Feature	TokenSpeed	vLLM	TensorRT-LLM	SGLang
调度器语言	C++ FSM	Python	C++	Python
KV 安全保证	编译期 (RAII)	运行时 check	运行时	运行时
模型定义语言	Python + Placement	Python	C++ / Python	Python
不对称 TP	✅ (3 groups)	❌	✅ (有限)	❌
Blackwell MLA	✅ CuTe DSL	✅ (采纳 TS MLA)	✅ native	❌
内核插件系统	✅ 5-band registry	❌	❌	❌
Retraction	✅ FSM first-class	✅ Python preemption	❌	✅
PD Disaggregation	✅ (Mooncake)	✅	✅	✅
Spec Decode	✅ EAGLE-3	✅	✅	✅
生产就绪度	Preview	Production	Production	Production
Agentic 优化	✅ 设计核心	⚠️ general	⚠️ general	⚠️ RadixAttention
开源	✅ MIT	✅ Apache	❌ (限制)	✅ Apache

8. Tech Debt & Code Quality #

Issue	Location	Severity	Impact
FluentLLM 传承代码	engine/, entrypoints/	Medium	架构中有非 TokenSpeed 原生的抽象层，增加理解成本
L3 prefetch 未接入	scheduler/operations/cache.cpp	Low	逻辑完整但从未被调用，死代码
内核 ops/ 部分不走 registry	ops/layernorm, kvcache, embedding	Medium	10 个 family 中只有 4 个 (attention, gemm, moe, quantize) 通过 register_kernel 注册
无 AOT binary 在 repo 中	tokenspeed-mla/objs/	Low	最优 MLA prefill 需要外部 .so，repo 内无法复现 blog 数字
测试覆盖不明确	test/	Medium	Python/C++ 测试存在但覆盖率未报告

9. Verdict & Recommendations #

应该关注的场景:

• 如果你在 Blackwell 上部署 MoE+MLA 模型（DeepSeek V3/V4, Kimi K2.5）的 agent 推理——TokenSpeed 展示了当前最佳的设计方向
• 如果你在设计推理引擎的调度器——FSM + RAII 模式值得借鉴
• 如果你需要为新模型快速适配并行——Placement 编译器是目前最优雅的方案之一

风险:

• Preview 状态，不适合生产部署
• 核心团队分布在多个组织（NVIDIA、AMD、Qwen、Together AI），长期维护承诺不明确
• MLA 内核的最优路径依赖未开源的 binary .so

Top 3 值得学习的设计:

1. C++ FSM + RAII 类型安全——将 KV cache 正确性从 convention 提升到 compilation，是推理引擎领域的一次方法论升级
2. Placement 编译器——DTensor 思想在推理场景的正确简化：保留注解的便利性，去掉 dispatch 的开销
3. MLA q-head 折叠——对 agent 场景 small-batch decode 的精准优化，体现了 workload-aware kernel design 的思路

10. Ecosystem Influence #

• MLA 内核被 vLLM 采纳 (PR #41778)，验证了 CuTe DSL kernel 的社区可移植性
• Placement 编译器 可能影响 vLLM/SGLang 未来的并行化设计方向（当前都是手写通信）
• C++ FSM 调度器 为推理引擎的 correctness-by-construction 设立了标杆
• 合作模式: LightSeek Foundation 联合 5+ 组织共建——如果成功，可能成为 foundation model inference 的 Apache Arrow 式协作项目