概要
TokenSpeed はオープンソースの LLM 推論エンジンであり、プロダクションレベルのエージェントワークロードを対象として、TensorRT-LLM クラスの性能と vLLM クラスの使いやすさを兼ね備えている。そのアーキテクチャは、C++ の有限状態マシンによる制御プレーンと Python の実行プレーンを分離し、コンパイル時に KV キャッシュとリクエストライフサイクルの安全性を保証する。カーネルシステムはプラガブル設計を採用しており、Blackwell GPU 上で最速クラスの MLA 実装の一つを含んでいる。
README プレビュー
\n \n\n\nTokenSpeed is a speed-of-light LLM inference engine designed for **agentic workloads**, with TensorRT-LLM-level performance and vLLM-level usability. The goal is to be the most performant inference engine for production agentic workloads.\n\nCore components:\n\n- **Scheduler**: C++ control plane and Python execution plane. Request\n lifecycle, KV cache ownership, and overlap timing are encoded as a\n finite-state machine, with safe KV resource reuse enforced by the type system at compile time.\n- **Kernels**: pluggable, layered kernel system with a portable public API and\n a centralized registry including one of the fastest **MLA**\n (Multi-head Latent Attention) implementations on Blackwell for agentic workloads.\n- **Entrypoint**: SMG-integrated AsyncLLM for low-overhead CPU-side request\n handling.\n\n## Differentiation\n\nTokenSpeed takes a fundamentally different architectural approach from existing LLM inference engines. The **PyTorch Foundation highlighted these architectural differentiators** when [welcoming TokenSpeed](https://pytorch.org/blog/pytorch-ecosystem-landscape-q3-update/) to the PyTorch Ecosystem Landscape:\n\n> TokenSpeed is an open source LLM inference engine and the first to separate the control plane from the execution plane. The control plane is implemented in C++ as a finite-state machine, using the type system to enforce safe resource management, including request lifecycles and KV cache state, at compile time rather than runtime. The execution plane is implemented in Python, enabling fast iteration and lowering the cognitive load for researchers and engineers. This architecture combines strong correctness guarantees in the core scheduling system with the development velocity of a high-level execution layer.\n>\n> TokenSpeed also treats kernels as a first-class, modular subsystem, separating them from the core engine through a portable public API, centralized registry and selection model, and an extensible plugin mechanism
FAQ (2)
トラブル対応triton sampling バックエンドを使用している場合、n>1 を指定した単一のリクエストが同一の completions を返すのはなぜですか?
既知のバグ: parallel_sampling.py の prepare_parallel_sampling_replica は、トークン化されたリクエストをシャローコピーし、rid だけを差し替えるため、n 個すべてのレプリカが同じ SamplingParams と seed を共有してしまう。triton サンプラーは tl.randint(seed, offset)(ops/sampling/triton/gumbel.py)からノイズを導出するため、すべてのレプリカが各ステップで同じトークンを生成する。リクエストは正しい使用量の計上とともに HTTP 200 を返すため、サイレントな失敗となる。回避策: (1) n: N の1つのリクエストの代わりに、n: 1 で N 個の個別のリクエストを送信する。リクエスト単位のサンプリングはそれぞれ異なる補完を生成する(temperature 0.7 および 1.2 で確認済み)が、prefill の共有とバッチ処理が失われるという代償を伴う。(2) NVIDIA GPU では、デフォルトの --sampling-backend flashinfer を使用する。その Philox RNG にはバッチ行が含まれる(curand_init(philox_seed, bx, ...))ため、レプリカは自然に分岐する。また、紛らわしい類似の原因も除外すること: バックエンドは NVIDIA 以外のプラットフォームではデフォルトで greedy となり、temperature/top_p/top_k/seed を無視しながらも 200 を返す。アクティブなバックエンドは curl http://<host>:<control-port>/get_server_info | grep sampling_backend で確認すること。これは GRPO 形式の RL ロールアウトにおいて重要であり、n>1 の同一の補完は分散ゼロのアドバンテージを生じさせ、学習信号が得られなくなる。
トラブル対応AMD gfx950 で高い並行性でサービングを行うと、一部のベンチマークステップがランダムに約1秒間ストールし、約36%遅くなるのはなぜですか?
ストールは、起動時ではなくサービング中にJITカーネルコンパイルが発生することが原因です。原因は2つあります。(1) スケジューラはプロンプトを最大8192トークンのprefillバッチに集約しますが、デフォルトのprefillグラフキャプチャは2048トークンまでしかカバーしていないため、大きなprefillバッチはキャプチャされていないパスを通り、サービング中にカーネルがコンパイルされます。(2) Attentionカーネルには多数のコンパイル時特殊化があり、それらはwarmup中に実行されません。修正方法: サーバー側のprefillグラフキャプチャの上限をスケジューラの最大prefillバッチサイズ(例: 8192トークン)に合わせて設定し、コンパイルがリクエスト中ではなくキャプチャ時に行われるようにします。さらに、本番環境で使用されるattention特殊化をカバーする、より代表的なwarmupリクエストを実行してください。診断のヒント: この問題の特徴は、TTFTのテール(p90/p99)に見られる約0.9秒の一定のストールで、中央値のTTFTとTPOTは変化せず、すべてのリクエストは正常に完了します。