TokenRouter speeds up token-level model routing by up to 64x

A team from the thu-nics lab (Tsinghua) has released TokenRouter, a system that lets a large and a small language model share a single generation: for each next token, the system decides which model writes it. The headline "up to 64x" figure is the ratio of TokenRouter's throughput (tokens per second across all parallel requests) to the official implementations of existing routing algorithms. The gain does not come from new math for choosing a model. It comes from engineering: the authors removed the overhead that eats almost all the time in raw implementations.

Request routing is a mature field; token routing is not

The author of the original post in @gonzo_ML notes that directing requests to different models is nothing new. His company, Intento, has been doing it for almost ten years. That is why he was surprised by the OpenRouter valuation and the Stripe deal: according to Menlo Ventures, the payments giant is buying the model-aggregator service.

Such a router looks at the whole request. A simple question goes to a cheap model, a complex one to an expensive model. Examples include commercial services and the open-source RouteLLM.

TokenRouter works at a different level. Within a single response, difficulty is uneven. Most tokens are trivial: word endings, connectives, boilerplate code. Only at certain forks in the reasoning is a strong model needed. Entropy serves as the difficulty criterion: the less certain the small model is about the next token, the higher the entropy. When entropy is high, the token goes to the large model. Otherwise, the small one writes it.

In theory, this lowers inference cost: the heavy model works only where it cannot be avoided. In practice, the scheme runs into overhead.

Where tokens per second disappear

A naive implementation calls one model, then the other, at every step, recomputes the context, and shuttles data between processes. All of these operations are costly. The post's author calls this the most interesting part of the work: without optimization, the surrounding machinery is very expensive.

TokenRouter is built on top of SGLang, an open-source engine for running LLMs, and consists of several parts:

• A three-function API. The routing algorithm is described through a small interface, so different routers plug in without rewriting the runtime.
• Asynchronous sub-servers. Each model has its own sub-server, and they run in parallel without waiting for each other.
• KV-cache preservation between calls. The KV-cache stores the already computed intermediate states of the context. Without saving it, the context would have to be recomputed from scratch at every model switch.
• A delayed-batching scheduler based on Markov chains. A GPU processes a batch of requests more efficiently than one at a time. The scheduler models the stream of calls to each sub-server and decides whether it is worth waiting a little while a batch accumulates.

For R2R, one of the existing token-level routers, the authors separately measured the contribution of CUDA graphs. This is an NVIDIA technology: a chain of GPU operations is recorded once and launched as a whole, without spending time on individual calls. Of the total 2.76x speedup, CUDA graphs account for 1.71x.

Numbers on olympiad math and agentic code

The system was tested on five routing algorithms. The tasks were of two types: reasoning on AIME2024 (problems from the American math olympiad) and agentic scenarios on SWE-Smith (a set of code-editing tasks for software agents).

At 4 parallel requests:

• TokenRouter: 76.6–290.3 tokens/s, depending on the algorithm.
• Token-level baseline on SGLang: 4.5–63.1 tokens/s.
• Official implementations: 0.7–78.3 tokens/s.
• Gain: from 2.01x to 64.15x.

On the R2R router specifically:

• 8 requests: 134.9 → 372.5 tokens/s (2.76x).
• 16 requests: TokenRouter outperforms R2R by 3.61x at the same concurrency.
• Against R2R running a single request, the gain is 18.58x, while per-user speed rises 1.13x.

The last pair of figures shows the practical point. The system can be loaded with many users without losing speed for each of them.

What it means in practice, and the catch

The authors show that token-level routing can be hidden behind a single OpenAI-compatible endpoint. To the client it is an ordinary API, while inside, two models split the generation. In terms of the speed-to-accuracy trade-off, such a scheme may beat request-level routers like RouteLLM. For companies that pay for inference, this is a new lever: the price per token falls in proportion to the share of "easy" tokens, while quality at the hard forks is preserved.

There are two limitations an engineer should consider before adopting it:

• Each sub-server has to store the KV-cache of the entire sequence. GPU memory use nearly doubles, because the cache exists for both the small and the large model.
• The scheduler's Markov solver has been tested only up to 4 parallel requests. How it will behave across hundreds of sessions in production is not yet known. The figures for 8 and 16 requests were obtained on R2R.

The code itself is openly available (thu-nics/TokenRouter). Teams that serve their own models can test the approach on their own workloads. The main lesson of the work is a practical one: the idea of token routing has been known for a long time, but without careful systems engineering, its gains are lost in the overhead.