TAO Makes Tool-Return Readiness a First-Class Signal for Agent Serving

Introduction

Agentic LLM services do not behave like ordinary continuous-generation workloads. A session may run a model on clustered GPUs, call a tool on distributed CPUs, wait for the result and then return to the model for another generation phase. A scheduler that only sees the GPU queue can therefore make poor decisions: a session waiting on a tool may occupy valuable GPU state, while other sessions are repeatedly loaded and evicted from the KV cache.

The TAO framework—Tool-Aware Joint Scheduling—treats the tool phase as part of the scheduling problem rather than as an external pause. Its central question is not only which request needs GPU time, but also which request is likely to become runnable soon and which node can resume it most efficiently.

How TAO works

• Interruption at tool boundaries. When a session enters an external tool call, TAO pauses it and releases GPU capacity. The session moves to a shared waiting pool and can re-enter scheduling after the tool result arrives, instead of holding resources while it cannot generate.
• Triplet session valuation. For each session-node pair, TAO combines near-term tool-return readiness, node-specific KV-cache credit and token footprint. Readiness is estimated from runtime hints and historical tool durations. Cache credit decays over time, preventing old locality from dominating decisions indefinitely.
• Joint admission and placement. Rather than routing sessions one by one with a greedy heuristic, TAO uses an asynchronous dynamic program to select waiting sessions and place them under GPU capacity constraints. The reported scheduling overhead is below 0.12% of runtime, with asynchronous revalidation handling changes after a plan is produced.

Results and trade-offs

The evaluation covers DeepSeek-V4 Flash, Qwen3-32B and GLM-4.7 across 12 model-workload-concurrency settings. Against the best evaluated state-of-the-art baseline for each metric, TAO reduces latency by up to 26.8%, lowers mean time to first token by up to 36.3% and raises decode throughput by up to 29.5%.

The results also expose an important systems trade-off. In a highly saturated workload, TAO may deliberately delay admission rather than immediately push more sessions into the GPU engine. That can increase client-visible TTFT, but it protects active KV-cache state from thrashing, preemption and recomputation. In the AutomationBench case with 128 workers, mean TTFT is reported as 19.1 seconds for TAO versus 3.2 seconds for Dynamo. Yet overall task latency falls from 1,417.2 seconds to 563.1 seconds, while decode throughput rises from 482.4 to 1,013.2 tokens per second.

Why it matters

TAO’s broader contribution is a different definition of scheduling context. An agent-serving system must reason about GPU capacity, cache locality and the timing of external tools at the same time. The benefit of joint optimization can also compound across turns: a better early placement may preserve more KV cache and avoid fragmentation, improving the states seen by later scheduling decisions.

This approach is especially relevant to tool-heavy, bursty services operating under tight cache budgets. Its practical value will still depend on the quality of tool-duration estimates, the degree of hardware heterogeneity and the chosen balance between immediate TTFT and end-to-end completion time. Integrating such a policy into an existing inference engine may be as important as the scoring model itself.

Source: Hugging Face Daily Papers