Implementation chains: the numbers that decide an engine component's design

Implementation rounds ask for code, and the design choices in that code come from a number. State the number before writing the function: it is what justifies the data structure, and an interviewer who asks "why not the simple version" is asking for it.

What concatenation copies

inputs   Llama 3.1 8B, 128 KiB of KV per token, generating 4,096 tokens
         by appending to the cache with torch.cat at every step
copies   step t copies all t tokens: 1 + 2 + ... + 4,096 = 8.39M token copies
bytes    8.39M x 131,072 B = 1.0 TiB moved to produce 512 MiB of cache
sanity   quadratic in length; a preallocated [B, Hkv, max_len, d] buffer
         written in place moves each token once

Preallocate and write at each row's current length. The causal mask then needs the offset: query t of a chunk sits at absolute position len - n_new + t.

Drilled in: GQA with a preallocated KV cache

When the allocator needs a block

inputs   16-token blocks; a sequence at 1,000 tokens
blocks   ceil(1,000 / 16) = 63, the last one holding 8 tokens
append   up to 1,008 tokens: still 63 blocks, no allocation
next     token 1,009 crosses the boundary: block 64
shared   a write into a partly filled block with refcount above 1 copies it
         first; only the first writer copies
sanity   allocation is a boundary event, so the hot path of a decode step
         usually allocates nothing

A refcount per physical block, a free list and a block table per sequence are the whole structure. The copy-on-write rule is the part to test.

Drilled in: Paged KV block manager with copy-on-write

Why top-p needs a short circuit

inputs   a 128,256-token vocabulary; probabilities sorted, then a float32
         cumulative sum
result   the final cumulative value can end below 1.0: 0.99976 for a
         peaked distribution in a float32 check
search   with p = 1.0, "first index where the sum reaches p" finds nothing
fix      return early for p >= 1 and for temperature 0; softmax in float64
sanity   128,256 rounded float32 terms lose 1e-5 to 2e-4 of the total,
         more as the distribution gets more peaked

Keep every token whose probability is at least the probability at the cutoff index, so tied tokens are kept or dropped together.

Drilled in: Top-p sampling with ties and temperature

The token budget bounds a step

inputs   Llama 3.3 70B FP8 on TP=4 H100 SXM, 1,979 TFLOP/s per GPU at peak
2,048    2 x 70.55e9 x 2,048 / (4 x 1,979e12) = 36.5 ms of compute at most
8,192    the same at 8,192 tokens: 146 ms
order    decodes first (one token each), then the rest of a partial prefill,
         then new requests first come, first served
sanity   every running decode waits for the whole step, so the budget is an
         inter-token latency limit as well as a throughput knob

Spending the budget in a fixed order is what makes the scheduler testable: a test can assert exactly which requests ran in each step.

Drilled in: Continuous batching scheduler with a token budget

What a radix hit saves

inputs   a shared 2,560-token system prompt on Llama 3.1 8B BF16, H100 SXM
KV       2,560 x 131,072 B = 320 MiB stored once
prefill  2 x 8.03e9 x 2,560 / 989e12 = 41.6 ms skipped per hit at peak
sanity   the tree pins the matched path with a reference count while a
         request runs, and evicts only unlocked leaves

Split an edge wherever a lookup diverges inside it. When a leaf's last child is evicted, the parent becomes a leaf and joins the eviction heap at once.

Drilled in: Radix-tree prefix cache with LRU

Position ids in BF16

inputs   RoPE angle = position x frequency, computed for position 100,000
BF16     8 significant bits: from 65,536 to 131,072 the spacing is 512
result   100,000 rounds to 99,840
sanity   every position in that range is off by up to 256, so long-context
         tokens rotate as if they were somewhere else

Form position ids as integers and compute cos and sin in float32 or wider. Test that chunked plus token-by-token rotation matches rotating the whole sequence at once, bit for bit.

Drilled in: RoPE for a decode step with position offsets

INT4 packing

inputs   groups of 128 weights along K; scale = absmax / 7; values clipped
         to [-8, 7]; two nibbles per byte, even index in the low nibble
storage  0.5 B per weight + 2 B FP16 scale / 128 = 0.515625 B per weight
70B      70.55e9 x 0.515625 B = 36.4 GB, from 141.1 GB in BF16
sanity   the scale costs 3.1% on top of the packed nibbles

Multiply the activations by each group's integer weights first and apply that group's per-column scale once per group, not once per element.

Drilled in: INT4 packing and dequantizing matmul

Ring all-reduce traffic

inputs   one 8,192-wide BF16 hidden state, 16,384 B; TP=8
steps    p - 1 = 7 reduce-scatter steps, then 7 all-gather steps
sent     each rank sends 2 x (p - 1) / p x 16,384 = 28,672 B
sanity   bytes per rank approach twice the buffer whatever p is; the
         2 x (p - 1) = 14 sequential steps are what grow

At decode sizes the 14 dependent steps cost more than the bytes. That is the case for one-shot all-reduce kernels on NVLink.

Drilled in: Ring all-reduce simulation