SAGA: Workflow-Atomic Scheduling for AI Agent Inference on GPU Clusters

SAGA is a distributed scheduler that improves the efficiency of compound AI agent workflows on GPU clusters by shifting from request-level to program-level scheduling, which preserves intermediate KV cache states and reduces task completion time by 1.64x despite a tradeoff in peak throughput.

The Gist

Imagine you are running a busy kitchen (a GPU cluster) where chefs (AI agents) are trying to cook complex, multi-course meals (AI tasks).

Currently, most kitchen managers (existing schedulers like vLLM) treat every single order as a completely separate, one-off event. If a chef needs to chop vegetables, then wait for the oven to preheat, then chop more vegetables, the manager forces the chef to:

1. Cook the first batch.
2. Throw away all the chopped veggies and the dirty knives (the KV cache) because the chef is "waiting" for the oven.
3. When the oven is ready, the chef has to chop the exact same vegetables all over again from scratch.

This "start-over" cycle happens dozens of times per meal. It wastes massive amounts of time and space, making the kitchen 3 to 8 times slower than it needs to be.

SAGA is a new kitchen manager that changes the rules. Instead of looking at individual orders, SAGA looks at the entire recipe as one single unit. Here is how it works, using simple analogies:

1. The "Recipe Book" (Agent Execution Graphs)

Instead of guessing what the chef will do next, SAGA reads the recipe book (the Agent Execution Graph).

• The Problem: The chef stops to wait for the oven (a "tool call"). Old managers assume the chef is done and clear the counter.
• SAGA's Fix: SAGA knows the recipe says, "After the oven, we need to chop onions again." So, it tells the chef: "Keep the chopped onions and the knife on the counter. Don't wash them yet."
• The Result: When the oven is done, the chef picks up right where they left off. No re-chopping. SAGA predicts this so well that it performs almost as perfectly as a manager who could see the future (a theoretical "optimal" manager).

2. The "VIP Table" Strategy (Session-Affinity Batching)

Imagine a chef is working on a complex 10-course meal.

• The Problem: In the old system, if the chef gets busy, the manager might send the next step of the meal to a different chef at a different station. The new chef has to re-read the whole recipe and re-chop the veggies because they don't have the first chef's notes.
• SAGA's Fix: SAGA says, "This entire 10-course meal belongs to Chef A at Station 1." Even if Chef A is waiting for the oven, the next step is reserved for them. If Station 1 gets too crowded, SAGA might move the whole meal to a new station, but it brings the "notes" (the cache) with it so the new chef doesn't have to start over.
• The Result: The kitchen stays organized, and chefs don't waste time re-doing work.

3. The "Fairness" Rule (Agent Fair Share)

Imagine a restaurant with two types of customers:

• Customer A: Orders a simple burger (a short task).
• Customer B: Orders a massive, 50-course banquet (a long, complex agent task).
• The Problem: Old managers often prioritize the burger because it's quick to finish. The banquet customer waits forever, getting frustrated.
• SAGA's Fix: SAGA looks at the whole banquet. It realizes, "If we keep feeding the burger, the banquet will never finish." It ensures that the banquet gets enough attention to finish on time, even if it means the burger waits a little longer. It guarantees that everyone gets their full meal, not just the quick snacks.

The Trade-Off (The "Speed vs. Quality" Balance)

SAGA is incredibly fast at finishing individual complex meals (reducing the time to finish a task by 1.64 times). However, because it spends time organizing and keeping things ready for the next step, it can't churn out as many total meals per hour as a manager who just throws everything into a blender and ignores the recipe.

• The Paper's Claim: SAGA is about 30% slower at maximum raw volume (throughput) compared to the "churn-and-burn" style.
• Why it matters: The paper argues this is a good trade-off. Most AI agents are interactive (like a coding assistant or a browser bot) where users care about how fast the task finishes, not how many tasks the server can theoretically squeeze in.

Summary of Results

When tested on a real 64-GPU supercomputer:

• Speed: Tasks finished 1.64 times faster than the current best standard (vLLM with prefix caching).
• Memory: The kitchen used its counter space (GPU memory) 22% more efficiently, meaning it could handle more complex recipes without running out of space.
• Reliability: 99.2% of tasks finished within their promised time limits, even when the kitchen was chaotic and crowded.

In short, SAGA stops AI agents from throwing away their work every time they pause, ensuring they can pick up exactly where they left off, making complex AI tasks feel much snappier and more reliable.

Technical Summary

1. Problem Statement

Current GPU cluster schedulers (e.g., vLLM, SGLang) are optimized for single-shot inference, treating each Large Language Model (LLM) request as an independent unit. However, modern AI Agents operate via iterative "Thought-Action-Observation" loops, executing tens to hundreds of chained LLM calls per task, interspersed with external tool invocations (code execution, web browsing, database queries).

This creates a fundamental mismatch with existing scheduling systems, leading to three critical inefficiencies:

1. KV Cache Regeneration Overhead: Current systems evict Key-Value (KV) caches during tool-call idle periods (which can last from milliseconds to minutes). When the agent resumes, the cache must be regenerated, wasting 38% of execution time and increasing latency by 3–8×.
2. Fragmented Memory Utilization: Independent scheduling leads to poor memory locality. GPU memory utilization averages only ~42% because caches are fragmented and evicted prematurely, despite the potential for high reuse.
3. Lack of Task-Level Fairness: Existing schedulers optimize for per-request latency or throughput. In multi-tenant environments, this causes "starvation" for agents with long, complex workflows, as short requests monopolize resources.

