Skip to content

Astraea: A State-Aware Scheduling Engine for LLM-Powered Agents

🕒 Published (v1): 2025-12-16 06:55 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Astraea is a lifecycle-centric inference serving engine for LLM agentic workflows that replaces per-segment scheduling with global Job Completion Time (JCT) optimization. It combines a Stateful-MLFQ scheduler (informed by historical request state and predicted future behavior) with an adaptive KV cache manager, reducing average JCT by up to 25.5% over vLLM-FCFS on heterogeneous agentic workloads.

Problem

Existing LLM serving systems (e.g., vLLM) treat each segment of a multi-stage agentic workflow as an independent task, optimizing local per-segment latency while ignoring API call durations and chained dependencies. This causes Head-of-Line blocking and suboptimal global JCT, because API calls can dominate total latency (seconds vs. milliseconds for decode) yet are invisible to the scheduler.

Method

Astraea builds on vLLM and adds three coordinated components:

  1. Service Time Predictor: Offline-profiled prefill latency model + segment-level generation-length oracle + category-mean API latency statistics to annotate each segment with estimated compute and I/O times before scheduling.

  2. Stateful-MLFQ Scheduler: A hierarchical preemptive multi-level feedback queue using token cost (not wall-clock time slices) as migration thresholds. New requests enter the highest-priority queue; compute-heavy requests are demoted, I/O-triggered requests are promoted. Within a queue, segments are sorted by an enhanced HRRN score: (accumulated_parent_wait + segment_proc_time) / segment_proc_time, which approximates SRPT under equal wait but prevents starvation as wait time grows. Preemption is at segment (not iteration) granularity.

  3. Adaptive KV Cache Manager: Monitors GPU memory against a high-watermark threshold. Below threshold: Preserve policy (keep KV cache on GPU). Above threshold: selects among Preserve/Discard/Swap by minimizing a memory-waste cost model (W_preserve = T_api · C_self · M, W_discard = T_recompute · C_batch · M, W_swap = 2·T_swap · C_batch · M).

The scheduling problem is proven NP-hard (reduction from 1|prec|ΣCj), motivating the heuristic design.

Key Contributions

  • Formal NP-hardness proof for the LLM agent scheduling problem under precedence constraints.
  • Stateful-MLFQ: first scheduler to unify a request's full historical compute/I/O profile with future predictions for global JCT optimization.
  • Adaptive KV cache manager that dynamically picks Preserve/Discard/Swap per-request based on quantified memory-waste cost, not static policy.
  • Prototype implemented on vLLM with modular, non-interfering integration.

Results

  • GPT-J (6B, single A100), 30% GPU memory, QPS=5: Astraea JCT 122.88s — 19.1% below Infercept, 25.5% below vLLM-FCFS.
  • Vicuna-13B (2×A100), 30% GPU memory, QPS=5: Astraea JCT 140.02s vs. vLLM-SJF 171.16s (−18.2%) and vLLM-LAS 153.96s (−9.1%).
  • Vicuna-13B, 90% GPU memory, QPS=5: Astraea JCT 132.53s — 13.4% below Infercept, 16.0% below vLLM-FCFS.
  • Stability under load (30% memory): vLLM-SJF JCT degrades 51.9% from low→high QPS; Astraea degrades only 16.2% (3.2× more stable).
  • Ablation (Stateful-MLFQ alone, no adaptive cache, GPT-J, 30%, QPS=3): JCT 137.66s — 19.3% below SJF, 2.3% below FCFS, confirming scheduler-only benefit.
  • Scheduler overhead: 0.0006% of average JCT — negligible.
  • Baselines: vLLM-FCFS, vLLM-SJF, vLLM-LAS, Infercept (FCFS).

Limitations

  • API latency prediction uses category-mean statistics; highly variable or novel API categories will degrade scheduling quality.
  • Generation-length prediction uses a segment-level oracle (ground-truth token counts) — not a deployable predictor; real-world accuracy would lower gains.
  • Evaluated on only two model sizes (6B, 13B) and one hardware configuration (A100 80GB); behavior at 70B+ or with tensor parallelism across many GPUs is untested.
  • Preemption at segment granularity (not token level) limits scheduling responsiveness within a running segment.
  • The NP-hardness proof motivates heuristics but no approximation ratio is established for Stateful-MLFQ.

Relevance to Harnesses / Meta-Harnesses

Astraea directly addresses the runtime orchestration layer that underlies any multi-step agentic harness: it shows that naively wrapping LLM calls in a loop (ReAct, LangChain, AgentGraph) creates scheduling pathologies at the serving infrastructure level that no harness-level logic can fix. For meta-harness designers, this is a systems-level complement — the harness controls what agents do and when they call tools, while Astraea controls how those calls are prioritized and whether intermediate state (KV cache) survives the wait. The adaptive KV cache manager is especially relevant to harnesses that issue long-latency tool calls (image generation, search), where the choice of Preserve vs. Discard vs. Swap directly affects whether a harness iteration resumes at low cost or pays full recomputation. This work validates that lifecycle-aware state tracking — the same principle behind stateful meta-harnesses — must be applied at the serving layer as well as the orchestration layer.