TAO speeds up AI agents: first token up to 36% faster, tasks up to 27% faster

When an AI agent fixes a bug in a repository, it does more than think. The model issues a command, a compiler, tests or a search engine on separate processors execute it, the result goes back to the model, and the cycle repeats dozens of times. Researchers at Shanghai Jiao Tong University call this setup CGDC: a GPU cluster for inference plus CPUs scattered across the network for tools.

They propose TAO, a scheduler that decides which agents to run on the GPU and on which specific node, based on when each agent's tool response will return. Compared with the best competitor on each metric, TAO cuts task execution latency by up to 26.8% and mean time to first token after a session resumes (TTFT) by up to 36.

3%, and raises generation speed (tokens per second) by up to 29.5%.

Why GPUs and cache suffer while an agent waits for a tool

While a tool runs, the agent doesn't need the GPU. The session can be paused and its slot given to others. But a pause has a cost. The model keeps a KV-cache in memory: the already computed representation of the conversation history. If the session returns to the same node quickly, the prefix is reused and the whole context doesn't have to be reprocessed. If it waits too long, the cache is evicted, and recomputation eats up capacity.

The authors demonstrated this on DeepSeek-V4 Flash traces from two H100 nodes: the share of cache that survives drops noticeably as the observation ages, with the curve plotted over an interval of up to 1 800 seconds.

The authors say existing systems make two mistakes:

• They use tool runtime only as an indirect signal for cache retention or token counting, not as a standalone indicator that a session is ready.
• They score each session separately with heuristics. But sessions compete for the same capacity, and greedily choosing one can crowd out a better combination of others.

Three numbers in strict order and a knapsack problem

For each pair of a waiting session and a node, TAO computes a triple of numbers that are compared in order, like words in a dictionary:

• Cache credit. How many tokens of this session are observed to sit on the given node. Older observations decay exponentially with a set half-life. If there is no data about the node, the credit is zero.
• Return readiness. A binary flag: will the tool return within the next scheduling window? The prediction is built from hints given by the tools themselves (Claude Code has Monitor, OpenCode has a Bash timeout message), parsed sleep commands and duration history. If none of these exist, a default of five seconds is used. Stale predictions get temporary trial chances so that a session doesn't stay "ready" forever.
• Token count. If the first two are equal, longer contexts are preferred.

The order is fixed, while the values change over time. That is why no manual weights in a sum are needed.

An exact dynamic programming algorithm then does the selection. The problem resembles the classic multiple knapsack: each node has a remaining capacity, each session has an admission cost, and the goal is to maximize total value. The solution is exact for a frozen snapshot of the state, although rounding capacity may exclude plans that were feasible before rounding. The scheduler runs in a separate asynchronous process and checks the plan against live state before applying it. According to the authors, preparation and computation take less than 0.12% of runtime in a test with four replicas.

What the H100 measurements showed

They tested three models: DeepSeek-V4 Flash on two H100 nodes, Qwen3-32B on four, and GLM-4.7 on one. The workloads were: SWE-bench Lite (fixing real bugs) via mini-SWE-agent, Terminal-Bench 2.0 with Claude Code for hard terminal tasks, and SWE-verified for Qwen under high concurrency. Each node holds one replica of the model.

They compared TAO against four named approaches: the vLLM scheduler with round-robin distribution, the same scheduler with cache awareness, NVIDIA's Dynamo (KV-aware), and ThunderAgent (by token count and "decreasing" packing).

On SWE-bench Lite, relative to the best competitor in each column:

• task latency is lower by 2.4–26.8%;
• TTFT is lower by 7.4–36.3%;
• generation speed is higher by 5.1–29.5%.

On Terminal-Bench, TAO leads on latency and TTFT at all three concurrency levels. On Qwen3-32B, it has the lowest TTFT at all three numbers of parallel sessions. The maxima are taken across 12 model–workload–concurrency configurations, so the gain for any specific combination is more modest.

What engineers should do with this

The main takeaway for anyone running inference for agents themselves: treat tool runtime as scheduler data. If you already measure call durations, some of that information can be turned into GPU priority and a placement decision.

The authors claim the approach carries over to multi-agent setups and cyclic agents, as well as to ordinary chat workloads. That is their claim, and the material doesn't make clear whether separate measurements back it up. All figures were obtained on H100s and specific model combinations, so results may differ on other hardware and with a different tool profile.