Transformer Inference and Context Engineering

The context window is the most important resource in LLM inference, and most engineers manage it poorly. A 128K-token window does not mean you can throw 128K tokens at a model and expect good results. Long contexts are expensive, slow, and often counterproductive.

This article covers how to engineer contexts efficiently: prompt structure, token budgets, KV cache behavior, and why long contexts frequently hurt rather than help.

Token Budget Accounting

Before writing any prompt, know your token math.

pythonCopy
import tiktoken

enc = tiktoken.encoding_for_model("gpt-4o")

def count_tokens(text: str) -> int:
    return len(enc.encode(text))

def budget_report(system_prompt: str, user_message: str,
                  retrieved_context: list[str], max_output: int = 512):
    system_tokens = count_tokens(system_prompt)
    user_tokens = count_tokens(user_message)
    context_tokens = sum(count_tokens(c) for c in retrieved_context)
    total_input = system_tokens + user_tokens + context_tokens
    total_budget = 128_000  # gpt-4o

    print(f"System prompt:      {system_tokens:>6} tokens")
    print(f"User message:       {user_tokens:>6} tokens")
    print(f"Retrieved context:  {context_tokens:>6} tokens ({len(retrieved_context)} chunks)")
    print(f"Total input:        {total_input:>6} tokens")
    print(f"Reserved for output:{max_output:>6} tokens")
    print(f"Available:          {total_budget - total_input - max_output:>6} tokens")
    print(f"Budget utilization: {total_input / total_budget:.1%}")
    return total_input

Run this before deploying. If p95 prompt length exceeds 80% of the context window, you have a production risk.

Prompt Structure for Efficient Retrieval

Where you put information in the prompt matters. LLMs suffer from "lost in the middle": they attend better to content near the beginning and end of the context.

pythonCopy
def structure_rag_prompt(question: str, chunks: list[str],
                          system_base: str) -> list[dict]:
    # Put the most relevant chunk first (not last)
    # Put instructions at the end of the system prompt (recency bias helps)
    context_text = "\n\n---\n\n".join(
        f"[Source {i+1}]\n{chunk}" for i, chunk in enumerate(chunks)
    )
    return [
        {
            "role": "system",
            "content": (
                f"{system_base}\n\n"
                f"Retrieved context:\n{context_text}\n\n"
                "Answer the user's question using only the retrieved context. "
                "If the context does not contain the answer, say so explicitly."
            ),
        },
        {"role": "user", "content": question},
    ]

Practical ordering rule: most relevant chunk first, question at the end, instructions immediately before the question. This places high-signal content at both ends of the "lost in the middle" range.

KV Cache Behavior and Cost

The KV cache stores key-value tensors for all input tokens. For cloud APIs (like OpenAI), cached tokens cost less. For self-hosted models, cache memory limits how many concurrent users you can serve.

pythonCopy
def estimate_kv_cache_memory(
    n_layers: int,
    d_model: int,
    n_heads: int,
    seq_len: int,
    batch_size: int,
    dtype_bytes: int = 2,  # fp16
) -> float:
    """Returns KV cache memory in GB for a given config."""
    d_head = d_model // n_heads
    kv_per_layer = 2 * batch_size * n_heads * seq_len * d_head * dtype_bytes
    total = n_layers * kv_per_layer
    return total / (1024 ** 3)

# Llama 3.1 8B: 32 layers, d_model=4096, 32 heads
# At seq_len=8192, batch=8:
mem = estimate_kv_cache_memory(32, 4096, 32, 8192, 8)
print(f"KV cache: {mem:.1f} GB")  # ~32 GB - equals the model weights

This is why long contexts are expensive: the KV cache memory scales linearly with sequence length. At 128K tokens, the KV cache for a 70B model exceeds 100GB. Most production deployments cap max_model_len significantly below the theoretical maximum.

Context Window Utilization Patterns

Not all long contexts are created equal:

Good use of long context: A legal contract where any clause might be relevant. You must include the full document because you do not know which clause the question targets.

Bad use of long context: A RAG system that dumps the top-20 retrieved chunks into the context. Most chunks are noise. Retrieve fewer, better chunks instead.

pythonCopy
# Anti-pattern: retrieve many, dump all
chunks = retriever.query(question, top_k=20)
context = "\n".join(c.text for c in chunks)  # 10K tokens of mixed relevance

# Better: retrieve few, rerank, keep the best
from sentence_transformers import CrossEncoder
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

chunks = retriever.query(question, top_k=20)
scores = reranker.predict([(question, c.text) for c in chunks])
top_chunks = [chunks[i] for i in sorted(range(len(scores)),
                                         key=lambda x: -scores[x])[:5]]
context = "\n\n".join(c.text for c in top_chunks)  # 2.5K tokens of high relevance

Reranking reduces context size by 4x and improves answer quality simultaneously.

Why Long Contexts Do Not Always Help

Research (Liu et al., 2023 "Lost in the Middle") showed that LLMs consistently fail to retrieve information from the middle of long contexts, even when it is explicitly present. The quality of "needle in a haystack" retrieval degrades linearly as the context grows past ~8K tokens for most models.

Practical implication: if your task requires finding a specific fact in a large document, prefer chunked retrieval over full-document context. Full-document context is best when the task requires synthesizing information across the entire document, not finding individual facts.

Common Mistakes

Not accounting for conversation history in token budgets. Multi-turn conversations accumulate context. A conversation with 10 turns can easily exceed 20K tokens. Implement a sliding window or summarization strategy.

Padding retrieved context to fill the context window. More context does not help and frequently hurts. Retrieve 3–5 highly relevant chunks, not 20 marginally relevant ones.

Ignoring prompt caching economics. OpenAI and Anthropic cache repeated prompt prefixes. If your system prompt is 2K tokens and sent on every request, you pay full price every time unless you use their caching APIs. Prefix caching can reduce costs by 50–80% for chat applications.

Where to Go Next

See also: [transformers-attention-in-practice], [rag-foundations-retrieval-quality], [inference-optimization-performance]