Algorithms for ML Engineers

Suppose you ask an internal engineering assistant:

Why did the auth tests fail?

Its permission filter returns four authorized text chunks (small passages) in a local fixture. Only one mentions an expired test setup, so finding it means inspecting all four. That is a reasonable place to start. What changes when the same search has 50,000 candidates?

That scan works. The question is how much work it repeats as the corpus grows. An algorithm is a recipe for completing a task; complexity describes how its work or memory grows as its input grows. We'll keep the four-row search in view while counting the work a larger corpus would repeat.

What grows when the corpus grows

Before comparing lookups, let's be clear about what each one receives. The SQL permission filter, optionally reinforced by row-level security (per-row access rules), has already returned only auth project chunks permitted for this request. Four symbols name the inputs that later code will grow with:

A permission filter can shrink n, and a wider embedding increases d. Returning two citations controls k, but it doesn't skip scoring every remaining candidate. That last distinction will matter when we replace a full sort with a heap.

Count a complete scan

Let's start with a retriever we can inspect. It scores a chunk by counting question words that appear in its text. The four rows below stay with us throughout the chapter. Before you run the code, make a prediction: the question matches one row, but how many rows will the function inspect?

chunks = [
    {"id": "auth-fixture", "text": "auth tests fail when fixture expires"},
    {"id": "ci-timeout", "text": "ci job timeout after slow integration test"},
    {"id": "cli-flag", "text": "readme documents old command line flag"},
    {"id": "deploy-runbook", "text": "rollback service after canary error"},
]

def lexical_scan(question: str, allowed_chunks: list[dict]) -> tuple[list[tuple[int, str]], int]:
    terms = set(question.lower().split())
    matches = []
    checked = 0
    for chunk in allowed_chunks:
        checked += 1
        overlap = len(terms & set(chunk["text"].lower().split()))
        if overlap:
            matches.append((overlap, chunk["id"]))
    return sorted(matches, reverse=True), checked

matches, checked = lexical_scan("auth tests fixture", chunks)
print("chunks_checked", checked)
print("best_chunk", matches[0][1])

chunks_checked 4
best_chunk auth-fixture

chunks_checked isn't elapsed time. It makes one repeated action visible: inspecting a candidate chunk. Four authorized chunks mean four inspections; forty thousand mean forty thousand. We can now see exactly what grows when this scan handles a larger corpus.

Real text work also depends on chunk length. A configured chunk-size limit lets us use candidate count as a useful first model. Without that bound, words read per chunk become another growing input. A counter exposes repeated work; it doesn't measure how quickly a machine performs it.

Find that repeated operation before reaching for an index.

Big-O names the growth shape