2. Methodology: The SAGA Architecture

SAGA proposes a paradigm shift from request-level to program-level (workflow-atomic) scheduling. It treats the entire agent workflow as a single schedulable unit rather than a sequence of independent requests. The system consists of three main layers:

A. Agent Execution Graphs (AEGs)

SAGA models agent workflows as directed graphs where nodes represent LLM inference steps and edges represent dependencies (including tool calls).

• Construction: AEGs are derived from framework hints (e.g., LangChain callbacks) or inferred via pattern matching if hints are unavailable.
• Utility: The graph allows the scheduler to predict future steps, transition probabilities, and the likelihood of cache reuse across tool-call boundaries.

B. Three Core Mechanisms

1. Workflow-Aware KV Cache Management (WA-LRU):

   • Instead of standard LRU, SAGA uses a Workflow-Aware LRU (WA-LRU) eviction policy.
   • It calculates an eviction score based on: (1) Recency, (2) Predicted Reuse Probability (derived from the AEG), and (3) Cache Size.
   • Tool-Call-Aware TTL: The system dynamically adjusts the Time-To-Live (TTL) of cached KV states based on the specific tool being called and current memory pressure. This prevents premature eviction during long tool waits.
   • Speculative Prefetching: If the next step is predictable, SAGA prefetches the required KV cache during the tool execution phase to overlap I/O and computation.

2. Session-Affinity Batching with Work Stealing:

   • Affinity Routing: Requests belonging to the same agent session are routed to the same GPU worker to maximize cache hits.
   • Work Stealing: To prevent load imbalance (stragglers), a global coordinator performs randomized work stealing. If a worker is idle while others are overloaded, it steals tasks. Crucially, the system migrates the KV cache state (using mechanisms like Llumnix) to preserve locality, ensuring the stolen task doesn't lose its cached context.

3. Agent Fair Share (AFS):

   • A new fairness metric based on Task Completion Time rather than request count.
   • It uses Lyapunov drift analysis to guarantee that tenants with urgent, long-running tasks receive proportional resources, preventing starvation. The system dynamically adjusts scheduling priorities based on the "urgency" (remaining work vs. deadline) of each agent task.

3. Key Contributions

• Workflow-Atomic Abstraction: The first system to treat the entire agent program as the first-class scheduling unit, explicitly surfacing workflow structure (DAGs) to the scheduler.
• Theoretical Bounds:
  • Proved that WA-LRU achieves a competitive ratio within 1.31× of Bélády's optimal offline policy (the theoretical best possible cache eviction) on production traces.
  • Provided a formal proof for Agent Fair Share (AFS) using Lyapunov drift, guaranteeing bounded deviation in task completion times under bounded demand heterogeneity.
• Integrated System Design: Combines distributed scheduling, predictive caching, and fairness guarantees into a unified framework, addressing gaps left by prior work (e.g., vLLM's prefix caching, KVFlow's eviction, Llumnix's migration).

4. Experimental Results

Evaluated on a 64-GPU cluster (8 nodes × 8 NVIDIA A100s) using SWE-bench (coding agents) and WebArena (browser agents) workloads.

• Latency Reduction:
  • SAGA reduces Task Completion Time (TCT) by 1.64× (geometric mean) compared to the state-of-the-art vLLM v0.15.1 with Automatic Prefix Caching (APC).
  • Compared to vanilla vLLM, the improvement is 3.01×.
  • Specifically, on SWE-bench, TCT dropped from ~612s (vLLM) to 203s (SAGA).
• Memory Efficiency:
  • GPU memory utilization increased from ~42% (vLLM) to 71.3% (SAGA).
  • KV cache regeneration time was reduced from 38% of total time to just 8%.
• Fairness:
  • Achieved 99.2% SLO attainment (tasks meeting deadlines) under multi-tenant interference, compared to 67.3% for vLLM.
  • Eliminated the "long tail" latency for light tenants, ensuring consistent performance across heavy and light workloads.
• Trade-off:
  • SAGA optimizes for latency and task completion, resulting in a ~30% reduction in peak throughput compared to throughput-optimized batch scheduling. The authors argue this is an acceptable trade-off for latency-sensitive interactive agent deployments.

5. Significance

• Bridging the Gap: SAGA demonstrates that the inefficiencies in compound AI systems are not inherent to the models but are artifacts of request-level abstractions. By exposing workflow structure, schedulers can approach offline-optimal efficiency.
• Empirical Validation of Theory: The paper provides the first empirical competitive ratio bound (1.31×) for online cache management in agent workflows, validating that workflow-awareness is sufficient to nearly match optimal offline policies.
• Production Readiness: The system is implemented as an extension to vLLM and supports major agent frameworks (LangChain, AutoGen) without requiring code changes to the agents themselves.
• Future Direction: It establishes a new baseline for "Compound AI Systems," suggesting that future infrastructure must evolve from serving static requests to orchestrating dynamic, stateful workflows.

In summary, SAGA solves the critical latency and efficiency bottlenecks of AI agents by shifting the scheduling unit from the individual request to the entire workflow, leveraging predictive caching and fairness-aware resource allocation to achieve near-optimal performance on GPU clusters.