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

Paper

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

Original: Shuhuai Li et al., “TAO: Tool-Aware Joint Scheduling for LLM-Based Agentic Serving” (arXiv:2610.37230v1, preprint). The following account is based solely on the attached original text.

Introduction

Agents alternate between model inference and external tool execution. While a tool runs, that session’s GPU demand becomes suspendable; how this gap is used and which node the session is attached to when it returns determine response time. This paper proposes TAO, a framework that binds both problems into a single scheduling problem.

1. Problem: GPUs idle while tools run

1.1 Sessions and CGDC

Let us call the entire interaction of an agent a session. A session is a unit of work that alternates between model requests and tool calls. During a model request, the session occupies one GPU node; during tool execution, it waits for results outside the inference service.

The paper calls this arrangement CGDC (Clustered GPUs with Distributed CPUs). GPUs are clustered in nodes and share inference capacity, tools run on distributed CPUs outside them, and tool CPU/RAM is managed separately from the GPU admission budget. In other words, there is time that does not count against the GPU budget.

Roles are predefined as well. The scheduler records each session’s state and controls session admission, while the serving engine manages residency and eviction of physical KV cache (source: §2.1).

1.2 The expensive thing is “not being able to reuse the KV cache”

If a session’s next request can reuse the KV prefix produced by the previous turn, it skips prefill computation entirely. But whether that prefix is still alive is determined by “which node it was attached to” and “how much time has passed.” If a session returning after a tool finishes is placed on the same node, the cache is alive and it is fast; if it is sent to another node or pushed out by capacity pressure, prefill must be recomputed.

Conversely, evicting an active inference session requires immediate recomputation. When recomputation and restoration repeat, capacity is wasted on state that must be rebuilt (source: §2.1, §2.2).

1.3 interrupt and restart

At a tool boundary the scheduler can do two things.

• interrupt: release the model admission of a session running a tool and return node resources.
• restart: reassign a session whose tool has finished to an appropriate node and continue the existing interaction.

Both actions are just session state transitions, decoupled from the inference engine implementation. The paper’s starting point is that these two decisions ultimately reduce to the same session evaluation. Whether to interrupt is decided by ranking session values, and interrupted sessions enter the waiting pool together with newly arriving sessions. The paper handles joint evaluation in this waiting pool and the resume decision together (source: §2.2).

1.4 Two things existing schedulers miss

Existing serving systems mostly follow an “evaluate-then-decide” pattern. They evaluate node features at the session level, then decide where and when to serve.

Original Figure 1. The left panel shows the CGDC structure, and the right panel shows the two problems the paper targets.

• Problem 1 — evaluation misses tool-return readiness. Tool timing was already used before, but only as a proxy for KV reuse or token count. For example, predicting tool duration to set KV retention time, or discounting a session’s token demand over time while a tool runs. These values get mixed into fixed scoring rules or hand-picked thresholds.
• Problem 2 — decisions are made with per-session heuristics. The per-session scoring and routing rules of llm-d and SGLang allocate inference opportunities one session at a time. But sessions compete for the same node capacity, so admitting one session can remove the slot where a more valuable combination of other sessions would have fit. Labels in the right panel such as “Short but Low Cache,” “Urgent but Tool Running,” and “Session Still Waiting in Pool” point to this illusion (source: §1, §2.2).

2. A small example: while one session waits 12 seconds

To get the concept, follow a hypothetical coding session S. The numbers in this section are a hypothetical example for understanding and are not reported in the original.

Two things emerge from this example. One is that “whether it returns to a node with a live cache” determines response speed, and the other is that “if you know when the tool will finish,” you can raise S as a GPU placement candidate just before that and eliminate the wait. The paper turns these two into values called ‘cache credit’ and ‘return readiness.’

3. Method

3.1 Two-stage pipeline: evaluate → decide

