Weight reuse explains why prefill and decode need different tuning

TL;DR: A larger batch lets more tokens share the cost of reading the weights. It also adds computation and cache reads, so a step can take longer even while the server produces more tokens per second overall. Compare total throughput with the gaps between tokens for each user before increasing concurrency.

Where you are. Middle of the request module. The previous lesson followed three prompt IDs through a small Transformer. Continue that request through two decode steps, then use the changing shapes to explain a batch-size sweep.

Continue the same request: what gets recomputed?

The previous lesson's toy model has two blocks, two KV heads per block, and two coordinates per head. Prefill processes IDs [1, 4, 6] and selects the first output ID, 2. At that moment each block has keys and values for three processed positions. The newly selected ID has not passed through the model yet.

Feed ID 2 back at position three. Each block computes its new query, key and value; its query reads the three earlier K/V entries plus the new entry. Older queries and feed-forward outputs do not need to be recomputed. The same learned matrices still transform the new position. The second output is again ID 2, but it occupies a different position and generally receives different contextual features.

The CPU check compared those one-row decode logits against a fresh four-row causal forward pass. They agree within the stated float64 tolerance. This follows from causality, fixed parameters and the same positional treatment. Changing the model revision, position offsets or attention mask can invalidate that reuse. Training-mode dropout would also change the comparison.

Count the state explicitly using a hypothetical two-byte K/V representation. The CPU code uses eight-byte float64; the following is a separate payload calculation, not its measured allocation:

KV per processed position = 2 (K and V) × 2 blocks × 2 heads × 2 coordinates × 2 bytes
                          = 32 bytes
sanity: Q is temporary; adding it would count three tensors instead of two.

Stop after the third output and the model has processed five positions, not six. Its last selected ID needs no K/V unless generation continues. This is logical payload for ordinary causal decoding with no prefix sharing, speculation or preallocation. Real allocation includes block rounding, reserved growth and workspace. The KV Cache Sizing reference covers those boundaries.

Start with the weight-only estimate

Why does processing several positions together help? A linear layer uses the same weights for each position. If the GPU can read a group of weights once and use it across B positions, those positions share the memory cost.

Let N be the number of participating weights and q the bytes stored per weight. For this estimate, assume the weights do not remain in on-chip cache between passes and each tile of weights can be reused across all B positions. During prefill, those positions may come from one prompt. During ordinary decode, each active sequence contributes one position.

weight traffic ≈ N × q bytes
linear work ≈ 2 × N × B FLOPs
weight-only intensity ≈ 2NB / Nq = 2B/q FLOPs per byte
for BF16: q = 2 bytes, so intensity ≈ B FLOPs per byte
sanity: batch one gives 1 FLOP/byte, not 2 FLOPs/byte

This estimate describes linear layers, not all transformer work. Attention uses a different operation shape. Activations, output writes and limited tile reuse increase traffic. MoE batches can activate more distinct experts as B rises, so their weight traffic need not remain fixed. The matrix multiplication performance guide explains the role of matrix dimensions and reuse; use the full traffic expression when the omitted tensors are large.

Suppose the device's compute-to-bandwidth ratio, or ridge point, is 300 FLOPs/byte. Below this ratio, the memory bandwidth gives the lower performance ceiling; above it, arithmetic capacity does. With two-byte BF16 weights, the estimate becomes B FLOPs/byte:

The crossover in this simplified example is B ≈ 300. A different precision, hardware profile or attention workload moves it. This is not a universal minimum batch size for efficient serving.

Add the traffic that grows with every user

The weight calculation makes larger batches look attractive because the weight bytes stay fixed while more positions use them. The cache changes that picture. Each additional sequence brings its own history, and full-attention decode reads that history's keys and values.

To include those reads, suppose every sequence retains T tokens and each token needs k bytes of KV state. A batch of B sequences then reads approximately B × T × k cache bytes as well as the weights:

traffic per step ≈ Nq + B × T × k
memory time ≈ traffic / sustained bandwidth
linear compute time ≈ 2NB / sustained operation rate
optimistic step floor ≈ max(memory time, linear compute time)
aggregate output rate ≈ B / step time
per-sequence output rate ≈ 1 / step time

These expressions omit attention FLOPs, KV writes, launches and communication. They also assume the cache representation is read once per step; actual kernels can incur additional traffic. For unequal contexts use the sum of retained KV bytes across sequences instead of B × T × k. Sliding-window, latent and recurrent states need architecture-specific terms.

