TAO: Tool-Aware Joint Scheduling for LLM-Based Agentic Serving

AI summary

Overview

Research area: Distributed computing and machine learning systems — specifically, scheduling and serving infrastructure for LLM-based agents (arXiv category cs.DC), spanning LLM inference serving, KV-cache management, and combinatorial optimization.

Technical level: Advanced. The paper formalizes scheduling as a lexicographic multiple-knapsack / generalized-assignment problem solved by dynamic programming, and assumes familiarity with terms such as KV-cache reuse, prefill, time to first token (TTFT), and tensor-parallel replicas. The paper's high-level argument, however, is explained in accessible terms.

Scope (one sentence): The paper introduces TAO, a tool-aware joint scheduling framework that decides which waiting agent sessions to admit and which GPU node to place each on, using a three-part session value (decayed KV-cache credit, binary tool-return readiness, token footprint) maximized by an exact dynamic program.

What This Paper Is About

LLM agents alternate between model inference on clustered GPUs and calls to external tools (such as compilers and retrievers) that run on distributed CPUs. This creates gaps between LLM requests that can be used to reallocate GPU capacity — a setting the authors call Clustered-GPUs with Distributed-CPUs (CGDC). Existing serving systems, the authors argue, overlook when a tool call is about to return when valuing a session, and use heuristic, per-session rules to decide admission and placement, producing myopic and suboptimal decisions. TAO's goal is to fold tool timing and cache reuse into one unified session value, then jointly select and place sessions under shared capacity constraints.

Key Contributions

1. Identifying the CGDC serving scenario. The authors name and characterize the setting where clustered GPUs provide shared inference capacity while external tool calls on distributed CPUs introduce variable delays, making a session's demand for GPU computation interruptible.
2. A tool-aware session evaluation interface. TAO values each session–node pair as a lexicographic triplet: decayed KV-cache credit for that specific node, a binary tool-return readiness signal, and token footprint. Readiness is estimated from exposed timing hints, parsed sleep commands, and historical tool durations.
3. Joint selection and placement via dynamic programming. Instead of scoring sessions one at a time, TAO solves a capacity-state dynamic program over all waiting sessions at once, maximizing the aggregate value under per-node residual capacity, with asynchronous execution and live-state revalidation before the plan is applied.
4. Empirical evaluation across models, workloads, and concurrency levels. Across the 12 reported model–workload–concurrency settings, TAO reduces latency by up to 26.8%, reduces mean TTFT by up to 36.3%, and increases decode throughput by up to 29.5%, each relative to the best evaluated state-of-the-art baseline for that metric.

Main Findings