TAO operates in two stages. First it attaches a value to each session–node pair in the waiting pool (evaluation), and based on those values a dynamic program selects and places waiting sessions all at once (decision). Because evaluation and decision are separated, changing the value function leaves the scheduling procedure intact (source: §3.1).

  flowchart LR
    OBS["Observation: remaining KV on node b1 and<br/>remaining tool time of session S"] --> VAL["Value computation: k=decayed cache credit, r=1 (returning soon), t=token count"]
    VAL --> DP["DP: stores the best value per capacity state and<br/>assigns S to b1"]
    DP --> CHK["Re-check against live state just before apply<br/>(discard the plan if the selected session's signature or eligibility changes)"]
    CHK --> RUN["b1 reuses S's prefix<br/>skips prefill computation, generates first token"]
    DP -. "insufficient remaining capacity on b1" .-> BFD["BFD fallback: place on b2<br/>prefill must be recomputed"]
    BFD --> CHK

The concrete numbers in this diagram, like the example in Section 2, are not reported in the original. Each box is a state change. The first box is the observations (remaining cached token count and remaining tool time), the second is the result of turning them into values, the third is the process of recovering a placement from the values, and the last is what happens in actual execution. The dotted line is the case where capacity is insufficient and it takes another path. The serving engine’s role (physical KV allocation, reuse, eviction) does not disappear. The scheduler only reads on top of it “where the cache remains and when it can be reclaimed” as values. What shrinks is prefill recomputation, which disappears through cache reuse, and return waiting, which disappears through readiness.

3.2 Three values attached to each session–node pair

For each session i and node b pair, the paper attaches three values (source: §3.2, Eq. (2)).

• k_{i,b} — cache credit of session i remaining on node b. It shrinks as the observation ages.
• r_i — return readiness. A 0/1 signal indicating whether session i is likely to request the GPU again soon because its tool is about to finish.
• t_i — token count of session i. Represents the accumulated context size.

The key design is that these three are compared in lexicographic order rather than mixed in a weighted sum. Cache credit is examined first; if equal, readiness; if still equal, token count. The order itself is fixed, and only the values change over time. The paper justifies this order for three reasons (source: §3.2). First, prefix reuse eliminates prefill recomputation outright, so cache credit comes first. Second, when decay makes credits similar, readiness distinguishes sessions that are “returning soon” from those “still far away.” Third, when credit and readiness tie, larger accumulated context is preferred, but the optimizer also accounts for the capacity cost that context occupies when determining the overall assignment set.

The assignment constraints are as follows (Eq. (1)). Session i can go to only one node ($\sum_b x_{i,b} \le 1$), cannot exceed node b’s remaining capacity C_b ($\sum_i w_i x_{i,b} \le C_b$), and ineligible pairs cannot be assigned. The goal is the assignment that lexicographically maximizes the sum of values over all pairs (Eq. (3)).

3.3 Credit decays with observation age

Using the observed cached token count k⁰ as-is leaves old observations as a frozen signal. The paper computes credit by multiplying by $2^{-\tau/h}$ according to the time τ elapsed since observation (source: §3.3, Eq. (4)).

The half-life h is set to 10 seconds, and the computed credit is rounded before being used in evaluation. The basis is measurement. In a DeepSeek-V4 Flash trace on two nodes each with 4×H100, the binned median of observed cached token retention declines clearly over time (original Figure 3). The median was drawn with a technical exponential fit over populated bins, and each point carries a bootstrap 95% confidence interval. Giving large weight to old observations leads to wrong decisions, so decay is necessary.

The half-life is only a scheduler parameter and is distinct from the time constant fitted from the trace. Physical retention and eviction of the actual cache remain the serving engine’s responsibility.

3.4 Readiness tells “returning soon” as 0/1

Readiness identifies sessions that will need the GPU again soon. If the tool interface exposes progress or a completion deadline, it uses that first. Claude Code’s Monitor delivers background task events, and OpenCode reports termination after a Bash timeout (source: §3.4). Without direct timing evidence, it estimates the remaining time from historical tool duration records. The estimation order is per-tool history → global history → 5-second default. It collects the sample count, mean, sample variance, and observed range of completed calls, builds a conservative estimate with a finite-sample margin, subtracts the time already elapsed, and clips at zero (Eq. (5), Appendix A).