Consider an illustrative 8-billion-weight BF16 model and 2 TB/s sustained bandwidth. Its fixed weight traffic is 16 GB. Assume the model, the largest retained state below and workspace fit the available device memory. Each sequence adds 0.25 GB of KV reads per step and other work is small enough for the memory estimate to dominate:

batch 32 traffic = 16 GB + 32 × 0.25 GB = 24 GB
step estimate = 24 / 2,000 = 0.012 s
aggregate rate = 32 / 0.012 ≈ 2,667 tokens/s
per-sequence rate = 1 / 0.012 ≈ 83 tokens/s
weight/KV crossover = 16 GB / (0.25 GB per sequence) = 64 sequences
sanity: at batch 64 the cache and weight traffic are equal, not absent

Read the table in both directions. Going from one sequence to 32 raises total output from about 123 to 2,667 tokens/s, because each step now produces 32 tokens. But the step itself grows from 8.125 to 12 ms, so each user receives tokens less often. More total output does not necessarily mean a faster response for one person.

These are calculated bounds under the stated assumptions, not measured throughput. As B grows, the weight cost is shared more widely, but each sequence still adds 0.25 GB of reads. The memory-only rate therefore approaches 2,000/0.25 = 8,000 tokens/s. Device capacity and the compute ceiling can stop the increase much earlier. Longer contexts add more reads per sequence and lower that limiting rate.

Make the serving decision from both rates

The useful batch size depends on what users can tolerate. If requests fit in memory and a measured 12 ms token gap meets the product's target, the extra throughput may make batch 32 worthwhile. Test that choice with the actual context lengths and arrivals. Reduce concurrency or revisit the schedule if token gaps become too long or requests wait too long for their first token.

Also account for any time spent collecting a batch before execution. That waiting period, called a batching window, adds delay before the first step begins. The step-time calculation above does not include it.

Orca schedules at iteration granularity so completed sequences leave and new work can enter. That changes which sequences share the next iteration; it does not make the cache bytes or waiting time disappear.

Do this before moving on

For the tiny request, double the number of blocks while preserving every other dimension. Per-position K/V becomes 64 bytes; after prefill and two ordinary decode steps it occupies 320 bytes. The final selected output still has no computed K/V. Explain why doubling output allowance may require extra reservation before those bytes are actually used.

Repeat the table with 0.5 GB of KV reads per sequence. At batch 32, traffic becomes 32 GB, step time 16 ms, and aggregate throughput 2,000 tokens/s. The weight/KV crossover moves to batch 32. Explain why adding faster arithmetic alone would not improve this memory estimate, then name one measurement that could invalidate the estimate: actual kernel traffic, achieved bandwidth or attention time.

Go deeper

• Prefill vs Decode carries the same arithmetic with the model dimensions filled in.
• Continuous Batching explains how requests enter and leave the running batch between iterations.
• KV Cache Sizing gives you the kv term as a formula, which the next module turns into a capacity plan.
• Arithmetic Intensity by Operation tabulates where these operations land so you can check your reasoning.
• How big can a batch be before decode is compute-bound? is exactly the derivation above, asked as an interview question.
• Tokens per second for a 70B single stream is the batch-one case worked end to end.

Key takeaways

• Weight-only BF16 intensity is B FLOPs/byte under ideal reuse; it does not include attention or activation traffic.
• Aggregate rate is B/step-time while per-sequence rate is 1/step-time; they can move in opposite directions.
• In the example, batch 32 gives a 12 ms memory estimate and about 2,667 aggregate tokens/s.
• With 16 GB weights and 0.25 GB KV per sequence, the traffic crossover is 64 sequences; doubling KV halves that batch size.
• Choose concurrency using measured TTFT, token gaps and memory headroom, not the weight-only roofline alone.

Check yourself

Answer before you look. Recalling it is what makes it stick; recognising it does not.

1. 1The example moves from batch 8 to batch 32. Aggregate throughput rises, but step time increases from 9 ms to 12 ms. What happens to one active user?

   Its output rate is unchangedIts output rate rises fourfoldIts output rate falls from about 111 to 83 tokens/sIts first token must arrive sooner

2. 2At what batch size does KV traffic equal 16 GB of weights if each sequence contributes 0.5 GB?

   Answer it, then reveal

3. 3Does a throughput plateau prove decode became compute-bound?

   Answer it, then reveal

4. 4After a three-position prompt and three emitted output tokens, how many positions have computed K/V in ordinary cached decoding?

   One: only the current query is retainedSix: every emitted token already has its own K/VFive: prefill plus two consumed output tokensThree: output tokens never create K/V

5. 5Does reusing the KV cache eliminate the model’s learned projections and feed-forward work?

   Answer it, then reveal