• Headline gains on DeepSeek-V4 Flash with SWE-Lite: relative to the best competing method in each column, latency falls by 2.4–26.8%, TTFT by 7.4–36.3%, and decode throughput rises by 5.1–29.5%. The largest figures (26.8% latency, 36.3% TTFT, 29.5% decode) occur at 64 workers.
• Terminal-Bench results are not uniformly better: TAO leads on latency and TTFT at all three concurrency levels for DeepSeek-V4 Flash, but its decode-throughput improvement is reported as −8.8% at 16 workers, +1.3% at 32 workers, and −6.2% at 64 workers.
• Qwen3-32B at higher concurrency: on SWE-Lite, TAO achieves the lowest TTFT at all three worker counts (96, 128, 160) and the highest decode throughput at each. On SWE-verified, TTFT drops by 12.9–29.9% and decode throughput rises by 4.0–7.8%; latency falls 8.3% and 9.0% at 96 and 160 workers, but is 2.5% higher than ThunderAgent at 128 workers.
• A larger model with 8×H100 on one node (GLM-4.7): on SWE-Lite at 160 workers, improvements reach 15.5% in latency, 14.05% in TTFT, and 20.48% in decode throughput. Across all 500 SWE-verified tasks, TAO reduces mean task latency by 16.2–18.9% and client TTFT by 15.7–38.4%, while increasing decode throughput by 22.8–24.7% relative to ThunderAgent.
• Component attribution for cache-aware valuation (218-task fixed-workload replay with DP placement held fixed): cache-token ordering raises prefix-cache hit rate by 0.19 percentage points and reduces mean pause wait by 4.5%; adding cache-credit decay raises hit rate by 0.31 percentage points and decode throughput by 1.68%.
• Component attribution for tool readiness: enabling readiness in paired SWE-Lite runs increases decode throughput by 4.03% and reduces tool-end-to-first-token delay by 10.88%, with mean pause wait falling from 1.067 s to 0.972 s.
• Readiness prediction accuracy: across 10,967 completed tools, the final duration prediction has a mean absolute error of 0.410 s, with 82.9% of errors within 0.5 s.
• Joint optimization vs. a short-first heuristic: the dynamic program improves the lexicographic objective in 6 of 23 replay snapshots with multiple candidates (26.1%) and matches the heuristic in the other 17; across all 35 reconstructable snapshots it improves 6 (17.1%) and matches 29 (82.9%).
• Ablations (DeepSeek-V4 Flash, 2 nodes with 4×H100 per node, 32 workers): full TAO reaches 1417.97 tok/s decode and 96.31% cache hit. Removing cache value gives 1398.35 tok/s; removing readiness gives 1363.06 tok/s; using ThunderAgent's pause ordering gives 1355.20 tok/s; replacing the DP with best-fit-decreasing (BFD) gives 1233.72 tok/s. Both valuation and decision making contribute.
• Regressions on mixed workloads at high concurrency: on AutomationBench (800 tasks) at 128 workers, TAO is reported at −7.0% latency, −504.0% TTFT, and −11.2% decode — showing the approach does not dominate everywhere.
• Low scheduling overhead: recorded preparation plus completed-solve work is below 0.12% of runtime in the four-replica sweep.
• Temporal mechanics: cache credit decays with a half-life of h = 10 seconds; readiness uses a 1-second horizon, requires prediction timestamps refreshed within 30 seconds, and falls back to a five-second duration default when no history exists.

Methodology in Plain English

The authors start from the observation that an agent session doing a tool call does not need the GPU — so it can be safely interrupted, releasing its slot, and restarted later. An interrupted session joins a waiting pool alongside newly arriving sessions. The scheduling question becomes: at each restart opportunity, which waiting sessions should get GPU capacity, and on which node?

TAO answers this in two decoupled stages, which the authors describe as an improvement on the standard "evaluate-then-decide" pattern.

Stage one — value each (session, node) pair. Every candidate assignment gets a three-part score, compared in a fixed priority order:

1. Cache credit — how much reusable KV-cache state already exists on that node, discounted over time via an exponential half-life of 10 seconds, since older observations are less useful when a tool returns quickly. A node with no matching evidence gets zero credit.
2. Tool-return readiness — a yes/no flag for whether the session's tool call is expected to finish within a short horizon (1 second), or has already finished. The predictor prefers direct timing evidence such as exposed progress events or parsed sleep commands, then falls back to tool-specific history, global history, or a five-second default, subtracting elapsed time. A conservative finite-sample margin is added to the historical mean (using a Maurer–Pontil-style bound with δ = 0.2). Overdue predictions get temporary "probe" opportunities rather than staying marked ready forever.
3. Token footprint — the accumulated context size, which breaks ties but also costs capacity.

Because the three components are compared lexicographically rather than summed, no weighting coefficients need tuning.

Stage two — pick and place together. TAO discretizes session admission costs and node residual capacities by a quantum and runs a dynamic program over capacity states, processing candidates one at a time and choosing between skipping a candidate or assigning it to an eligible node. This is exact for the frozen, discretized snapshot, though rounding can exclude plans that were feasible before discretization. The planner runs in an isolated asynchronous worker; before applying a plan it checks that selected sessions still match their recorded state signatures and remain eligible, then applies only the longest prefix that still fits current capacities — otherwise the plan is discarded. If the budget is exceeded, it falls back to a BFD plan.