When the predicted remaining time enters the horizon H=1 second, r=1. The decision uses three things together (Eq. (6)). If tool completion has already been observed, r=1. If the predicted deadline has passed but completion has not yet been observed, it grants only a temporary probe opportunity instead of remaining “ready” indefinitely. Otherwise, it decides only if the predicted timestamp was updated within 30 seconds. Observed completion takes precedence over prediction.

This signal is tied to the scheduling period (1 second). A session with readiness 1 is raised as a GPU placement candidate and can emit its first token as soon as the tool finishes.

3.5 Select and place at once with dynamic programming

Once all values are assigned, the decision problem becomes “dividing items among multiple knapsacks.” Each node b is a knapsack with remaining capacity C_b, and each session i is an item with admission cost w_i and per-node value v_{i,b}.

The problem with existing approaches is that they solved this per session. Putting sessions one by one into the node that looks good in sequence blocks the slot where a more valuable combination could fit later.

The paper splits capacity into fixed units (quantum Δ), then, adding candidates one at a time, stores “the best value obtainable at each capacity state with the j candidates so far” with a DP that finds the exact optimum over a frozen, discretized snapshot (source: §3.5, Eq. (7)).

Here $\bar{w}_i = \lceil w_i/\Delta\rceil$, $\bar{C}_b = \lfloor C_b/\Delta\rfloor$, and $u_{j,b}$ is the term holding the value and count when candidate j is put on node b. Each candidate tries both “skip” and “assign to b,” so all feasible solutions are covered. However, because of rounding, a plan that was feasible before discretization may be excluded.

The computation runs in an isolated asynchronous worker, and just before applying, it re-checks against the live state (whether the selected session’s signature and eligibility are unchanged), then applies only the longest plan prefix that current capacity allows. If the signature changes, it discards the entire plan. The preparation and pool readiness cost is less than 0.12% of total runtime in the 4-replica sweep (source: §3.5, Appendix I).

Per-node cache observations are decayed with O(nm) preprocessing and folded into values, and nodes with the same cache can be grouped into g ≤ m capacity dimensions (physical KV blocks are not moved). When the state budget is exhausted, it falls back to best-fit-decreasing (BFD) placement. The admission cost w_i is the memory demand of the session’s total token count minus shared token capacity, and C_b is updated at each scheduling point as the node’s remaining capacity (Appendix B.1).

What the DP actually chooses differently

Original Figure 10. Each program is distinguished by a different hue, translucent spans are that program’s admission weight, and dark overlays are retained cache credit. The horizontal axis is admission weight (thousand tokens), and the height of each node’s horizontal bar is that node’s limit.

• (a) short-first affinity greedy: node 1 leaves 13.6k cache with 8.4k free, node 2 has 20.1k free.
• (b) dynamic programming: it leaves 16.8k cache (+3.2k, +23.2%), node 1’s free space shrinks to 1.4k and node 2’s free space grows to 27.1k.

While maintaining the same total safe token retention, the DP preserves more cache credit. The exact values are 16,791 tokens versus 13,624 tokens (source: Appendix H).

Original Figure 11. Each session’s width is its admission weight (including the 512-token reserve), and the number beside it is the program ID and cache credit.

• (a) short-first affinity greedy: admits 6 programs and defers program 177, a long-context one. Cache 10,397 tokens.
• (b) dynamic programming: admits 7 including program 177. Cache 14,949 tokens.

That is, the heuristic fills slots with short programs and pushes out a valuable session, while the DP admits that session too (source: Appendix H).

4. Results

4.1 Metrics and comparison conditions

There are three metrics. Latency is the average end-to-end task time (including model inference + tool execution), TTFT is the average time to first token, and decode throughput is generated tokens per second (TPS). In the component analysis, it additionally measures prefix-cache hit rate, latency from tool completion to first token, and pause wait. Pause wait is the time an interrupted session waited, including tool-call latency.

