Command Palette
Search for a command to run...
TAO: TOOL-AWARE JOINT SCHEDULING FOR LLM-BASED AGENTIC SERVING
TAO: TOOL-AWARE JOINT SCHEDULING FOR LLM-BASED AGENTIC SERVING
Shuhuai Li Dongdong Ge Yinyu Ye Shengdong Ye Jianghao Lin
Abstract
Agentic services alternate between model inference on clustered GPUs and external tools on distributed CPUs. In this scenario, existing systems (1) overlook tool-return readiness in service evaluation, wasting limited resources, and (2) use heuristic strategies to decide admission and placement on GPUs, resulting in suboptimal scheduling decisions. To bridge this gap, we present TAO, a tool-aware joint scheduling framework for LLM-based agentic serving. Its value interface combines near-term tool readiness, node-specific cache credit with time-aware decay, and token footprint; a dynamic program jointly selects and places waiting sessions under capacity constraints. Across the 12 reported model–workload–concurrency settings, TAO reduces latency by up to 26.8%, reduces mean time to first token (TTFT) by up to 36.3%, and increases decode throughput by up to 29.5% relative to the best evaluated state-of-the-art baseline for each metric.
One-sentence Summary
Researchers from Shanghai Jiao Tong University and Stanford University propose TAO, a tool-aware joint scheduling framework for LLM-based agentic serving that combines near-term tool readiness, node-specific cache credit with time-aware decay, and token footprint in a dynamic program to jointly select and place waiting sessions under capacity constraints, achieving across the 12 reported model-workload-concurrency settings up to 26.8% lower latency, up to 36.3% lower mean time to first token, and up to 29.5% higher decode throughput relative to the best evaluated state-of-the-art baseline for each metric.
Key Contributions
- TAO is a tool-aware joint scheduling framework for LLM-based agentic serving that combines near-term tool readiness, node-specific cache credit with time-aware decay, and token footprint into a unified session value interface.
- A dynamic program jointly selects and places waiting agent sessions under snapshot capacity constraints, while session interruption and asynchronous revalidation connect the resulting admission and placement plan to live execution.
- Across 12 reported model-workload-concurrency settings, TAO reduces latency by up to 26.8%, mean time to first token by up to 36.3%, and decode throughput by up to 29.5% relative to the best evaluated state-of-the-art baseline for each metric.
Introduction
The authors study LLM-based agent serving in a Clustered-GPUs with Distributed-CPUs (CGDC) scenario, where external tool calls create variable gaps between inference requests and make GPU demand interruptible. This matters because tool execution now contributes substantial latency, so scheduling must account for both inference capacity and when sessions will need GPU resources again. Prior systems typically use evaluate-then-decide heuristics that value sessions individually and treat tool timing mainly as a surrogate for KV-cache reuse or token count, leading to myopic and suboptimal placement when sessions compete for shared capacity. The authors propose TAO, a tool-aware joint scheduling framework that introduces tool-return readiness into session evaluation and uses dynamic programming to jointly optimize admission and placement across the waiting pool, improving latency, TTFT, and decode throughput over state-of-the-art baselines.
Method
The authors present TAO, a framework that combines agentic session evaluation with decision making via dynamic programming for joint selection and placement. The pipeline separates session evaluation from decision making, allowing the evaluation interface to evolve while preserving the scheduling procedure. As shown in the figure below:
At each restart scheduling decision, the waiting pool P contains agent sessions awaiting inference admission, including interrupted sessions and new arrivals. Let B be the nodes in the frozen snapshot. Each agent session i has an admission cost wi, and each node b has residual capacity Cb. Let E⊆P×B denote eligible assignments, and let xi,b∈{0,1} indicate whether session i is admitted to node b. The feasible set is defined as:
X={x:b∈B∑xi,b≤1,i∈P∑wixi,b≤Cb,xi,b=0 if (i,b)∈/E}.TAO assigns each agent session-node pair a lexicographic value triplet:
vi,b=(ki,b,ri,ti),where ki,b is the decayed cache credit, ri is the binary readiness signal, and ti is the token count. The system jointly selects and places sessions to maximize their aggregate value:
x⋆∈arglexmaxx∈Xi∈P∑b∈B∑xi,bvi,b.The evaluation triplet order remains fixed while its values evolve over time. KV-Cache credit takes precedence because prefix reuse directly avoids repeated prefill. As unrefreshed credit decays, tool-return readiness distinguishes agent sessions nearing a return from those still far from their next model turn. Finally, token footprint favors larger accumulated contexts among plans tied on cache credit and readiness.
To estimate reusable KV state, TAO uses an observed cached-token count or estimates it from observed token volume and cache rate. The decline in cache-token retention over time motivates discounting older observations, as illustrated in the figure below:
Older observations receive less credit according to the time-aware decay formula:
ki,b(t)=ki,b02−τi,b/h,τi,b=max(0,t−ti,bobs).Here ki,b0 is the observed cache credit and h is its decay half-life, set to 10 seconds. This discount makes the admission value responsive to observation age.
For tool-return readiness, the predictor combines hints from parsed sleep commands and historical tool durations. For a call with elapsed time ei(t), the historical estimate Df gives:
di(t)=max{0,Df−ei(t)}.Binary readiness is defined as:
ri(t)=⎩⎨⎧1,ai(t)bi(t),ai(t)gi(t)1[zi(t)≤t+H],oi(t)=1,oi(t)=0,zi(t)≤t,otherwise.Here oi(t) indicates observed tool completion, ai(t) indicates an established tool execution candidate, and gi(t) indicates valid prediction timestamps. zi(t) is the predicted completion or exposed return deadline, and H is the horizon set to 1 second. Observed completion takes precedence over prediction checks.
The optimizer is a capacity-state dynamic program for the multiple-knapsack and generalized assignment structure. Let Δ be the capacity quantum, wˉi=⌈wi/Δ⌉, and Cˉb=⌊Cb/Δ⌋. For candidate j, let Dj(s) be the best objective attainable from the first j candidates at the used capacity s, and let uj,b embed the value and count for assigning j to node b:
Dj(s)=lexmax{Dj−1(s),b eligiblemaxlex[Dj−1(s−wˉjeb)+uj,b]}.Only nonnegative states within Cˉ are allowed. Skipping or assigning each candidate exhausts the feasible choices, providing exactness for the frozen, discretized problem. Node-specific cache observations are decayed and folded into assignment values during preprocessing, and cache-equivalent nodes can be grouped into aggregate capacity dimensions. The planner runs in an isolated asynchronous worker and revalidates before application.
Experiment
The experiments evaluate TAO against vLLM round-robin and cache-aware schedulers, Dynamo, and ThunderAgent across DeepSeek-V4 Flash, Qwen3-32B, and GLM-4.7 on agentic workloads including SWE-bench Lite, SWE-verified, and Terminal-Bench 2.0. Across models, concurrency levels, and topologies, TAO consistently achieves the strongest overall tradeoff among task latency, time to first token, and decode throughput. Component breakdowns show that cache-aware valuation, tool readiness signals, and joint scheduling each improve cache reuse, reduce pause wait, and accelerate model responses after tool completion. Ablations further confirm that removing any of these components reduces throughput, cache hit rate, and responsiveness, validating the combined design.
TAO is the strongest overall choice across the tested models, concurrency levels, and topologies, combining consistently low latency and TTFT with higher decode throughput. On DeepSeek-V4 Flash, it improves latency, TTFT, and decode throughput on SWE-Lite and leads latency and TTFT on Terminal-Bench. On Qwen3-32B, it leads TTFT and decode throughput at higher worker counts, with latency reductions on most SWE-verified settings. TAO improves latency, TTFT, and decode throughput on DeepSeek-V4 Flash SWE-Lite relative to the best competing method at each worker count shown. On Qwen3-32B SWE-verified, TAO has the lowest TTFT and highest decode throughput across worker counts and reduces latency at most evaluated settings.
In a high-concurrency GLM-4.7 setting on one node with eight H100s, TAO outperforms ThunderAgent and other vLLM-based schedulers on SWE-Lite, with the largest relative gains at 160 workers. The pattern includes lower mean latency and client TTFT alongside higher decode throughput. On the full SWE-verified set, TAO similarly reduces latency and TTFT while increasing decode throughput compared with ThunderAgent. TAO achieves lower latency, lower TTFT, and higher decode throughput than ThunderAgent across all reported SWE-Lite worker counts, with the strongest improvements at 160 workers. Compared with Dynamo and vLLM scheduler variants, TAO cuts latency and TTFT by large margins at every concurrency level in the GLM-4.7 SWE-Lite run. In the larger GLM-4.7 evaluation, TAO reduces SWE-verified mean task latency and client TTFT while raising decode throughput relative to ThunderAgent.
On the mixed Terminal-Bench and SWE-Lite workload, TAO achieves the lowest latency and highest decode throughput across 32, 64, and 128 workers relative to Dynamo, vLLM scheduler variants, and ThunderAgent. Its TTFT is lowest at 64 and 128 workers and only slightly behind the best result at 32 workers. TAO records the lowest latency at every evaluated worker count. Decode throughput under TAO is higher than all listed baselines at every evaluated worker count. TAO has the lowest TTFT at 64 and 128 workers, while at 32 workers it is close to the best observed TTFT. At 128 workers, TAO shows a notable latency and TTFT advantage over ThunderAgent while maintaining higher decode throughput.
The full TAO configuration delivers the strongest overall efficiency on SWE-Lite, with the highest decode throughput and cache-hit rate and the lowest mean TTFT, tool-end-to-first-token delay, and pause wait. Removing cache value or readiness modestly reduces throughput and responsiveness, with readiness loss having a larger effect than cache value loss. The largest regression occurs when decision making is replaced by best-fit-decreasing, which yields the lowest decode throughput and the slowest latency metrics. The complete TAO configuration outperforms all variants in decode throughput, cache-hit rate, mean TTFT, tool-end-to-first-token delay, and mean pause wait. Ablating readiness reduces decode throughput more than removing cache value, while cache-hit rate remains relatively stable across variants. Replacing DP resumption with best-fit-decreasing causes the largest performance drop, cutting decode throughput and producing the highest mean and P90 latency timings.
The experiments evaluate TAO across multiple model configurations, benchmarks, and concurrency levels, including DeepSeek-V4 Flash, Qwen3-32B, GLM-4.7, SWE-Lite, SWE-verified, Terminal-Bench, and mixed workloads, comparing it with ThunderAgent, Dynamo, and vLLM-based schedulers. TAO consistently achieves lower latency and time-to-first-token while maintaining higher decode throughput, with the largest gains in high-concurrency GLM-4.7 runs and mixed workloads. Ablation studies confirm that the full TAO configuration is important for efficiency and responsiveness, where removing readiness or cache value causes modest regressions and replacing decision making with best-fit-decreasing produces the largest drop.