Behind the Stack, Ep 6 - How to Speed up the Inference of AI Agents

Optimizing AI Agent Inference Latency and Cost with Prefix Caching 

Introduction: The Latency Problem in AI Agents

AI agents are transforming everything from customer support to autonomous workflows. But under the hood, most AI agent architectures suffer from one major problem: growing latency and cost at scale.

Each reasoning step adds more tokens to the input, and because most systems (especially API-based or naive self-hosted setups) resend the entire prompt history on every call, AI agents end up:

• Repeating compute from earlier steps
• Wasting GPU cycles
• Scaling inference cost and latency quadratically

Even modern caching APIs fall short - they don’t cache intermediate thoughts, tool results, or agent memory effectively.

The Solution? Prefix Caching for AI Agents

Prefix caching is a feature available in advanced self-hosted AI inference engines like vLLM, SGLang, and TGI. It allows your AI agents to reuse previously computed context efficiently, cutting down latency and cost - without changing the logic of your agent.

In this post, you’ll learn:

• Why traditional AI agent chains are inefficient
• How prefix caching works inside LLM inference
• When and how to deploy it
• What infrastructure patterns support it best

If you're running multi-step AI agents, this is a foundational optimization strategy.

Understanding Agent Chains in LLMs

A typical AI agent combines:

• A user prompt
• A large language model (LLM)
• Optional tools like APIs, search engines, or databases

As the agent runs, it appends reasoning steps, tool results, and responses to the prompt. For example:

User: What's the weather in Tokyo?

AI Agent: Let me check that.

Tool (weather API): Sunny and 27°C.

Assistant: It’s currently 27°C and sunny in Tokyo.

Every new reasoning step or tool call adds more tokens to the input. Each interaction includes the entire conversation history, all prior messages, internal thoughts, and tool results, sent again to the LLM.

This behavior compounds quickly. For an N-step agent, the token input size grows linearly, but the total compute and latency grow quadratically due to full-context reprocessing at each step.

Latency and Cost Breakdown

To understand the performance impact, consider the inference pipeline in a decoder-only LLM:

• Prefill phase: Model processes input tokens and builds key-value attention cache (KV cache)
• Decode phase: Model generates output tokens, attending to the KV cache

The prefill phase dominates for long prompts - often consuming over 90% of total inference time. When the prompt history grows with each step, prefill cost quickly becomes the bottleneck.

Let’s model a 5-step agent chain with 200 initial tokens and 300 tokens added at each step:

• Step 1: The model processes 200 input tokens. Cumulative compute is 200 tokens.
• Step 2: With 300 more tokens added, the input grows to 500 tokens. Cumulative compute is now 700 tokens.
• Step 3: The input reaches 800 tokens. Cumulative compute: 1,500 tokens.
• Step 4: Input size is now 1,100 tokens. Cumulative compute rises to 2,600 tokens.
• Step 5: Final input size hits 1,400 tokens. Cumulative compute totals 4,000 tokens.

In total, 4,000 tokens are reprocessed across steps, much of it being redundant work from earlier steps. This illustrates how prompt growth drives up inference cost, especially during the prefill phase.

Why Traditional Caching Isn't Enough

API-level caching (like that offered by OpenAI or Anthropic) works for stateless queries, but not for stateful AI agents.

Limitations of traditional caching for agents:

• No memory of intermediate reasoning
• No reuse of tool responses
• No optimization of the prefill (KV cache) stage
• Each LLM call processes the entire token history again

In short: standard caching does little to solve the token bloat that AI agents accumulate during long chains of reasoning.

What Is Prefix Caching? A Game-Changer for AI Agent Performance

Prefix caching addresses the problem directly by retaining and reusing KV cache blocks from earlier steps. Modern inference engines like vLLM, SGLang, and TGI support this capability.

When the new input shares a prefix with a prior prompt, the model:

• Detects the reused prefix
• Skips recomputation for those tokens
• Uses cached attention keys and values directly

This results in:

• Constant-time reuse for shared prompt history
• Significantly reduced prefill time
• Lower overall latency per agent step

Practical Benefits: Why Prefix Caching Matters for AI Agents

In a standard setup:

Step 1: User prompt -> Process all tokens

Step 2: User + Reasoning + Tool -> Reprocess entire context

Step 3: User + Reasoning + Tool + More Reasoning -> Reprocess again

With prefix caching:

Step 1: User prompt -> KV cached

Step 2: Append Tool response -> Only new tokens processed

Step 3: Append More Reasoning -> KV reused again

Latency per step remains flat. Token processing scales linearly. This significantly reduces response times in long-running agent chains.

Deployment Patterns for Scalable AI Agent Systems 

To implement prefix caching effectively:

1. Use KV-aware inference backends like vLLM or SGLang
2. Maintain session affinity to ensure cached prefixes stay on the same GPU
3. Avoid prompt mutations that change earlier tokens
4. Load-balance by prefix similarity in distributed systems

Advanced setups may:

• Track KV state across replicas
• Partition workloads into “prefix-heavy” vs “prefix-light” requests
• Route high-cache sessions to the same replica to maximize reuse

When Should AI Teams Use Prefix Caching?

Prefix caching is most impactful in:

• AI agents using tools across multiple steps
• Long, multi-turn user conversations
• High-volume, self-hosted inference setups
• Custom agent frameworks built on open LLMs

It’s less useful for:

• Single-shot queries
• Stateless RAG pipelines
• Embedding-only workloads

Conclusion: Make AI Agents Fast and Cost-Efficient 

AI agents are inherently stateful, which creates repeated work unless you optimize for it. Prefix caching gives your infrastructure a way to reuse prior context, avoid redundant compute, and unlock scalable AI agent performance.

It’s not just a nice-to-have - it’s a core optimization for production-grade agent systems.

If you're running or planning to deploy multi-step AI agents, integrating prefix caching is one of the most impactful things you can do to reduce:

• Latency
• GPU utilization
• Cloud inference costs