Why this matters

AI applications have performance characteristics that differ from traditional web services. A single LLM call takes 500ms to 5 seconds. Embedding a document corpus takes minutes. Loading a large context into a prompt has a measurable cost in both tokens and latency. Without systematic profiling, engineers end up optimizing the wrong thing — rewriting application logic that takes 10ms while ignoring the 3s provider call that dominates wall clock time.

This lesson teaches you where to look, how to measure accurately, and which optimizations actually move the needle.

Core concepts

Profiling LLM-heavy code

The first rule of optimization: measure before you change anything. In AI systems, most time is spent in three places:

1. Provider API calls (network latency + model inference time)
2. Embedding and vector operations
3. Document loading and preprocessing

Use time.monotonic() for wall-clock measurements and Python's cProfile for CPU profiling:

import cProfile
import pstats
import io
import time


def profile_pipeline(fn, *args, **kwargs):
    pr = cProfile.Profile()
    pr.enable()
    start = time.monotonic()
    result = fn(*args, **kwargs)
    elapsed = time.monotonic() - start
    pr.disable()

    stream = io.StringIO()
    ps = pstats.Stats(pr, stream=stream).sort_stats("cumulative")
    ps.print_stats(20)  # top 20 functions by cumulative time

    print(f"Total wall clock: {elapsed:.2f}s")
    print(stream.getvalue())
    return result

For async code, use structured timing around await points:

import asyncio
import time
from contextlib import asynccontextmanager


@asynccontextmanager
async def timed(label: str):
    start = time.monotonic()
    try:
        yield
    finally:
        elapsed_ms = int((time.monotonic() - start) * 1000)
        print(f"{label}: {elapsed_ms}ms")


async def instrumented_pipeline(query: str, client, retriever):
    async with timed("retrieval"):
        chunks = await retriever.search(query, top_k=5)

    async with timed("prompt_assembly"):
        prompt = assemble_prompt(query, chunks)

    async with timed("llm_call"):
        response = await client.generate(prompt)

    return response

This reveals immediately whether retrieval or generation dominates. Most engineers are surprised the first time they profile: retrieval often takes longer than generation.

Caching strategies

Caching LLM responses is the highest-leverage performance optimization available. An LLM call that costs 2s and $0.002 becomes a cache hit that costs 0.1ms.

Exact-match cache with TTL:

import time
import hashlib
import json
from typing import Optional


class TTLCache:
    def __init__(self, max_size: int = 1000, ttl_seconds: float = 3600) -> None:
        self._store: dict[str, tuple[dict, float]] = {}
        self._max_size = max_size
        self._ttl = ttl_seconds

    def _key(self, payload: dict) -> str:
        serialized = json.dumps(payload, sort_keys=True)
        return hashlib.sha256(serialized.encode()).hexdigest()

    def get(self, payload: dict) -> Optional[dict]:
        key = self._key(payload)
        if key not in self._store:
            return None
        value, expires_at = self._store[key]
        if time.monotonic() > expires_at:
            del self._store[key]
            return None
        return value

    def set(self, payload: dict, result: dict) -> None:
        if len(self._store) >= self._max_size:
            # Evict the oldest entry (simple LRU approximation)
            oldest_key = next(iter(self._store))
            del self._store[oldest_key]
        key = self._key(payload)
        self._store[key] = (result, time.monotonic() + self._ttl)


cache = TTLCache(max_size=500, ttl_seconds=1800)


async def cached_generate(client, payload: dict) -> dict:
    cached = cache.get(payload)
    if cached is not None:
        return {**cached, "cached": True}
    result = await client.generate(payload)
    cache.set(payload, result)
    return result

For production, use Redis instead of an in-process dict so cache survives service restarts and is shared across instances.

Memory management with large contexts

Python's garbage collector handles most memory management, but large context windows create specific pressure points:

Token budget management: Track tokens explicitly rather than discovering context limit errors at runtime.

def estimate_tokens(text: str) -> int:
    # rough approximation: 1 token ≈ 4 characters for English text
    return len(text) // 4


