Caching In Agentic Llm Pipelines

Reproducing Variance: Caching in Agentic LLM Pipelines | AI21

Skip to Main Menu

Skip to Main Content

Skip to Footer

Back to Blog

-->

Back to Blog

TL ; DR When running experiments on agentic workflows – especially agents with multiple parallel LLM calls, caching breaks because LLMs are non-deterministic. What these workflows need is a caching mechanism that supports two contradictory behaviors: reproducibility and variance . We designed a cache key that encodes each LLM call’s position inside the pipeline to make it resistant to changes in the pipeline’s execution order. The result is cleaner and broader experiments: you can change one prompt and use cached results for the rest; A/B-test any component against identical upstream outputs; and run best-of-N inference without multiple branches collapsing to the same answer.

Caching in a non-deterministic world

You don’t need us to tell you that agents are difficult to evaluate due to some core characteristics. They’re:

Complex (many processes, often parallel)

Dynamic (their course of action is determined on the fly by LLMs)

Long-running (sessions can take minutes or hours)

Expensive

Caching, then, becomes go-to recourse for managing evaluation experiments.

But here’s the catch (sorry):

Caching was built for a deterministic world – same input, same output, every time. As non-deterministic entities, LLMs don’t play by these rules. You call the same model with the same prompt and you actually expect a different result. So much so that you can even control the degree of difference via temperature.

And yet most of our engineering toolbox carries that deterministic assumption baked in. When you apply these constructs to LLMs without rethinking them, things break in ways that don’t look broken at first.

That’s exactly what happened in our research lab. Our algorithm developers – who were running agent evals – started reporting that “the cache is broken.” When our team of software engineers took a look, we understood that what was missing was a way to guarantee both reproducibility and variance in the cache for every run of an agent eval. So we built it. This blog details that path from problem to solution.

Reproducing variance: Defining the problem

When we’re building and testing agents in our lab, we work with Maestro to optimize them at runtime. At inference time, it can utilize best-of-N sampling : generating N independent attempts at a task and returning the best one, so that the accuracy of the agent output isn’t degraded by the noisiness of any single LLM call. In addition to improving agent accuracy, that approach has also been instrumental for generating indicators of agent performance, such as variance estimations or the metric _Success@K_ , as demonstrated in our previous post on how we topped the SWE-bench leaderboard .

So it was concerning when we started getting messages like “the cache is broken!!” Looking closer, it became clear that the problem was more nuanced than that: Researchers were expecting reproducibility from the cache – but because of LLMs’ non-deterministic nature, running the same agent again recomputes – but doesn’t reproduce – results. So we were seeing things like:

When re-running an experiment, expecting unchanged prompts to be cached (cache hits) and the experiment to run quickly – but then getting different results (cache misses).

Expecting variance (cache misses) when running an experiment for success@K, but then getting the same result (cache hits).

What our researchers needed was a mechanism that supports two different cache behaviors – reproducibility and variance – within the same pipeline, selectively deciding when to reuse results and when to recompute. If we could build that, we could unlock how to reproduce a best-of-N run.

Diagramming contradictory cache behavior

To reproduce a best-of-N agent run, it would be ideal if we could just toggle on / off reproducibility.

Reproducibility can be shown in the simplest example of caching behavior is two LLM calls with the same prompt – we can achieve that by turning on the cache toggle.

Another characteristic of an agentic workflow that we get for free here is LLM call composability – when the output of one LLM call is the input of another:

Best-of-N sampling, on the other hand, necessitates variance even within a single run. Here, every new Generate branch must miss; if two branches return the same result, it can’t really be said that the agent tried N times.

So we’d have to toggle off reproducibility to get this variance – but then we still haven’t reproduced our best-of-N run. Our task is to solve the seemingly paradoxical question: How do you reproduce a best-of-N run when needed? If we could pull it off, it would look something like this:

Designing the cache key

At this point, it was clear that a simple on/off toggle wouldn’t be enough to achieve that delicate combination of reproducing variance. We would need to use the LLM API settings/parameters to design a cache key – the string attached to the LLM call which determines which calls share results and which don’t – that could fulfill those contradictory conditions.

The cache key is significant because, besides a workable on/off toggle for reproducibility, it’s the only lever we have left available to affect the cache behavior.

Two core assumptions shaped our solution design:

We can attach a cache key to each LLM call

A cache hit requires both the prompt and the cache key to match

For this process, we’ll refer to the “write me a poem” LLM call as the “generation” step, and the “Is a valid Haiku” as the “validation” step.

Attempt 1: Using the step name for the cache key

If we set it up so that the cache key just uses the step name, all three Generate branches can share the same key. So when reproducing, Generate₂ hits Generate₁’s cached result and so on. Three branches, one answer. This indicates that our cache key is not specific enough: We need a way to distinguish between the three different Generate calls.

Attempt 2: Step name + counter

One natural fix for making it easy to distinguish between calls: number each call as it finishes. Run 1 works perfectly – each generation branch receives a number. In the next run we can reuse the gen_0, gen_1, and gen_2 keys to reproduce this variance

Where this breaks is when we compose LLM calls on top of each other. Consider the following pipeline:...