There are four comparison targets: vLLM’s round-robin (RR) and cache-aware (CA) schedulers, the KV-aware scheduler Dynamo, and ThunderAgent, which uses BFD placement with token-count scheduling. Deployments are DeepSeek-V4 Flash on 2 nodes × (4 H100 per node), Qwen3-32B on 4 nodes × (2 H100 per node), and GLM-4.7 on 1 node × 8×H100 (TP8). The workload is SWE-bench Lite (SWE-Lite) using the mini-SWE-agent harness; DeepSeek additionally runs Terminal-Bench 2.0 with Claude Code, and Qwen additionally runs SWE-verified (source: §4.1).

The paper reports up to 26.8% latency reduction, up to 36.3% TTFT reduction, and up to 29.5% decode throughput increase across 12 model–workload–concurrency settings, with the comparison baseline being the best evaluated baseline per metric (source: abstract).

4.2 Main results: where it won and where it lost

In the DeepSeek-V4 Flash SWE-Lite results, each cell is the improvement over the best competing value for that metric (source: original Table 1).

The improvement grows with concurrency. At 64 workers, TAO’s latency is 185.410 s, far better than the best competing 253.228 s, TTFT is 0.224 s versus 0.351 s, and decode throughput is 2,021.523 versus 1,561.307 tok/s. At 16 workers the gaps are narrower: 113.264 s versus 116.011 s, 0.135 s versus 0.146 s, and 934.967 versus 889.590 tok/s. That is, the more sessions compete for node capacity, the more pronounced the effect of “optimal placement at once” and “valuation reflecting cache and readiness” (source: §4.2).

Looking at the other settings’ improvements together makes the overall picture clearer. The values below transcribe the Rel. Imprv. rows of Table 1, and negative numbers are cells where TAO fell behind (source: original Table 1).

In 5 of the 36 cells in Table 1, TAO falls behind the best competing value. DeepSeek Terminal-Bench leads latency and TTFT at all three concurrency levels but trails decode at 16 workers (−8.8%) and 64 workers (−6.2%). Qwen3-32B SWE-Lite trails at 96 workers decode (−5.1%) and 160 workers latency (−0.2%), and SWE-verified trails only at 128 workers latency (−2.5%). The conclusion is closer to “improves several metrics together in most conditions” than “wins every condition on every metric.”

The large-model scaling experiment, GLM-4.7 (1 node 8×H100, TP8, 300 tasks), recorded 15.5% latency reduction, 14.1% TTFT reduction, and 20.5% decode increase at SWE-Lite 160 workers, and across all SWE-verified 500 tasks recorded 16.2–18.9% latency reduction, 15.7–38.4% client TTFT reduction, and 22.8–24.7% decode increase. The comparison target is ThunderAgent, the SOTA on each metric (source: original Table 2, §4.2).

4.3 Serving dynamics: don’t cram everything in at once

Original Figure 4. This is a running-request-count trace when GLM-4.7 runs 300 SWE-Lite tasks at 160 workers (TP8), with each run aligned to zero at the first client request. TAO finished the workload at 60.0 minutes, ThunderAgent at 70.9 minutes, vLLM CA at 164.3 minutes, vLLM RR at 176.8 minutes, and Dynamo at 189.4 minutes. TAO keeps the number of concurrently running requests low throughout the middle of the run while finishing earlier. That is, it stably admits additional sessions while keeping the concurrent request count low. Corresponding traces for 96 and 128 workers (TAO 53.9 and 59.2 minutes, ThunderAgent 60.1 and 66.2 minutes) are in Appendix D (source: §4.2, Appendix D).

4.4 Which parts actually contributed

Component ablation. The setting is SWE-Lite, DeepSeek-V4 Flash, 2 nodes × 4×H100, 32 workers (source: original Table 4, §4.4).

The largest drop is the last row. Reverting the DP to simple BFD placement drops decode from 1417.97 to 1233.72 tok/s and increases mean TTFT from 0.154 to 0.195 s. This means that good valuation alone is not enough if placement is myopic; performance falls sharply, supporting the design rationale that “evaluation” and “decision” play different roles. Removing readiness lowers decode to 1363.06 and increases tool end→first token to 0.227 s. The effect of removing cache value is relatively small (1398.35). It is also notable that the ablation effects are not monotonic. The “without cache value and readiness” row (1373.51) shows higher decode than the “without readiness only” row (1363.06). However, compared with the gap to the last DP→BFD row, this irregularity is not large.