How it was tested. The authors compared TAO against four baselines — Dynamo (KV-aware scheduling), vLLM Scheduler with round-robin, vLLM Scheduler with cache-aware policy, and ThunderAgent (token-count scheduling with best-fit-decreasing placement) — on DeepSeek-V4 Flash (2 nodes, 4×H100 each), Qwen3-32B (4 nodes, 2×H100 each), and GLM-4.7 (one node, 8×H100, TP8). Workloads included SWE-bench Lite via the mini-SWE-agent harness, Terminal-Bench 2.0 with Claude Code, SWE-verified, AutomationBench, and a mixed workload. Metrics were mean end-to-end task latency (including tool execution), mean TTFT, and decode throughput in tokens per second, plus prefix-cache hit rate, tool-end-to-first-token delay, and pause wait in component analyses.

Why This Matters

Impact on research. The paper reframes agentic serving as a joint admission-and-placement problem rather than a per-session routing problem, and argues that tool timing is a first-class scheduling signal rather than merely a proxy for KV reuse or token count. It also contributes an explicit statement of the CGDC deployment model, a clean separation between session evaluation and decision making (so the value interface can evolve while the optimizer stays fixed), and a conditional loss bound for the asynchronous apply-and-revalidate execution model.

Real-world applications:

• Coding agents that compile, run tests, and query repositories between model turns (SWE-bench Lite, SWE-verified, Terminal-Bench, and Claude Code scenarios in the paper).
• Production LLM serving platforms that must multiplex many concurrent agent sessions over a fixed GPU fleet while tools run on separate CPU infrastructure.
• Browser-automation and workflow agents (AutomationBench in the paper) where many short tool calls interleave with model requests.
• Mixed fleets that also serve ordinary chatbot traffic, since the authors state TAO is compatible with non-agent chatbot scenarios.

Industry relevance. The baselines compared against — vLLM's round-robin and cache-aware schedulers, SGLang-style prefix-locality policy, NVIDIA Dynamo, and ThunderAgent — are the schedulers practitioners actually deploy. Improvements of up to 26.8% in latency and 29.5% in decode throughput, achieved with scheduling overhead reported below 0.12% of runtime, come without changing the inference engine; the authors describe the approach as composing with engine-level KV-cache management and demonstrate integration with TokenCake. The reported regressions (Terminal-Bench decode throughput and AutomationBench at 128 workers) also matter for practitioners deciding when not to use it.

Future Directions

• Understanding and closing the reported regressions. TAO does worse than baselines on decode throughput for Terminal-Bench (−8.8% at 16 workers, −6.2% at 64) and on all three metrics for AutomationBench at 128 workers (−7.0% latency, −504.0% TTFT, −11.2% decode). The paper does not explain these cases, leaving a clear open question about which workload regimes the value triplet misprices.
• Strengthening the readiness predictor. Mean absolute error of 0.410 s against a 1-second scheduling horizon, with only 82.9% of errors within 0.5 s, suggests room to improve the historical estimator or exploit richer tool interfaces; the authors note their finite-sample margin is a heuristic that provides no calibrated per-call coverage.
• Scaling the exact dynamic program. The authors note that rounding can exclude plans feasible before discretization and describe fallbacks (BFD, longest-prefix application, plan discarding) when execution outruns the snapshot. How the approach behaves at larger node counts and with wider capacity dimensions than those evaluated is not reported.
• Extensibility the authors mention but leave open. The paper says TAO "can be generalized and extended to different workloads such as multi-agent orchestration and Looped Agents," and discusses the relationship to program-aware scheduling in Agentix and workflow-aware cache management — but does not evaluate these settings.

Target Audience

Most valuable to systems researchers and engineers working on LLM inference serving, agent infrastructure, and cluster scheduling — particularly those deploying vLLM, SGLang, or Dynamo-style stacks who want cache locality and tool-call timing handled in one admission decision. It is also relevant to applied optimization researchers interested in online knapsack-style scheduling with lexicographic objectives, and to graduate students in distributed systems or ML systems looking for a well-specified problem formulation plus an end-to-end evaluation on real agent benchmarks. Readers without background in KV-cache management and serving metrics such as TTFT will need to consult the cited systems work first.

Authors’ 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.