Big-O gives an asymptotic upper bound: how a cost can grow as an input gets large, ignoring fixed multipliers and small-input details. The label alone doesn't say whether you're describing a worst case, an average case, or an amortized sequence of operations, so name both the operation and the kind of bound.^{[1]Reference 1big-O notationhttps://xlinux.nist.gov/dads/HTML/bigOnotation.html}

Under the bounded-chunk simplification, this scan visits every authorized chunk. Its candidate-inspection work is tight: Θ(n), which also means O(n). Read Θ as "theta": it says the count grows linearly, not merely that linear growth is an upper limit.

Watch the count while repeating the same four rows. Repetition here changes workload size, not the meaning or quality of the documents:

base_chunks = [
    {"id": "auth-fixture", "text": "auth fixture expired"},
    {"id": "ci-timeout", "text": "ci timeout"},
    {"id": "cli-flag", "text": "old cli flag"},
    {"id": "deploy-runbook", "text": "canary rollback"},
]

def count_checks(question: str, allowed_chunks: list[dict]) -> int:
    terms = set(question.split())
    checks = 0
    for chunk in allowed_chunks:
        checks += 1
        _ = terms & set(chunk["text"].split())
    return checks

for copies in (1, 2, 4, 8):
    corpus = base_chunks * copies
    print(f"n={len(corpus):>2} checks={count_checks('auth fixture', corpus):>2}")

n= 4 checks= 4
n= 8 checks= 8
n=16 checks=16
n=32 checks=32

Doubling n doubles the repeated checks. Building a set of question terms still costs something, but that overhead doesn't grow once per candidate. Big-O helps us keep our attention on the part that can become expensive at scale.

It would be misleading to call the whole assistant "O(n)." Keep the claim attached to one operation: the candidate-inspection part of this bounded-length scan is Θ(n). The SQL filter, result ordering, network calls, and answer generation each have their own costs.

There's another operation inside the function: it sorts its r matching chunks before returning them. Here r is the number of matches, and r <= n. The scan is Θ(n), while sorting is O(r log r). If every chunk matches, the full function can be O(n log n). Looking at the steps separately keeps the scan from hiding an expensive sort.

Choose a bound that says what you mean

The growth class tells us how work scales. These three labels answer a different question: what kind of bound are we describing?

Θ is a tight growth class: both an upper and a lower bound hold, up to fixed multipliers for sufficiently large inputs.^{[2]Reference 2Thetahttps://xlinux.nist.gov/dads/HTML/theta.html} The pair loop later performs exactly n * (n - 1) / 2 comparisons, so its loop count is Θ(n²); writing O(n²) alone states only the upper bound.

O is still useful when you promise only an upper bound, but it isn't automatically a worst-case label. Average-case analysis depends on an input distribution, expected-case analysis averages randomness or hash behavior, and amortized analysis accounts for operation history.

A correct label still doesn't choose the lookup. The scan we counted grows with every authorized chunk. When does the data let a cheaper search win?

Pick a search that matches the data

A cheaper lookup wins only when you already have the structure that search needs. Finding an exact key in an unsorted list, finding one in a sorted list, and finding similar meaning are different jobs. The Data Structures for AI lesson matched those jobs to containers. Here we'll compare the work they require. An approximate-nearest-neighbor (ANN) index changes which candidates we inspect, so measure its recall alongside latency. In the graph row, V counts nodes and E counts links:

For 50,000 sorted IDs, binary search needs at most 16 interval decisions because $2^{16} = 65{,}536$. Sorting those IDs on every request would erase the benefit, so sorted order or a hash index must be built and maintained outside the hot path.

Binary search also can't rank embeddings. Semantic similarity doesn't provide one scalar ordering where nearby vectors are guaranteed to sit next to each other. That is a different input and a different index problem.

A hash set is the right tool when you already have the exact ID. It can't answer "which stored prompts start like this one?" That is a prefix problem. The Data Structures lesson built a character trie for token, tool, and top. This chapter counts that work, then compresses the trie into a structure serving engines actually use.

Binary search halves a sorted range

The four authorized chunk IDs are already sorted alphabetically. cli-flag happens to sit at the midpoint, so this search gets a lucky one-step hit. On a miss, the ordering still tells us which half could contain the target.

lo is the first possible index and hi is one past the last. The remaining range is therefore [lo, hi). Integer division // 2 chooses a midpoint; moving a boundary discards the part that can't contain the target.

sorted_ids = ["auth-fixture", "ci-timeout", "cli-flag", "deploy-runbook"]

def binary_search(sorted_ids: list[str], target: str) -> tuple[int | None, int]:
    lo, hi = 0, len(sorted_ids)
    steps = 0
    while lo < hi:
        mid = (lo + hi) // 2
        steps += 1
        if sorted_ids[mid] == target:
            return mid, steps
        if sorted_ids[mid] < target:
            lo = mid + 1
        else:
            hi = mid
    return None, steps

index, steps = binary_search(sorted_ids, "cli-flag")
print("index", index)
print("steps", steps)
print("full_scan_checks", len(sorted_ids))

index 2
steps 1
full_scan_checks 4

The search finishes in one midpoint check here. An early-stopping linear search would find cli-flag on its third check; the earlier ranking scan would inspect all four. Binary search's worst-case guarantee comes from halving the remaining interval, so its step count stays O(log n) even when n is tens of thousands. A midpoint check may perform both == and <; we're counting visited midpoints, not individual comparison operators.

Exact counts still need care. This half-open implementation can require 17 midpoint checks for 65,536 items because the final one-element interval needs its own comparison. The list has to stay sorted, and a new ID inserted in the middle of a Python list still costs O(n).

Prefix match, tries, and radix trees

Imagine several requests that begin with the same system prompt and tool schema. If a lookup rereads that shared prefix for every stored request, it pays for those tokens again and again. A trie stores one symbol per edge: lookup of a prefix of length p walks p edges, then enumeration costs whatever the returned completions require. Reaching the prefix doesn't get more expensive just because unrelated keys were added elsewhere.

Long unique suffixes waste nodes. The words token and tokenize share token; tokenize then adds a chain i -> z -> e that never branches. A radix tree, also called a Patricia trie or compressed prefix tree, collapses each non-branching run into one edge that holds a sequence. Lookup still follows exact symbols, but it compares a whole shared run at once.

The names are easy to mix up, so keep the jobs in view. A radix tree answers a longest exact prefix. An embedding index finds approximate nearby vectors. Radix sort is a third, separate algorithm.

View larger

Three stored prompts share the system and tool tokens, then split. The compressed edge stores system + tools once. A new request that starts the same way follows that edge before checking its remaining tokens. Keep that picture in mind while the next program counts the work with token IDs.

Count the difference on toy token IDs. 1, 2 is the shared system-and-tools prefix. 3 is user A. 4, 11 is user B's whole remaining path.

prompts = [
    [1, 2, 3, 10],
    [1, 2, 4, 11],
    [1, 2, 3, 12],
]
query = [1, 2, 3, 13]

def naive_prefix(query: list[int], prompts: list[list[int]]) -> tuple[int, int]:
    comparisons = 0
    best = 0
    for prompt in prompts:
        shared = 0
        for left, right in zip(query, prompt):
            comparisons += 1
            if left != right:
                break
            shared += 1
        best = max(best, shared)
    return best, comparisons

# first token -> (edge tokens, child map)
tree = {
    1: (
        (1, 2),
        {
            3: ((3,), {10: ((10,), {}), 12: ((12,), {})}),
            4: ((4, 11), {}),
        },
    )
}

def radix_prefix(query: list[int], node: dict) -> tuple[int, int]:
    comparisons = 0
    index = 0
    current = node
    while index < len(query):
        child = current.get(query[index])
        if child is None:
            break
        edge, nxt = child
        shared = 0
        mismatch = False
        for offset in range(min(len(query) - index, len(edge))):
            comparisons += 1
            if query[index + offset] != edge[offset]:
                mismatch = True
                break
            shared += 1
        index += shared
        if mismatch or shared < len(edge):
            break
        current = nxt
    return index, comparisons

naive_len, naive_cmp = naive_prefix(query, prompts)
radix_len, radix_cmp = radix_prefix(query, tree)
print("shared_tokens", naive_len, radix_len)
print("naive_comparisons", naive_cmp)
print("radix_comparisons", radix_cmp)

shared_tokens 3 3
naive_comparisons 11
radix_comparisons 3

The naive scan rereads 1, 2 for every stored prompt, so three keys take 11 token comparisons. The radix walk reads the shared edge once, follows the 3 branch, and stops after 3 explicit token comparisons. Both find the same longest prefix of length 3. The tree also makes three dictionary lookups, including the failed lookup for 13; those aren't included in the comparison counter.

Indexing query[index + offset] avoids copying the remaining query at each edge. Writing query[index:] inside the loop would create a new list every time. Many short edges could then cause quadratic copying even while the token-comparison counter stayed linear.

Add a thousand other prompts that share only [1, 2] and the scan keeps paying per key. With expected constant-time child lookup, the tree's descent takes O(p + 1) work for p matched tokens. The extra 1 accounts for detecting a mismatch, including an immediate miss. Enumeration after the match is a separate cost, because returning more completions means doing more work.

For token sequences $x$ and $y$, the longest common prefix is the greatest number $p$ of initial tokens that agree. In this notation, $x_{1:j}$ means the first $j$ tokens; $j=0$ includes the empty prefix:

Matching tokens can make previously computed model state reusable, but only when the model and other inputs that produced that state also match. The unmatched suffix still needs processing; a prefix hit doesn't make the whole request free.

Serving uses the same idea on token IDs. A prefix cache can reuse the prompt-processing work, called prefill, for a compatible cached prefix. SGLang's RadixAttention organizes token prefixes in a radix tree.^{[3]Reference 3SGLang: Efficient Execution of Structured Language Model Programshttps://arxiv.org/abs/2312.07104} vLLM instead hashes full blocks of cached model state. Its keys include preceding blocks and additional identity such as adapter or multimodal inputs, and block boundaries limit the reused length. Exact token agreement is necessary, but doesn't by itself establish safe reuse across models or cache scopes.^{[4]Reference 4Automatic Prefix Cachinghttps://docs.vllm.ai/en/stable/design/prefix_caching/}

You don't need to implement the cache here. Keep its algorithm test: if the question is longest exact prefix, count work in p. If the question is nearby meaning, you're back in the $n \times d$ scoring column.

Vector scoring adds another input

We have answered one question: which prompts share an exact prefix? To find text with nearby meaning, semantic retrieval commonly represents the question and each chunk as embedding vectors, then computes a score. The NumPy and tensor shapes lesson treated a vector as coordinates and a dot product as paired multiplies plus adds. Here we'll count those multiplies in plain Python so the growing input stays visible.

Each exact dot-product score inspects d coordinates. When both vectors have Euclidean norm one, meaning $\sqrt{\sum_i x_i^2}=1$, their dot product equals cosine similarity. This describes their magnitude, not their number of coordinates. The lab vectors below aren't normalized that way, so treat this as a raw dot product. Scoring all n allowed chunks performs work proportional to n * d.

Before running the code, predict the counter rather than the winner: four vectors with three coordinates each require how many coordinate multiplications, even when two query coordinates are zero?

chunk_ids = ["auth-fixture", "ci-timeout", "cli-flag", "deploy-runbook"]
chunk_vectors = [
    [1.0, 0.2, 0.1],
    [0.0, 1.0, 0.2],
    [0.8, 0.1, 0.4],
    [0.1, 0.2, 1.0],
]
question = [1.0, 0.0, 0.0]

def dot(left: list[float], right: list[float]) -> tuple[float, int]:
    score = 0.0
    multiplies = 0
    for a, b in zip(left, right, strict=True):
        score += a * b
        multiplies += 1
    return score, multiplies

scores = []
coordinate_multiplies = 0
for vector in chunk_vectors:
    score, used = dot(question, vector)
    scores.append(score)
    coordinate_multiplies += used

best = max(range(len(scores)), key=lambda i: scores[i])
print("best_chunk", chunk_ids[best])
print("scores", [round(s, 1) for s in scores])
print("coordinate_multiplies", coordinate_multiplies)

best_chunk auth-fixture
scores [1.0, 0.0, 0.8, 0.1]
coordinate_multiplies 12

The question vector is [1.0, 0.0, 0.0], so only the first coordinate of each chunk contributes to each sum. The other multiplies still happen; they add zero. zip(..., strict=True) pairs coordinates and raises ValueError if the lengths differ. max(..., key=lambda i: scores[i]) returns the index with the largest score. A dense matrix multiply performs the same $n \times d$ coordinate products, often much faster than a Python loop.

Even though only coordinate 0 of the question vector is nonzero, this dense loop still performs all 12 multiplications (4 chunks $\times$ 3 coordinates). A dense matrix operation doesn't automatically exploit arbitrary zero values. Sparse representations and specialized kernels can change the work, but they require a different execution path.

With 50,000 chunks and 768 coordinates, the same pass uses 38,400,000 multiplications before it chooses citations. The exact score pass is Θ(n * d) in time.

Scoring work is only half the picture. Those float32 vectors take 50,000 * 768 * 4 = 153,600,000 bytes, about 146.5 MiB: Θ(n * d) input memory, separate from temporary workspace. The build-it function later materializes its allowed rows, adding O(n) temporary references; a streaming exact scan can keep only the query vector and running winners in extra memory.

Return top-k without sorting every result

The answer needs only k = 2 citations, but a full sort orders every candidate. That's the mismatch to keep in mind as we choose a selection method. In product code, treat k <= 0 as an empty result so heap logic never reads from an empty list.

scores = [
    (1.0, "auth-fixture"),
    (0.0, "ci-timeout"),
    (0.8, "cli-flag"),
    (0.1, "deploy-runbook"),
]

top_two = sorted(scores, reverse=True)[:2]
print("top_two", top_two)
print("fully_reordered_scores", len(scores))

top_two [(1.0, 'auth-fixture'), (0.8, 'cli-flag')]
fully_reordered_scores 4

The result is correct, but two losers were sorted even though they'll never be shown. Full sorting costs O(n log n) comparisons in the worst case and materializes all n scored results. When k is small, what structure could keep the weakest winner easy to replace without ordering every candidate?

In Data Structures for AI you already used a min-heap so the weakest retained result sits at the root. A bounded min-heap stores only the top-k values seen so far. A stronger candidate can replace the current weakest winner. Repairing the heap follows one path through a tree with k entries. That path has at most about log2(k) parent-child steps: three for k = 8, ten for k = 1,024.

Python's heapq helpers implement that structure in the lab. Tuples compare left to right, so (score, id) breaks a score tie using the ID: this code favors the lexicographically larger string. Change the key if your requirement differs; for integer IDs, (score, -doc_id) makes smaller IDs win ties.

Here's the invariant that lets us trust the heap: after each input, it contains the strongest k items seen so far, or all the items if fewer than k have arrived. Once it's full, a candidate no stronger than the root can't displace a winner. A stronger candidate replaces the weakest winner, and the heap is repaired. The condition item > heap[0] matters because heapreplace replaces the root even when the incoming item is worse.^{[5]Reference 5heapq: Heap queue algorithmhttps://docs.python.org/3/library/heapq.html}

For 2 <= k <= n, maintaining the heap costs O(n log k) after scoring, with O(k) selection memory. Sorting the final winners adds O(k log k), still within that bound. For k = 1, one running maximum gives O(n) selection; for k > n, there are only n items to retain.^{[6]Reference 6Introduction to Algorithms (3rd ed.)https://mitpress.mit.edu/9780262033848/introduction-to-algorithms/}

The next function accepts the same scored rows and keeps at most two winners. Compare its result with the full sort:

from heapq import heappush, heapreplace

scores = [
    (1.0, "auth-fixture"),
    (0.0, "ci-timeout"),
    (0.8, "cli-flag"),
    (0.1, "deploy-runbook"),
]

def bounded_top_k(items: list[tuple[float, str]], k: int) -> list[tuple[float, str]]:
    if k <= 0:
        return []
    heap: list[tuple[float, str]] = []
    for item in items:
        if len(heap) < k:
            heappush(heap, item)
        elif item > heap[0]:
            heapreplace(heap, item)
    return sorted(heap, reverse=True)

top_two = bounded_top_k(scores, k=2)
print("top_two", top_two)
print("selection_values_stored", len(top_two))
assert bounded_top_k(scores, k=0) == []
for fixture in (scores, list(reversed(scores)), [(0.5, "a"), (0.5, "c"), (0.5, "b")]):
    for k in (1, 2, len(fixture), len(fixture) + 1):
        assert bounded_top_k(fixture, k) == sorted(fixture, reverse=True)[:k]

top_two [(1.0, 'auth-fixture'), (0.8, 'cli-flag')]
selection_values_stored 2

The heap improves selection, not exact vector scoring. You still paid O(n * d) to produce the four scores above, and you'd still pay 38.4 million coordinate products at n = 50,000, d = 768. A faster tail step can't fix an expensive upstream step.

Keep separate time and memory ledgers

The same algorithm can have one growth rate for work and another for storage. Before changing code, label each number as input memory, output memory, or temporary workspace:

A heap can cut selection memory from Θ(n) to Θ(k), but it can't shrink a vector matrix that the service keeps resident. Likewise, changing a nested loop to stream its counter can save memory while leaving its Θ(n²) time unchanged. Time and memory need separate fixes.

When exact vector scoring exceeds the compute budget, changing the sort won't help because the score pass comes first. We need to change which candidates the search examines, then compare that path against the exact baseline.

View larger

Spot a quadratic trap

Linear work doubles when the input doubles. Quadratic work grows much faster because each item is compared with many others. To spot the difference in a loop, ask whether the inner range grows along with the outer one.

A tempting cleanup step might compare every retrieved chunk against every other retrieved chunk to remove near-duplicates. For n chunks it evaluates n * (n - 1) / 2 pairs, which is Θ(n²). Before running the counter, estimate the jump from 16 to 1,000 chunks; the exact output makes the trap hard to ignore.

def pair_comparisons(n: int) -> int:
    comparisons = 0
    for left in range(n):
        for right in range(left + 1, n):
            comparisons += 1
    return comparisons

for n in (4, 8, 16, 1_000):
    print(f"n={n:>4} pairs={pair_comparisons(n):>6}")

n=   4 pairs=     6
n=   8 pairs=    28
n=  16 pairs=   120
n=1000 pairs=499500

At n = 1,000, a request-time pairwise pass makes 499,500 comparisons. That's a concrete failure case: a harmless-looking loop can consume the latency budget before generation begins. The count came from the changing inner range, not from the fact that the code has two for statements.

Two nested loops aren't automatically quadratic. If the inner loop visits exactly 32 tokens for each chunk, the count is 32 * n, which is Θ(n) because 32 doesn't grow with n. Count the iterations, not the indentation.

For exact duplicate removal, a set changes the operation. It remembers texts already seen instead of comparing every pair. Semantic similarity still needs a method that compares meaning or representation.

candidate_texts = [
    "auth fixture expiry",
    "auth test failure",
    "auth fixture expiry",
    "ci timeout retry",
    "auth test failure",
]

seen = set()
unique = []
membership_checks = 0
for text in candidate_texts:
    membership_checks += 1
    if text not in seen:
        seen.add(text)
        unique.append(text)

print("unique_chunks", len(unique))
print("membership_checks", membership_checks)

unique_chunks 3
membership_checks 5

This path performs exactly n membership checks. With normal hash-table behavior, each check takes expected O(1), so the pass takes expected O(n) total time. Collisions can make one check O(n). Hashing still reads each text, so unbounded string length remains part of the real cost.

It still can't tell whether two different strings mean nearly the same thing. That's a different requirement, so don't trade away semantic matching just to get a nicer complexity label.

The same square shape appears later in attention

The same counting habit will reappear in Transformer internals. For now, carry one preview: scaled dot-product attention computes a score from every query position against every key position. Vaswani et al. write that matrix as $QK^{\top}$, so a self-attention sequence with t token positions has t * t score cells.^{[7]Reference 7Attention Is All You Need.https://arxiv.org/abs/1706.03762}

You don't need the rest of the model yet. If the prompt doubles from 256 to 512 tokens, predict the cell count before looking at the program.

def score_matrix_size(tokens: int) -> tuple[int, float]:
    cells = tokens * tokens
    mib_if_float32 = cells * 4 / (1024 * 1024)
    return cells, mib_if_float32

for tokens in (128, 256, 512, 1_024):
    cells, mib = score_matrix_size(tokens)
    print(f"tokens={tokens:>4} cells={cells:>7} raw_matrix_mib={mib:>4.2f}")

tokens= 128 cells=  16384 raw_matrix_mib=0.06
tokens= 256 cells=  65536 raw_matrix_mib=0.25
tokens= 512 cells= 262144 raw_matrix_mib=1.00
tokens=1024 cells=1048576 raw_matrix_mib=4.00

Doubling tokens quadruples cells. Treat this as a per-layer, per-head order of magnitude for a dense float32 score grid for one sequence during prefill, before multi-head batching, fused kernels, or KV-cache decode optimizations.

If each query and key has $d_k$ coordinates (the per-head width in later chapters), forming that grid also performs $\Theta(t^{2} d_{k})$ multiply-accumulate work per head, while the score grid itself takes $\Theta(t^{2})$ values. This doesn't model a complete language model or serving memory yet. It transfers the algorithm skill: when you see an all-pairs matrix, ask whether an input is being squared.

Prefill versus decode: two computational regimes

An autoregressive Transformer processes the prompt before it starts generating, then adds one token at a time. These are two useful phases to count separately. The dense-attention costs below are per head, with head width $d_k$; they don't include projections, feed-forward layers, or the rest of the serving request:

1. Prefill (prompt evaluation): Process the prompt's $T$ positions together, possibly in chunks. Causal attention lets each position attend to itself and earlier positions. Dense attention work is $\Theta(T^2 d_k)$ per head, up to the causal constant factor. Matrix-matrix operations can reuse loaded values well, so large prefills often become compute-bound. Prefill contributes to time to first token (TTFT), which also includes queueing, preprocessing, and other work.
2. Decode (token generation): For one sequence, a step adds one query against a cached history of length $t$. Its attention work is $\Theta(t d_k)$ per head. Small decode batches often become bandwidth-bound while reading model weights and KV state. Larger batches, kernel choices, and hardware can change the bottleneck. Decode contributes to inter-token latency (ITL); batching and scheduling also affect the gaps observed by a client.

Neither phase has one universal FLOPs-per-byte value. A batched decode also combines work across sequences, so it needn't execute as one literal vector-matrix kernel. Fused attention such as FlashAttention can avoid storing the full score grid in device memory while retaining dense attention's quadratic arithmetic.^{[8]Reference 8FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesshttps://arxiv.org/abs/2205.14135}

View larger

Attention compute scaling versus KV cache memory

The KV cache stores past key and value activations so decode doesn't recompute them at every step. One dense-attention decode step is linear in cached length, but the cache grows as generation continues. Across $G$ generated tokens after a prompt of length $T$, attending to the growing history totals $\Theta((GT + G^2)d_k)$ work per head. Caching saves recomputation; it doesn't make an arbitrarily long continuation linear in total output length.

For an exact decode trace, start with eight completed positions already cached. Consuming one additional input appends its key and value before its causal attention reads nine positions, including itself. Two more consumed inputs read ten and eleven positions. Those three rows total 30 key reads and leave eleven cached positions. In general, G consumed inputs after T previously cached positions read G × T + G(G + 1)/2 keys per head. This counts consumed decoder inputs; prompt prefill already produces the logits for the first sampled output.

For Llama 3 8B's 32 layers, 8 Grouped-Query Attention (GQA) key-value heads, and head dimension $d_k = 128$, assume a 16-bit KV cache (2 bytes per element).^{[9]Reference 9Llama model specificationshttps://github.com/meta-llama/llama-models/blob/main/models/sku_list.py} Each stored token then requires:

The first factor of 2 counts separate keys and values; the last counts bytes per element. This model has 32 query heads, but groups of four share one key-value head. Cache storage therefore uses the 8 KV heads, not the 32 query heads. Each query head still computes its own attention scores.^{[10]Reference 10GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245[9]Reference 9Llama model specificationshttps://github.com/meta-llama/llama-models/blob/main/models/sku_list.py}

At a concurrent batch size of $B = 16$ streams, context length directly dictates your serving memory budget:

• At 1,024 tokens per stream: $16 \times 1{,}024 \times 128 \text{ KiB} = 2.0 \text{ GiB}$.
• At 4,096 tokens per stream: $16 \times 4{,}096 \times 128 \text{ KiB} = 8.0 \text{ GiB}$.
• At 8,192 tokens per stream: $16 \times 8{,}192 \times 128 \text{ KiB} = 16.0 \text{ GiB}$.

These counts include raw KV values only, with equal-length streams and no prefix sharing. Model weights, temporary workspace, cache metadata, and unused space in partially filled blocks consume additional memory. Cache quantization or sharing changes the count.

View larger

Measure scale and install a budget guard

Big-O shows how a workload can grow, but it can't tell you whether today's hardware meets the service objective. Count expected work to spot risk, then benchmark the real implementation on representative data and hardware. A guard can reject or reroute requests whose predicted work exceeds the supported budget.

The first guard uses a multiply count, not a made-up millisecond conversion. That keeps the example honest: hardware and kernels determine how many milliseconds those operations take.

def exact_score_work(candidate_chunks: int, embedding_dim: int) -> int:
    return candidate_chunks * embedding_dim

def within_budget(candidate_chunks: int, embedding_dim: int, max_multiplies: int) -> bool:
    return exact_score_work(candidate_chunks, embedding_dim) <= max_multiplies

budget = 1_000_000
for candidates in (500, 2_000, 50_000):
    work = exact_score_work(candidates, embedding_dim=768)
    print(candidates, work, within_budget(candidates, 768, budget))

500 384000 True
2000 1536000 False
50000 38400000 False

View larger

When a request exceeds the budget, authorization still applies and the answer still needs relevant evidence. Exact ranking over that many authorized candidates no longer fits this budget. Possible next experiments are a stricter legitimate filter, an offline-created index, or an approximate nearest-neighbor index evaluated for recall.

A later option is HNSW, a hierarchical proximity-graph approach to approximate nearest-neighbor search. It trades exact scoring of every vector for graph-guided candidate exploration. The original paper reports logarithmic scaling under its construction assumptions; treat that as a reason to measure recall and latency together, not as a free O(log n) replacement for the exact baseline.^{[11]Reference 11Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphshttps://arxiv.org/abs/1603.09320} That belongs in Vector DB Internals, not in this foundation exercise.

Guarding tail latency in production

Operation counts tell us how work grows. To enforce a wall-clock budget in production, we also need to learn where the time goes:

• Look beyond the median: The mean and p50 are different statistics, and neither describes the slowest requests. A hypothetical service could have a 20 ms median and an 800 ms p99. Measure the percentiles and deadlines your users need, then include queue time as well as computation in the budget.
• Profile Python collection: CPython uses reference counting plus a cyclic garbage collector. Collection can contribute pauses, but their duration depends on the object graph and runtime. Measure collection activity before changing it. gc.disable() stops automatic cyclic collection. Reference counting continues, and unreachable cycles can accumulate without a deliberate collection policy.^{[12]Reference 12gc: Garbage Collector interfacehttps://docs.python.org/3/library/gc.html}
• Bound batch formation time: A rule that waits for exactly eight arrivals can leave a low-traffic request waiting. Ray Serve's batch_wait_timeout_s limits how long it tries to fill a batch after the first request; Triton has its own max_queue_delay_microseconds setting. These settings don't cap total request latency when execution is already busy.^{[13]Reference 13Ray Serve: Dynamic Request Batchinghttps://docs.ray.io/en/latest/serve/advanced-guides/dyn-req-batch.html[14]Reference 14Triton Inference Server: Dynamic Batcherhttps://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/batcher.html}
• Measure preprocessing separately: Tokenization or cleaning can leave an accelerator waiting for input. Worker pools and overlapping preparation with execution can help, but they add scheduling, memory, and transfer costs. Compare the measured pipeline before and after rather than assuming more workers keep the GPU busy.

There is another useful retrieval pattern: retrieve, then refine. A quick first-stage score proposes a shortlist, then a slower scorer or cross-encoder ranks only those candidates. A one-stage top-k heap still uses one score for both selection and ranking. Refinement needs a second scoring function, and the shortlist must preserve authorization and enough recall for the final answer.

The fixture below makes that second stage visible using precomputed scores, not real model calls. Its rows have already passed the permission filter. The first stage reads four cheap scores; the second reads richer scores for only three shortlisted rows. A real reranker would compute those richer scores on demand. Here the counters track score reads so you can inspect which rows reach each stage.

from heapq import nlargest

authorized_rows = [
    {"id": "auth-fixture", "cheap_score": 0.91, "refine_score": 0.80},
    {"id": "ci-timeout", "cheap_score": 0.88, "refine_score": 0.55},
    {"id": "cli-flag", "cheap_score": 0.75, "refine_score": 0.95},
    {"id": "deploy-runbook", "cheap_score": 0.10, "refine_score": 0.20},
]

def retrieve_then_refine(
    rows: list[dict],
    shortlist_size: int,
    result_size: int,
) -> tuple[list[str], int, int]:
    shortlist = nlargest(
        shortlist_size,
        rows,
        key=lambda row: (row["cheap_score"], row["id"]),
    )
    refined = sorted(
        ((row["refine_score"], row["id"]) for row in shortlist),
        reverse=True,
    )[:result_size]
    return [row_id for _, row_id in refined], len(rows), len(shortlist)

citations, candidate_scored, refinement_scored = retrieve_then_refine(
    authorized_rows,
    shortlist_size=3,
    result_size=2,
)
print("candidate_stage_scored", candidate_scored)
print("refinement_stage_scored", refinement_scored)
print("citations", citations)
assert citations == ["cli-flag", "auth-fixture"]
assert candidate_scored == 4
assert refinement_scored == 3

# Offline reference: rank every authorized row under the final scorer.
reference_ids = {
    row["id"]
    for row in sorted(
        authorized_rows, key=lambda row: (row["refine_score"], row["id"]), reverse=True
    )[:2]
}
for shortlist_size in (2, 3):
    result_ids, _, scored = retrieve_then_refine(authorized_rows, shortlist_size, 2)
    recovered = len(set(result_ids) & reference_ids)
    print(f"shortlist={shortlist_size} refine_reads={scored} recovered={recovered}/2 ids={result_ids}")
    assert recovered == (1 if shortlist_size == 2 else 2)

candidate_stage_scored 4
refinement_stage_scored 3
citations ['cli-flag', 'auth-fixture']
shortlist=2 refine_reads=2 recovered=1/2 ids=['auth-fixture', 'ci-timeout']
shortlist=3 refine_reads=3 recovered=2/2 ids=['cli-flag', 'auth-fixture']

Here the cheap score chooses candidates and the richer score changes their order. A shortlist of two still returns two citations, but it drops cli-flag, the final scorer's best result. Refinement can reorder auth-fixture and ci-timeout; it can't bring back a row it never sees. The three-row shortlist keeps both reference winners.

Here recall at 2 against the final-score reference is the number of recovered reference IDs divided by 2: 1/2 versus 2/2. Computing that reference requires all four richer scores during offline evaluation. It doesn't belong in the live shortcut, and a perfect result on this one fixture doesn't establish corpus-wide recall. Compare authorized candidates under the same final scoring and tie rules, then measure shortlist quality before celebrating lower refinement cost.

Build it: an authorized exact top-k retriever

The guard runs before vector scoring, and the permission fixture makes the boundary visible: it contains three allowed auth rows, one private auth row, and one billing row. Only the allowed rows contribute to the scoring budget:

• Authorization happens first.
• Every authorized vector is scored exactly.
• A bounded heap retains only k citations.
• A preflight count rejects work above the configured multiply budget.

from heapq import heappush, heapreplace

rows = [
    {"id": "auth-fixture", "project": "auth", "principals": {"engineer:u3"}, "vector": [1.0, 0.0, 0.1]},
    {"id": "auth-ci-timeout", "project": "auth", "principals": {"engineer:u3"}, "vector": [0.2, 1.0, 0.1]},
    {"id": "auth-cli-flag", "project": "auth", "principals": {"engineer:u3"}, "vector": [0.8, 0.1, 0.4]},
    {"id": "auth-secrets-private", "project": "auth", "principals": {"engineer:u1"}, "vector": [0.0, 0.3, 1.0]},
    {"id": "billing-auth-fixture", "project": "billing", "principals": {"engineer:u3"}, "vector": [1.0, 0.0, 0.0]},
]

def dot(left: list[float], right: list[float]) -> float:
    return sum(a * b for a, b in zip(left, right, strict=True))

def authorized_exact_top_k(
    question: list[float],
    all_rows: list[dict],
    project: str,
    principal: str,
    k: int,
    max_multiplies: int = 1_000_000,
) -> tuple[list[str], int]:
    if k <= 0:
        return [], 0
    allowed = [
        row
        for row in all_rows
        if row["project"] == project and principal in row["principals"]
    ]
    required_work = len(allowed) * len(question)
    if required_work > max_multiplies:
        raise ValueError(
            f"exact scoring needs {required_work} multiplies; budget is {max_multiplies}"
        )
    heap: list[tuple[float, str]] = []
    coordinate_multiplies = 0
    for row in allowed:
        score = dot(question, row["vector"])
        coordinate_multiplies += len(question)
        item = (score, row["id"])
        if len(heap) < k:
            heappush(heap, item)
        elif item > heap[0]:
            heapreplace(heap, item)
    ranked = sorted(heap, reverse=True)
    return [chunk_id for _, chunk_id in ranked], coordinate_multiplies

citations, work = authorized_exact_top_k(
    [1.0, 0.0, 0.0],
    rows,
    "auth",
    "engineer:u3",
    k=2,
)
print("citations", citations)
print("coordinate_multiplies", work)
assert "auth-secrets-private" not in citations
assert "billing-auth-fixture" not in citations
assert work == 9
assert authorized_exact_top_k(
    [1.0, 0.0, 0.0],
    rows,
    "auth",
    "engineer:u3",
    k=0,
) == ([], 0)

try:
    authorized_exact_top_k(
        [1.0, 0.0, 0.0], rows, "auth", "engineer:u3", k=2, max_multiplies=8
    )
except ValueError as error:
    print("rejected:", error)
else:
    raise AssertionError("over-budget scoring should have been rejected")

citations ['auth-fixture', 'auth-cli-flag']
coordinate_multiplies 9
rejected: exact scoring needs 9 multiplies; budget is 8

The first call stays within budget, scores three allowed vectors, and returns two citations. The second call has room for only eight multiplies, so it stops before any dot product. Authorization still scans the input rows, and the guard doesn't bound that work or elapsed time. It protects one named stage, which is why we need to measure the rest of the request too.

For this lab, vectors contain finite numbers and must have the query's dimension. zip(..., strict=True) detects a mismatched dimension instead of silently dropping coordinates. Compare any optimized retriever against this exact ranking on the same allowed rows, with the same tie rule.

Check the growth by hand

You've seen these counts in code. Try each one without code now: write down the input that grows, the operation you're counting, and your expected result before checking the sketches.

1. A project has n = 2,500 authorized chunks, each with d = 384 coordinates. How many coordinate multiplications does an exact dot-product scan perform?
2. A full sort orders 10,000 scored results even though the assistant returns k = 5. Which part does a bounded heap reduce, and which part remains unchanged?
3. A near-duplicate cleanup compares each of 200 chunks with every later chunk. How many comparisons occur?
4. An attention preview changes prompt length from 256 tokens to 1,024 tokens. By what factor does the number of score cells grow?
5. Your guard allows one million coordinate multiplications. Does exact scoring fit for n = 4,000, d = 256?
6. Three cached prompts share the first 6 tokens with a new 8-token request, then diverge. About how many token comparisons does a naive per-key scan do if each key stops at the first mismatch?
7. A sorted catalog has 65,536 chunk IDs. What is the worst-case number of midpoint checks for binary search?

Solution sketches

1. 2,500 * 384 = 960,000 coordinate multiplications.
2. The heap reduces selecting the best five values from full-sort work to bounded top-k work. It doesn't remove the exact score pass that produced all 10,000 values.
3. 200 * 199 / 2 = 19,900 comparisons.
4. Token length grows by 1,024 / 256 = 4, so square score-cell count grows by 4 * 4 = 16.
5. No. 4,000 * 256 = 1,024,000, which exceeds the one-million budget.
6. Each of the 3 keys compares 6 matches plus 1 mismatch, so about 21 comparisons. A radix-tree walk would scale with the prefix length, not with n.
7. At most 17 midpoint checks. Although $2^{16} = 65{,}536$, this half-open binary search can spend a seventeenth comparison checking the final one-item interval. For example, searching for the first ID visits 17 midpoints.

Which input grows which cost

When a request slows down, start by asking which input grew and what operation now repeats. This table gives you a first place to look:

Keep these three habits as you move into later AI systems:

1. Name the operation and its growing inputs before stating a complexity.
2. Keep a smallest-correct baseline so optimization can be checked for correctness.
3. Couple latency work with product requirements: authorization, citation quality, and recall don't disappear when a faster algorithm arrives.