def fit_chunks_to_budget(
    chunks: list[str],
    max_tokens: int,
    reserve_for_output: int = 1024,
) -> list[str]:
    available = max_tokens - reserve_for_output
    selected = []
    used = 0
    for chunk in chunks:
        chunk_tokens = estimate_tokens(chunk)
        if used + chunk_tokens > available:
            break
        selected.append(chunk)
        used += chunk_tokens
    return selected

Explicit cleanup for large intermediate objects: When processing large document collections, delete large intermediate objects explicitly to help the garbage collector:

def process_large_corpus(doc_paths: list[Path]) -> list[dict]:
    results = []
    for path in doc_paths:
        raw_text = path.read_text()  # potentially large
        chunks = split_into_chunks(raw_text)
        del raw_text  # release the full text immediately
        for chunk in chunks:
            results.append(embed_chunk(chunk))
        del chunks  # release the chunk list
    return results

Streaming to avoid materializing large lists:

def stream_embeddings(chunks: Iterator[str], client) -> Iterator[dict]:
    for chunk in chunks:
        embedding = client.embed(chunk)
        yield {"text": chunk, "embedding": embedding}
        # chunk and embedding are released after yield

Concurrent processing patterns

Beyond async network calls, CPU-bound processing (document parsing, tokenization, score computation) can be parallelized with ProcessPoolExecutor:

from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path


def parse_document(path: Path) -> dict:
    # CPU-intensive: OCR, PDF parsing, tokenization
    text = extract_text(path)
    return {"path": str(path), "text": text, "tokens": len(text) // 4}


def parallel_parse(paths: list[Path], max_workers: int = 4) -> list[dict]:
    results = []
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(parse_document, p): p for p in paths}
        for future in as_completed(futures):
            try:
                results.append(future.result())
            except Exception as exc:
                print(f"Parse failed for {futures[future]}: {exc}")
    return results

Use ProcessPoolExecutor for CPU-bound work (parsing, tokenization, scoring). Use asyncio.gather for I/O-bound work (API calls, database queries). Do not mix them carelessly — running async code inside a ProcessPoolExecutor worker requires starting a new event loop per process.

Working example

A profiled, cached, memory-efficient document pipeline:

import asyncio
import time
import logging
from pathlib import Path
from typing import Iterator

logger = logging.getLogger(__name__)


class InstrumentedPipeline:
    def __init__(self, client, cache: TTLCache) -> None:
        self._client = client
        self._cache = cache
        self._timings: dict[str, float] = {}

    def _time(self, label: str):
        start = time.monotonic()
        return lambda: self._timings.__setitem__(label, time.monotonic() - start)

    async def run(self, doc_path: Path, query: str) -> dict:
        stop = self._time("doc_load")
        text = doc_path.read_text()
        stop()

        stop = self._time("chunking")
        chunks = list(split_into_chunks(text))
        del text
        stop()

        stop = self._time("retrieval")
        relevant = await self._retrieve(query, chunks)
        del chunks
        stop()

        stop = self._time("generation")
        result = await self._generate(query, relevant)
        stop()

        logger.info("pipeline_timings", extra=self._timings)
        return result

    async def _generate(self, query: str, context: list[str]) -> dict:
        payload = {"prompt": assemble_prompt(query, context), "model": "gpt-4o-mini"}
        return await cached_generate(self._client, self._cache, payload)

Common mistakes

Profiling in development with small inputs. A document that loads in 10ms in dev loads in 400ms in production with a cold file system. Profile with realistic data sizes.

Caching at the wrong layer. Caching prompt assembly output (which is nearly free) adds complexity without benefit. Cache expensive operations: provider calls, embedding generation, reranker inference.

Materializing entire corpora in memory. embeddings = [embed(c) for c in all_chunks] on 1M chunks requires gigabytes of RAM. Stream embeddings to storage as you generate them.

Not measuring cache hit rate. A cache that is never hit adds latency (cache key generation, serialization) without benefit. Instrument cache hits and misses.

Try it yourself

1. Add @asynccontextmanager-based timing to a pipeline you own. Log the timings as a structured event. Find which stage dominates.
2. Implement a TTLCache class. Verify that it returns cached values within the TTL and misses after expiry. Add hit/miss counters and log them.
3. Take a pipeline that processes documents serially. Rewrite the CPU-bound parsing stage to use ProcessPoolExecutor. Measure the speedup on a corpus of 50 documents.