Build the next-generation LLM inference engine oriented to Agentic reasoning
Why Does Agentic Reasoning Require a New Inference Engine
At the very beginning of the TokenSpeed project, the first question we needed to answer was: why does the Agent scenario require a completely new inference engine design.
The optimization goal of traditional Chat-class workloads is relatively single: a user initiates a request, and the model generates a complete full response in one go. In this mode, the inference system only needs to fill the batch as much as possible to maximize throughput, and the KV Cache is released after one loading, leading to inherently low utilization. However, the load characteristics of Agent scenarios are fundamentally different from those of Chat. Taking the coding agent as an example, the model needs to continuously think, call tools, obtain results, and continue thinking based on the results. The whole process consists of multiple rounds of short request loops. The first request may carry a prefix context of tens of KB, but the output of each subsequent loop is often only a few hundred tokens. This means that the inference engine must re-examine the cache reuse strategy, request scheduling mechanism, and the optimization objective function of the entire system.
Under this load pattern, users' sensitivity to latency increases significantly. If an Agent tool experiences significant lag during the inference phase, users will intuitively feel insufficient "responsiveness". Therefore, what really needs to be optimized in the Agent scenario is throughput under latency constraints. We internally defined a performance metric: the TPS per user is no less than 70, and under this constraint, we strive to maximize the throughput capacity of each GPU. Taking TensorRT-LLM as the benchmark, TokenSpeed is 9% faster than TensorRT-LLM in end-to-end latency in the minimum latency scenario where the batch size is equal to 1; at the typical operating point of 100 TPS per user, the throughput of each GPU is 11% higher.
The KV Cache reuse rate is a key variable that determines inference costs. Taking DeepSeek as an example, its cache reuse rate is reportedly over 99%, while most systems in the industry can only reach 70% to 80% at present. In the Agent scenario, if the same tens of KB of context can be hit repeatedly, it means that each hit can completely eliminate the prefill calculation, and the overall cost of the system will be greatly reduced. We verified the performance of TokenSpeed when the context length expands from 128K to 1M in tests, and the decoding throughput only drops by 16%. In addition, as more and more models adopt linear attention mechanisms, including Qwen's Gated DeltaNet and the KDA used by Kimi K3, the management mode of KV Cache is also changing. Linear attention models no longer maintain the traditional one-to-one correspondence between tokens and KV, but represent the context through a continuously iterated state, including SSM state and convolution state. How to efficiently update, save and restore these states in the inference engine is a new system design proposition.
The coupling between model architecture and hardware features is also deepening simultaneously. The MOE architecture has become the mainstream choice since the release of DeepSeek, but optimizing MOE models is not easy. Taking Kimi K3 as an example, this model has 896 experts. Although the Latent MOE technology is used to greatly reduce the writing cost and increase the proportion of activated experts to 16%, the complexity of routing and communication increases significantly. At the same time, the Blackwell generation of hardware introduces a series of new features such as Tensor Memory and a new generation of Tensor Core, and AMD's CDNA architecture also has its unique hardware capabilities. If the inference engine cannot make full use of these hardware features, the performance advantage of Blackwell over Hopper cannot be truly unleashed. Based on these observations, TokenSpeed has carried out in-depth optimization for Blackwell at the Kernel level to ensure that the hardware potential is fully tapped.
The co-design of Runtime and Kernel cannot be ignored either. No matter how well the scheduler is designed, if the underlying Kernel has performance bottlenecks or excessive coupling, the overall system still cannot reach the ideal state; conversely, no matter how extreme the Kernel optimization is, if the scheduler generates a large number of bubbles, the throughput cannot be fully utilized. In terms of architecture design, TokenSpeed has learned from this experience, and systematically decoupled and reconstructed scheduling, execution and underlying operators.
Runtime: Building a Scalable Inference Engine Architecture
The overall architecture of TokenSpeed is divided into three layers: Control, Runtime and Kernel. The control plane is written entirely in C++, the Runtime is implemented in Python, and the Kernel is independent as a sub-package. The core purpose of this layered design is to maintain the clarity of module boundaries. The control plane is only responsible for scheduling decisions and does not involve any data plane execution; the Runtime is only responsible for executing the execution plans issued by the control plane, including cache plans and forward calculation plans; the Kernel layer focuses on operator implementation. The three perform their respective duties, avoiding the maintenance dilemma caused by the mutual penetration of scheduling logic, execution logic and operator implementation in past inference engines.
There are practical historical reasons for choosing C++ instead of Rust for the control plane. Most of the team members have long-term development experience and are not very familiar with Rust. C++ fully meets the requirements in terms of performance and expressiveness, so we finally decided to use C++ for the control plane programming. The reason why Runtime retains Python is that the Python ecosystem is already very mature in the deep learning field, with high integration of libraries such as PyTorch and CUDA, and the development efficiency of the model adaptation layer is much higher than that of other languages. After the Kernel sub-package is made independent, NVIDIA and AMD experts can perform in-depth optimization for their respective hardware features without needing to understand the internal implementation of the scheduler or Runtime.
Another benefit brought by this layered architecture is the automation of parallel compilation. In the past development of inference engines, the code for model parallel splitting was often scattered in various model adaptation modules, and deeply coupled with attention implementation, MOE implementation, communication logic, etc. As the number of supported models increases, these scattered parallel codes become more and more difficult to maintain, and the access of any new model may involve modifications to multiple modules. TokenSpeed's approach is to completely separate the parallel splitting logic from the model adaptation code to form an independent Placement system. Model developers only need to declare the boundary of each module, the input and output types, and the attributes of the module in the parallel dimension, whether it is replicate, shard or partial, and the system will automatically handle the subsequent splitting and communication.
Specifically, each module can be mounted to different parallel groups. The attention module is mounted to the Attention Tensor Parallel group, and the MOE module is mounted to the MOE parallel group, and each module independently selects the parallel strategy. Based on these Placement definitions, TokenSpeed's type compilation system will traverse the execution plan of the entire decoding layer and compare the Placement differences between adjacent modules. When an inconsistency is found, the system automatically inserts the corresponding communication primitives: all-reduce is required from partial to replicate, all-gather is required from shard to replicate, and reduce-scatter is required from shard to partial. The output of the entire compilation process is a unified compiled_decoder_layer execution plan, which is accessed in three stages at runtime: the prefill stage performs pre-communication alignment, then executes model operators, and finally executes communication output. This automated parallel compilation not only greatly reduces the maintenance cost of model adaptation, but also makes fusion between adjacent operators possible because the type system provides pre-information, reducing the overhead of kernel launch, and also bringing positive benefits to performance.
The core of the control plane is a state machine that covers the full life cycle of a single request. From entering to completion, an inference request inevitably goes through several key states: submitted, prefilling, decoding, retracting/timeout, finished. Each state itself is a resource container, holding resources such as KV page references required for this state, the slot list of SSM/MLA, etc. When the state migrates, the ownership of these resources is synchronously transferred to the next state. The entire control plane is a continuously flowing state machine system.
State migration is completed through an event-driven mechanism. Each event carries the corresponding C++ operator, and during the state transition process, the ownership of resources is transferred through the unique_ptr move semantics of C++. The cleverness of this design is that it advances the runtime resource management problem to the compilation period. When the code is compiled, the compiler has actually verified the problems of resource leakage and resource allocation errors. Because all resource transfer paths are explicitly defined and checked in the type system. The resources themselves are managed in the RAII (Resource Acquisition Is Initialization) way. Whether it is the reference count of the Cache Block, or the reference count of the Device Node or Host Side, the corresponding resources will not be released in advance before the data movement or operation is completed.
The state machine design of the control plane is completely decoupled from the execution plane. The scheduler only needs to focus on when to advance the request from the prefilling state to the decoding state, and when to trigger retracting, without needing to care about how the underlying Kernel performs the specific calculation. This decoupling allows the scheduling strategy to be iterated independently, without frequent modifications due to adjustments of underlying operators. The Runtime layer receives the execution plan from the control plane, executes the calculation according to the sequence and resource configuration defined in the plan, returns the result to the control plane after execution, and drives the state machine to continue to flow.
In actual engineering, this state machine design significantly reduces race conditions and resource leakage problems in concurrent request scheduling. Resource management bugs that were common in past Python implementations, such as KV pages not being recycled in time after a request is canceled, or reference counts not being correctly decremented in exception paths, have been systematically solved in C++ implementation through compile-time checks and RAII mechanisms. This is also one of the core motivations for us to completely rewrite the control plane in C++.
KV Cache management is one of the modules that TokenSpeed has invested the most energy in architecture design. The prefix reuse in Agent scenarios has much higher requirements for the cache system than traditional Chat scenarios. The same Agent request will repeatedly carry the same prefix context in a complete workflow. If prefill needs to be recalculated in each loop, the inference cost will expand sharply.
In the first version of the implementation, we drew on the idea of SGLang and used Radix Tree as the management structure for prefix caching. The starting point of Radix Tree design is to maintain a tree structure from root to leaf, each node on the tree corresponds to a segment of token sequence, and the prefix matching of all requests is carried out on this tree. This design works well in the pure attention scenario, but when we started to adapt Kimi K3, we found that the semantics of Radix Tree conflict with some new attention patterns. Mamba-class models and Sliding Window Attention naturally break the hierarchical matching assumption that Radix Tree relies on, and their cache access pattern is not a simple linear extension from prefix to suffix.
This experience prompted us to refactor the KV Cache management system. In the process of supporting Kimi K3, we migrated all the original Radix Tree to the Flat KV design. The new cache system no longer perceives the attention type of the cache, and only requires the upper layer to report different types of caches to the system in a unified format, declaring their respective types and occupied space. The bottom layer uniformly packages all caches into blocks of the same size for flat storage management. The first version of K3 uses the method of equal logical block and physical block, and manages 1536 tokens as one block. However, this granularity leads to prominent problems of prefix reuse efficiency and intra-page fragmentation in long context scenarios. Later, we drew on the practice of vLLM and introduced a two-layer block design: the first layer is a logical block with a uniform block size of 128; the second layer is a physical block with a block size of 1536. The logical block of size 128 is used for prefix reuse and finer-grained cache matching, while the physical layer remains unchanged. This two-layer design achieved good optimization effects in subsequent tests.
Cache preemption and recovery is another key capability. TokenSpeed is now fully integrated with the KV Store system of Mooncake. When the video memory is insufficient, the entire KV state of the request is unloaded to the KV Store for persistence, and when the request is rescheduled back, it is reloaded through a lazy reload mechanism to continue the calculation. In terms of PD separation (Prefill and Decode separation), we initially implemented a layer-by-layer transmission method. But layer-by-layer transmission leads to very low RDMA bandwidth utilization, because the transmission of each layer requires independent network round-trip overhead. Later, we optimized it to allow users to specify several layers to be packaged into one block for transmission, that is, a customizable layered transmission strategy. This flexibility allows users to adjust the transmission granularity according to the specific network topology and model scale.
Kernel: Making MLA Fully Utilize GPU Performance
The Kernel layer optimization work in TokenSpeed is mainly promoted by NVIDIA and AMD experts in the team. At present, TokenSpeed's Kernel has reached state-of-the-art performance level on Blackwell, and has been integrated by the vLLM community as the default Kernel package on Blackwell.
MLA (Multi-head Latent Attention) is an attention mechanism widely adopted by domestic open source models. Although MLA saves writing cost by compressing KV into a single latent state, its calculation pattern is not naturally hardware-friendly. In the decode phase, in the first matrix multiplication of query and latent key, the dimension of query is 1, and the entire matrix multiplication actually degenerates into vector multiplication. This means that the computing resources of Tensor Core are seriously wasted on Blackwell. The second bottleneck is bandwidth. In Agent scenarios, the prefix length often reaches tens of KB or even hundreds of KB. Even with INT8 quantization, bandwidth data loading is still a limiting factor. The third problem is shape mismatch: the combination of the number of heads of MLA and the training length often leads to a tall and thin matrix multiplication, which has a negative impact on SM utilization and launch overhead.
Aiming at these problems, we redesigned the Kernel on Blackwell. First, make full use of the newly introduced on-chip memory resource of Tensor Memory. We specially wrote the allocator for Tensor Memory, allowing Tensor Core to directly read and write Tensor Memory, avoiding the overhead of intermediate data movement through registers. Secondly, an optimization technique of fold factor is introduced at the matrix multiplication level. Since the query dimension is 1 in the decode phase, the Tensor Core utilization is extremely low. We found that when the number of heads is less than 128, the query tokens can be folded into the head dimension. After folding, the effective head dimension is multiplied by the fold factor, the matrix dimension becomes higher, and Tensor Core can perform more efficient calculations. In addition, we also designed the Split KV mechanism to split the calculation into two stages: the first stage calculates the partial sum of Split KV, and the second stage performs the reduction operation. The Split parameter and the size of the workspace are automatically calculated and cached at runtime. The collaborative execution across