Cache-aware evaluation. In a fixed workload replay of 218 tasks, the DP placement was held identical to the ThunderAgent control while only the valuation was changed. Cache token ordering raised prefix-cache hit rate by 0.19%p and reduced mean pause wait by 4.5%. Adding cache credit decay on top raised hit rate by 0.31%p and decode throughput by 1.68%. Figure 5 shows the reduction in physical prefill, TTFT, and pause wait together (source: §4.3, original Figure 5).

Tool readiness. In paired 32-worker SWE-Lite runs, enabling readiness raises decode throughput by 4.03%, reduces latency from tool completion to first token by 10.88%, and lowers mean pause wait from 1.067 s to 0.972 s. The left CDF in Figure 6 shows the distribution from tool completion to first output token, and the right shows the prediction error distribution. Across 10,967 completed tool predictions, mean absolute error is 0.410 s and 82.9% are within 0.5 s (the return-latency CDF has 10,481 and 10,673 observations, and the prediction panel has 10,967 predictions) (source: §4.3, original Figure 6).

Concurrent placement. Of 23 replay snapshots with multiple candidates, 6 (26.1%) improved the lexicographic objective and the other 17 matched the short-first heuristic. Across all 35 recoverable cases including single-candidate cases, 6 (17.1%) improved and 29 (82.9%) matched. That is, on this criterion the DP was never worse, and it gained on snapshots with severe contention (source: §4.3, Appendix H).

4.5 Mixed workloads and supplementary evaluations

Mixed workload (source: original Table 3). In a Terminal-Bench + SWE-Lite mix (389 tasks), latency (+12.2/+23.2/+13.8%) and decode (+12.2/+8.6/+4.4%) improved at all of 32, 64, and 128 workers, and only 32-worker TTFT trailed at −5.5%. AutomationBench (800 tasks) improved all three metrics at 32 and 64 workers. At 128 workers, however, all three trailed: latency 563.114 s versus 523.676 s (−7.0%), TTFT 19.123 s versus 3.166 s (−504.0%), and decode 1,013.157 versus 1,140.862 tok/s (−11.2%). The direction of gain and loss varies by workload.

Low-concurrency fallback (source: Appendix E.1, Table 5). DeepSeek-V4 Flash was placed on a single TP8 node and swept from 4 to 64 workers. At 4–8 workers, TAO’s mean task latency is within 2.3% of the best baseline, decode throughput stays at or above 92% of the best, and mean TTFT is 15–20 ms higher than the best. At 32 and 64 workers, ThunderAgent falls far behind (at 64 workers, latency 475.033 s versus 226.268 s, TTFT 9.332 s versus 0.423 s). This confirms that TAO does not collapse even when scheduling headroom shrinks.

InferenceX replay (source: Appendix F, Table 7, Figure 9). DeepSeek serves a fixed corpus consisting of 393 AgentX traces, a 1-million-token context limit, and an 1,800-second profiling window on 2 nodes × 4×H100. TAO ranks first in decode throughput at all of 16, 32, and 64 workers (+4.2/+1.8/+10.1%), and TTFT ranks first only at 16 workers; at 32 and 64 workers the cache-aware scheduler is lower. In Figure 9, TAO’s prefix-cache hit rate ranks first at all worker counts.

Composition with another KV management engine (source: Appendix E.2, Table 6). Layering TAO on TokenCake increases decode throughput from 785.159 to 886.772 tok/s and from 1,292.584 to 1,347.373 tok/s at 16 and 32 workers, and at 64 workers improves latency from 509.018 to 274.892 s, TTFT from 1.590 to 0.593 s, and decode from 790.491 to 1,414.547 tok/s. However, the authors also report that in this run the program’s pause, resume, and migration never occurred, and they acknowledge that attributing the effect to the mechanism requires a separate run.

