TAO: Tool-Aware Joint Scheduling for LLM-Based Agentic Serving
Abstract
TAO improves scheduling for LLM agent services by considering tool readiness and using dynamic programming to place sessions, reducing latency and increasing throughput.
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.
Community
DeepInfer recently open-sourced how DeepSeek manages the KV cache with 3FS and balances prefill and decode loads. Since some of those load balancers remain unreleased, they propose their own effective router, TAO—a tool-aware joint scheduling framework designed for bursty agentic workloads under tight KV cache budgets.
Targeting agentic workloads that alternate between clustered GPUs and distributed CPU tools, TAO improves scheduling in three steps:
- Tool-Boundary Interruption: Pauses sessions during external tool calls to free up GPU capacity and places them into a shared waiting pool.
- Triplet Session Valuation: Scores each session–node pair without manual weight tuning by prioritizing time-decayed KV-cache credit, near-term tool-return readiness (predicted from runtime hints and historical tool durations), and token footprint.
- Joint DP Placement: Replaces greedy per-session routing with an asynchronous dynamic program (under 0.12% runtime overhead) that jointly selects and places waiting sessions under GPU capacity limits.
Tested on DeepSeek-V4 Flash, Qwen3-32B, and GLM-4.7, TAO cuts latency by up to 26.8%, lowers mean TTFT by up to 36.3%, and boosts decode throughput by up to 29.5% over leading baselines.
Really practical framing: treating tool-return readiness as a first-class scheduling signal next to KV-cache reuse makes a lot of sense for agent workloads, and the structured-history duration predictor reaching a 0.410 s mean absolute error over 10,967 tool calls is a nice detail. The ablation in Table 4 stood out to me: replacing DP resumption with best-fit-decreasing drops decode throughput from 1417.97 to 1233.72 tok/s, while the snapshot analysis shows the DP improving the objective in 6 of 23 multi-candidate snapshots and matching short-first in the rest. Is most of that gap coming from those few snapshots, or from the asynchronous revalidation around the plan? I also noticed the AutomationBench run at 128 workers, where mean TTFT rises to 19.1 s versus 3.2 s for Dynamo; what drives that case?
Thanks for the sharp questions! Here is a brief breakdown of both points:
- The 6/23 Multi-Candidate Snapshots (Table 4 vs. Section 4.3)
- Reconstructable baseline: The 6/23 (26.1%) metric reflects an offline, static replay where both schedulers face the exact same snapshot. Most scenarios are not same even in replay, it is therefore a conservative number for example showcase.
- The "Butterfly Effect": In online multi-turn serving (Table 4), once a single scheduling decision diverges, subsequent states branch immediately and are no longer identical. A single superior DP decision early on preserves more KV cache (e.g., +23.2% credit in Fig. 10) and avoids memory fragmentation, preventing downstream thrashing and queue clogging.
- Asynchronous revalidation impact: Revalidation overhead is minimal (<0.12% runtime, <0.4% plan discard/partial rate), so the throughput gap in Table 4 primarily stems from BFD's compounding suboptimality across turns rather than revalidation loss.
- Mean TTFT on AutomationBench at 128 Workers
- Client-visible TTFT vs. admission wait: The reported client TTFT includes the time spent waiting in the scheduler's pool. Dynamo admits requests aggressively into the GPU engine, yielding a lower initial TTFT but triggering massive KV-cache thrashing and preemption under saturation.
- Intentional admission gating: TAO purposefully holds returning sessions in the waiting pool until sufficient node capacity is available, shielding in-flight sessions' KV cache from eviction.
- The global trade-off: While upfront TTFT is higher, admitted sessions finish far faster without recomputation. As shown in Table 3, TAO cuts overall task latency from 1417.2 s to 563.1 s (2.5x faster) and doubles decode throughput (1013.2 vs. 482.4 tok/s).
Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper