On optimizing AI costs

AI inference costs can blow up in your face (including agentic coding tools) in many ways which come from lack of utilization of prompt caches, bloated context windows due to tools etc. A lot of these issues didn’t have any in-app solutions until the problems became so glaringly obvious that providers have started offering solutions to some of these problems.

LLMs convert user input to tokens which later form embeddings. The embeddings (8kB per token) can get quite large in size in the KV cache (0.5MB-1MB per token per layer) and the structure of the system is that you need to compute all the embeddings to produce the final output. To avoid frequent re-computation providers usually cache the computation in the GPUs. Labs price this ~10x less than uncached tokens. The cache is short-lived, ~5-10 min. This is why it’s so important to hit the cache. The article explains how I do it in >Alfred, and where it’s headed in the future.

The simple path where you just enter an AI chat application some messages and it returns a response, without any special capabilities works fine. Things get really interesting when you expose ability to call tools that depend on user defining or integrating them, especially when the app allows you to connect lots of MCPs.

So what affects the cache computation in modern day LLM applications? Normally the moving parts are- the system prompt, the tools fed available to the model, and the user messages and model’s own responses. A prefix is the order in which we establish a hierarchy among these parts, this is how the provider chooses how these parts become model context. Different providers establish hierarchy in different ways, and it usually evolves with time.

Since there’s no universal prefix, only whatever the provider rendered before running the prefill. We didn’t have tools when AI was introduced, they were introduced in March 2023 by OpenAI, 5 months after ChatGPT released with GPT 3.5, ~3 years after GPT 3.

Anthropic documents the hierarchy as tools first, followed by the system prompt and the messages, tool wrappers compile around the system prompt. You can set the boundary with cache_control breakpoints.

Say, you had a connected Github integration in your AI application which worked with this policy, and in the middle of the chat thread, you go and connect Sentry, and send a message again, the cache won’t be hit this time. To make matters worse, this might pollute the context with a bunch of unnecessary tools which you might not need in the entirety of the request, which would be highly detrimental for the quality of the responses generated.

Explicit cache control

Some model providers have since provided cache control options to continue from specific prefixes in your prompts. As per Anthropic docs, placing cache_control directly on individual content blocks allows for fine-grained control over exactly what gets cached.

Let’s say your AI application changes its system prompt based on memories from user facts (which are subject to change during the conversation): then even if your tools don’t change, you’d incur prefilling the context (everything from the changed block onwards) sometimes multiple times in a single chat thread.

If you place breakpoints in the code across multiple places, the provider would ensure we don’t invalidate the cache breakpoints preceding the one that missed. The order still stays provider defined; ie, tools - sys - history. So, although this solves a lot of problems, it doesn’t solve the problem applications like Alfred face almost everytime.

Deferred tool loading

In the previous example, we cannot optimize for changes in the tools array. This situation is more common than the labs anticipated in the beginning, particularly with the inception of MCP. Applications should be able to connect to hundreds of integrations, but doing so plainly would bloat the context window so much there’ll be no room for meaningful work.

Since this particular use case is becoming more common; some providers have already established support for this situation. Anthropic, for instance, has introduced defer_loading attribute on the tools parameter. From their docs:

“Every tool in the tools array, including user-defined tools, accepts optional properties that control how the tool is loaded, who can call it, and how its inputs are validated. These properties compose: you can set defer_loading and cache_control and strict on the same tool.”

As I have realised few times in the past, we (can) control what to submit; but can’t (as of yet) control how they’ll be processed by the llm (become model context). So, unless we self host our own models, customising templates will remain a distant dream.

Since most coding harnesses are closed source, we’d never know if they have all these features but if they are serious about cost optimisations, it won’t be long before they go leaps and bounds in optimisations. All this converges towards a balance between in-situ on demand retrieval of data and speed of response. This leads us to a feature where the model will emit outputs which will be stored as artifacts in the harness and it can run a script to query those outputs without we manually passing it to the model. And none of this needs to be confined to coding.