Sensitivity to long tool latency (source: Appendix G, Table 8). SWE-Lite sessions instantiated with MiniSWE (6,995 model requests, 6,942 tool records) were given per-session fixed additional latencies of 1–5 s, 2–10 s, and 4–20 s. Compared with readiness off, TAO reduced makespan by 20.72/16.42/19.28% and mean session latency by 21.57/20.57/17.91%. The longer the latency, the larger the fraction of cases where readiness changes the actual placement decision: 1.65% → 3.94% → 7.93%. Compared with an Oracle that provides the known true time, history-based prediction cuts mean absolute error by 99.53–99.69% (0.025–0.030 s) and approaches Oracle in runtime performance.

5. Limitations and applicability judgment

5.1 Limitations the paper itself states or that follow structurally

• Optimality holds only over the discretized snapshot. The DP is exact on the frozen, discretized snapshot, but rounding can exclude plans that were feasible before discretization (source: §3.5).
• It does not guarantee global online optimality. The loss from asynchronous application is presented as a conditional bound in terms of value fluctuation ε_{s,t}, opportunities A_{s,t} arising from feasibility changes, and execution loss T_{s,t} (Eq. (12)). $$\max_{x \in X^\Delta_t} U_t(x) - U_t(y_t) \le A_{s,t} + 2\epsilon_{s,t} + T_{s,t}$$ The comparison point of this bound is not the planning moment but the fresh snapshot optimum of the same state observed at application time. It therefore does not guarantee end-to-end latency or regret against a globally optimal online policy. The 0/1 jump in readiness does not automatically fit the exponential-decay assumption, so a separate jump allowance is needed (source: Appendix B.2, Appendix C).
• Readiness prediction is not a calibrated probability. It is only a scheduling heuristic with a finite-sample margin and does not provide calibrated coverage for individual calls. If neither per-tool nor global history exists, it relies on the 5-second default (source: Appendix A).
• Resource boundary assumption. It assumes that CPU and RAM for external tools are managed separately outside the GPU admission budget. Situations requiring additional resource constraints or interaction-dependent values need an extension of this admission model (source: §2.1, §3.4 — the latter is the reviewer’s interpretation).
• Overhead accounting is partial. The recorded scheduler work time covers state preparation and completed pools, so failed probes and BFD executions are not in the time account. That ratio is therefore not a measure of CPU utilization or request latency overhead (source: Appendix I).

5.2 Points a reviewer would note

• Scalability was not validated on node count. The DP’s state space grows with the number of nodes and capacity quantization. Nodes with the same cache are grouped into g ≤ m capacity dimensions, and when the state budget is exhausted it falls back to BFD. The measured fallback rates are 5.24% for DeepSeek 64 workers and 3.06–8.98% for the Qwen sweep. However, the reported measurements cover deployments up to 4 nodes (TP2×4), and deployments with far more nodes were not validated (source: §3.5, Appendix I, Table 9).
• Readiness depends entirely on signal quality. It works best in environments where the tool interface exposes progress and deadlines or where historical statistics are valid. For tools without such signals it relies on the 5-second default, and the accuracy guarantee there is limited (source: §3.4, Appendix A).
• There is no starvation-avoidance mechanism in the original. The third coordinate, token count priority, applies only when both cache credit and readiness tie. That a session with a larger token count gains placement priority in that case is shown by the admission of session 177 in Figure 11, but no starvation-prevention mechanism is specified. The fact that the snapshot is re-solved at each scheduling point appears to be a mitigating factor (source: §3.2; the mitigating factor is the reviewer’s interpretation).
• Coupling with the engine is still at the design stage. The composition with TokenCake worked, but pause, resume, and migration were not triggered, so whether that mechanism actually contributed is unconfirmed (source: Appendix E.2).

5.3 Applicability judgment

This design is particularly suited to agent serving with frequent tool calls, many sessions competing for node capacity, and predictable tool durations. SWE-bench-style coding agents, multi-agent orchestration, and Looped Agent are examples. The paper states that it is also compatible with chatbot load.

Conversely, in workloads with few tools or where prediction is difficult without tool signals, the benefit of valuation shrinks. Indeed, at AutomationBench 128 workers all three metrics trailed. In deployments with very many nodes, DP cost and benefit must be validated separately.

Figures in this post are taken from the original arXiv:2610.37230 (CC BY 4.0). Only size and format were changed.