Inference Optimization: KV Cache, Speculative Decoding, and Batching Strategies

A staging benchmark that lied

An engineering team at a Bengaluru fintech builds a UPI dispute-resolution assistant on top of a 70-billion-parameter chat model, paired with a 7-billion-parameter draft model from the same family for speculative decoding. In staging, with two or three engineers hitting the endpoint by hand, the numbers are exactly what the theory promises: roughly 2.1× more tokens per second than plain autoregressive decoding, for the same GPU, same model, same output distribution. The team ships it. Three weeks later, during the 8–9pm peak when a few hundred customers are simultaneously arguing with the bot about failed transactions, someone on-call notices token throughput has barely moved from the non-speculative baseline — and on the busiest node, it is actually worse than turning speculation off. Nobody changed the model. Nobody changed γ, the number of tokens the draft model proposes per round. The only thing that changed between the benchmark and the incident is how many requests were running on the GPU at once. This chapter explains exactly why that variable, batch size, is the one the staging benchmark forgot to test — and works out, from the same memory-bandwidth arithmetic this curriculum has already used for plain batching, the precise point at which speculative decoding stops being a free lunch and starts being a tax.

Arithmetic intensity, generalized to "positions per pass"

Two facts from earlier in this curriculum are the entire foundation here, and neither is re-derived: decoding a token is memory-bandwidth bound because the GPU must stream every parameter out of HBM to compute one step, and batching B independent sequences together is nearly free because those parameters are streamed once and reused for all B rows while they sit on-chip. The question this chapter adds is what happens when a single forward pass is not asked to produce one token per sequence, but several — which is exactly what a speculative-decoding verification pass does: it evaluates γ + 1 candidate positions per sequence (γ drafted positions plus one bonus position) in one shot.

Generalize the batching arithmetic with one new symbol, w: the number of token-positions processed per sequence in a given pass. Ordinary decoding and a single draft-model step both have w = 1. A speculative verification pass has w = γ + 1. For a model with N parameters stored at d bytes each, one pass over a batch of B sequences streams N·d bytes from HBM exactly once (independent of B and w, since every row reuses the same on-chip weights) and performs approximately 2N·B·w floating-point operations — the standard 2N-FLOPs-per-forward-token estimate for a transformer, the same convention this curriculum already uses for batching, extended here to w positions per sequence instead of one. Arithmetic intensity — FLOPs performed per byte moved — is therefore:

AI(B, w) = (2 * N * B * w) / (N * d) = 2 * B * w / d

Notice N cancels. Arithmetic intensity depends only on how many token-positions ride along per byte of weight streamed, never on how large the model itself is. A GPU has its own fixed ratio of peak compute to peak memory bandwidth, the ridge point R = FLOPS_peak / BW (in FLOPs per byte). Below the ridge, a pass is memory-bound and its latency is set purely by N·d/BW, unmoved by how many positions w ride along — this is precisely the regime the earlier speculative-decoding treatment relies on when it calls the extra γ positions "free." Above the ridge, the pass is compute-bound and its latency grows linearly with B·w. Setting AI(B,w) = R and solving for B gives the batch size at which a pass of width w crosses from free to costly:

B*(w) = R * d / (2 * w)

Because a verification pass has w = γ + 1 while ordinary decoding has w = 1, the verification pass's crossover batch size is exactly 1/(γ+1) of ordinary decoding's crossover — the wider the speculative window, the sooner, in terms of concurrent load, the "free" tokens stop being free.

Worked example 1: where does the crossover actually sit?

Take the fintech's hardware, an A100 80GB SXM4: FLOPS_peak = 312 TFLOPS of dense fp16 tensor-core throughput, BW = 2039 GB/s of HBM2e bandwidth (both from NVIDIA's published datasheet). At fp16, d = 2 bytes. Take γ = 4, matching the draft window this curriculum's speculative-decoding treatment already worked with.

FLOPS_peak = 312e12
BW = 2039e9
d = 2
gamma = 4

R = FLOPS_peak / BW
B_star_verify = R * d / (2 * (gamma + 1))
B_star_naive  = R * d / 2

print(round(R, 2), round(B_star_verify, 2), round(B_star_naive, 2))
# 153.02 30.6 153.02

Trace it: R = 312×10¹² / 2039×10⁹ ≈ 153.02 FLOPs/byte — a hardware constant, unrelated to any model. Plain decoding's own crossover, B*(1) = 153.02×2/2 = 153.02, says a single-token-per-sequence batch on this GPU does not become compute-bound until roughly 153 concurrent sequences — comfortably past what a chat workload usually runs per GPU, which is why plain batched decoding is treated as "always memory-bound" almost everywhere else in this curriculum. The verification pass's crossover, B*(5) = 153.02×2/10 = 30.60, is more than 5× lower — and 30 to 40 concurrent conversations on one GPU is an entirely ordinary peak-hour load, not an edge case.

Worked example 2: throughput across batch sizes, in tokens per second

Crossing the ridge doesn't just change a regime label; it changes what speculative decoding is worth in tokens per second, compared against plain batched decoding on the same hardware. Model one full speculative round as γ sequential cheap draft steps (each with w = 1, using the 7B draft model's own parameter count and its own memory/compute costs) followed by one wide verification step (w = γ + 1, using the 70B target model), each step's latency being the max of its memory time and compute time exactly as derived above. The expected tokens emitted per round for a given per-token acceptance rate α is the E[T] = (1 − α^γ+1)/(1 − α) result this curriculum already derived from the rejection-sampling accept/reject test; it is not re-derived here, only re-used, because batching changes nothing about a single sequence's own acceptance probabilities — the exactness guarantee of the accept/reject test is completely unaffected by how many other sequences share the GPU.

def round_costs(FLOPS_peak, BW, d, N, N_draft, gamma, alpha, B):
    mem_target = N * d / BW
    mem_draft  = N_draft * d / BW
    E_T = (1 - alpha**(gamma + 1)) / (1 - alpha)

    compute_verify = 2 * N * B * (gamma + 1) / FLOPS_peak
    T_verify = max(mem_target, compute_verify)

    compute_draft_step = 2 * N_draft * B * 1 / FLOPS_peak
    T_draft_step = max(mem_draft, compute_draft_step)

    T_round = gamma * T_draft_step + T_verify
    T_naive = max(mem_target, 2 * N * B * 1 / FLOPS_peak)

    thr_spec  = B * E_T / T_round
    thr_naive = B / T_naive
    return thr_naive, thr_spec, thr_spec / thr_naive

N, N_draft = 70e9, 7e9
gamma, alpha = 4, 0.75
for B in [1, 8, 32, 64, 81, 128, 256]:
    thr_naive, thr_spec, speedup = round_costs(312e12, 2039e9, 2, N, N_draft, gamma, alpha, B)
    print(B, round(thr_naive, 1), round(thr_spec, 1), round(speedup, 3))

# 1   14.6    31.7   2.179
# 8   116.5   253.9  2.179
# 32  466.1   983.5  2.110
# 64  932.1   1141.5 1.225
# 81  1179.7  1181.3 1.001
# 128 1864.2  1241.1 0.666
# 256 2228.6  1259.1 0.565

Every line traces to the formulas above it. At B = 1 and B = 8, both far below the 30.6 crossover, the verification pass is still fully memory-bound, so throughput scales linearly with B and the speedup sits at the full 2.179× the accept/reject dynamics with α = 0.75 predict — matching, as it should, the low-concurrency case this curriculum's speculative-decoding chapter already worked through. At B = 32, just past the crossover, the verification pass has started paying real compute time and the speedup has already slipped to 2.110×. By B = 64 it is down to 1.225×. Solving thr_spec(B) = thr_naive(B) numerically lands the break-even point at B ≈ 81.1 — past that, speculative decoding is not neutral, it is actively slower than turning it off, because every one of the γ discarded or partially-wasted draft positions is now competing for the same compute units the target model's own useful FLOPs need. At B = 256, a plausible peak-hour batch on a busy inference node, speculation is running at 0.565× — not a smaller win, a genuine regression, because the naive path has kept climbing toward its own much higher 153-sequence crossover while the verification path plateaus early.

The other cost of batching speculation: a ragged batch's KV cache

The throughput story above assumes every sequence in the batch commits the same amount of new cache content each round, which is never actually true. Every sequence draws its own accepted length K_i independently, from zero (the very first drafted token is rejected) up to γ (every drafted token survives, plus a bonus token). Two sequences that entered a round at the identical position can leave it at four different positions. This is a genuinely different problem from the memory-allocation question this curriculum's KV-cache chapter already answers — that chapter is about where a variable-length cache physically lives in GPU memory; this is about the fact that, mid-batch, a speculative round writes cache content the target model's own verdict may then partially invalidate.

Concretely: the draft model, running sequentially for γ steps, computes and writes key/value vectors for all γ of its proposed positions as it goes — it has no way to know in advance which of its own guesses the target model will accept. When the target model's single verification pass finds the first rejection at position k + 1, every K/V entry the draft model wrote for positions past that point was computed from a token sequence that is no longer valid; committing it to the sequence's persistent cache would let a rejected token's key/value vectors influence attention for tokens generated afterward, silently breaking the exact-distribution guarantee the accept/reject test exists to protect. Those entries have to be identified and discarded — not merely left unread, but excluded from whatever block-table or contiguous region the sequence's persistent cache occupies — before the next round can begin. A serving engine's iteration-level scheduler, the same style of per-step rescheduling this curriculum covers for continuous batching, has to track a separate "committed length" per sequence every round, not just a single shared batch-wide step counter, precisely because speculative decoding makes every sequence's cache grow by a different, only-known-after-verification amount.

Diagram: ragged commits and the throughput crossover

Misconception: "More draft tokens per round always recovers a fading speedup"

Faced with the throughput collapsing as concurrency rises, the instinctive fix an engineer reaches for is to increase γ — more drafted positions per round should mean more chances to accept, and the earlier single-sequence treatment did show E[T] rising with γ. That instinct is backwards once the batch is already past its compute-bound crossover, and the numbers make the direction of the error concrete. Holding the fintech's setup fixed at B = 128 (already well past the 30.6 crossover) and raising γ from 4 to 6 to 8 to 12, the speedup does not recover — it falls further, from 0.666× to 0.537× to 0.444× to 0.323×. The reason is that E[T] grows only slowly with γ once α < 1 — it is bounded above by 1/(1 − α), a ceiling the geometric decay in the acceptance-length distribution approaches quickly — while the verification pass's compute time, once compute-bound, grows linearly in γ with no ceiling at all. Past the crossover, every additional drafted position is close to guaranteed extra GPU time and only marginal extra accepted tokens. The correct lever at high concurrency is not a wider draft window; it is either a smaller γ (shrinking w pulls the crossover back up, since B*(w) is inversely proportional to w), a better-aligned draft model (raising α so more of E[T]'s ceiling is actually reached at the γ you already have), or turning speculation off entirely above the batch size where the staging benchmark's assumptions stop holding.

Active recall

Attempt every question before reading its answer.

1. Using the idea of arithmetic intensity, explain in one or two sentences why speculative decoding's γ "free" positions per pass stop being free as batch size B grows.
2. A different inference GPU has FLOPS_peak = 200 TFLOPS (fp16) and BW = 3000 GB/s, and a team runs γ = 6. Compute the ridge point R and the verification crossover batch size B*.
3. A weaker-matched draft model drops α from 0.75 to 0.6, with γ = 4 unchanged. Recompute E[T]. Does this move the throughput crossover batch size B* = 30.6 at all? Why or why not?
4. The fintech's platform team, migrating for memory savings, quantizes the target model from fp16 (d = 2 bytes) to int8 (d = 1 byte) on the same A100, which also roughly doubles that GPU's dense tensor-core throughput for int8 versus fp16 (624 TOPS int8 versus 312 TFLOPS fp16). Holding N, N_draft, γ = 4, and α = 0.75 fixed, trace the full ripple: what happens to the ridge point R, to the crossover B*, to the memory time for one target pass, and to the speedup at a given batch size B?
5. In Panel A of the diagram, Sequence D has K = 0 — its very first drafted token is rejected. By how many tokens does its KV cache actually grow this round, and what has to happen to the key/value vectors the draft model already computed for its positions 2 through 5 before the next round can start?
6. A colleague, seeing throughput fall at high concurrency, proposes doubling γ from 4 to 8 while the batch stays fixed at B = 128. Using the worked numbers in this chapter, is this the right fix? What should they do instead?

Answers

1. A verification pass's arithmetic intensity is 2Bw/d, which grows with both batch size B and window width w = γ + 1; while that intensity sits below the GPU's own compute-to-bandwidth ridge point, the pass is limited purely by the fixed cost of streaming weights once, so extra positions really do ride along free, but once B grows large enough to push the intensity past the ridge, the pass becomes compute-bound and its latency starts scaling with B·w, so every extra drafted position now costs real, non-free GPU time.

2. R = 200×10¹²/3000×10⁹ = 66.67 FLOPs/byte. B* = R·d/(2(γ+1)) = 66.67×2/(2×7) = 133.33/14 ≈ 9.52. This GPU's lower compute-to-bandwidth ratio pushes the crossover far lower than the A100's 30.6 — on this hardware, even a modest batch of 10 concurrent sequences is enough to make the verification pass compute-bound.

3. E[T] = (1 − 0.6⁵)/(1 − 0.6) = (1 − 0.07776)/0.4 = 0.92224/0.4 = 2.3056 tokens/round, down from 3.051 at α = 0.75. No, B* does not move: B* = Rd/(2(γ+1)) depends only on the hardware's ridge point, the precision d, and γ — nowhere does α appear in that formula. α changes how many tokens a round is worth once you know its cost, not whether the verification pass itself is memory- or compute-bound; a worse draft model makes the win smaller everywhere but does not shift the batch size at which the win starts shrinking.

4. The ridge point doubles: R_int8 = 624×10¹²/2039×10⁹ ≈ 306.03, exactly 2× the fp16 value. And because B* = Rd/(2(γ+1)) depends only on the product R·d, and d has halved to 1 byte, that product is unchanged — 306.03×1 = 306.03, identical to fp16's 153.02×2 = 306.03 — so the verification crossover stays at B* ≈ 30.6. That much survives regardless of the draft model, because B* is a target-model-only quantity: on this hardware int8 tensor cores compute roughly twice as fast, and that speedup exactly cancels the halved byte count in the ratio that sets the crossover. But the question quantizes only the target model, holding the 7B draft model at fp16, and that detail matters for everything downstream of B*, because a full speculative round is not just a verification pass: T_round = gamma * T_draft_step + T_verify. Only T_verify, built from the target model's memory and compute costs, shrinks under quantization; T_{draft_step} is computed from the draft model's own, still-fp16 memory and compute costs and does not shrink at all. Re-running round_costs() with the target model's inputs quantized (FLOPS_peak = 624e12, d = 1) but the draft model's left at fp16 (FLOPS_peak = 312e12, d = 2) gives thr_naive that doubles exactly at every B, as expected (14.6→29.1 at B = 1, 1864.2→3728.5 at B = 128), but thr_spec that grows by less than 2× (31.7→49.4 at B = 1, 1241.1→2282.9 at B = 128), because the unshrunk γ·T_{draft_step} term now makes up a larger share of T_round than before. The speedup therefore gets worse at every batch size, not "completely unchanged": 1.695× at B = 1 (versus 2.179× at matched precision), 0.885× at B = 81 (versus ≈1.00×), 0.612× at B = 128 (versus 0.666×). The break-even batch size itself moves down, from B ≈ 81.1 to B ≈ 68.9, because thr_naive is racing ahead on a faster, quantized target model while thr_spec is held back paying the same fp16 draft-step cost as before. Quantizing only the target model makes the whole system faster in absolute terms, but it narrows, not preserves, the concurrency window in which speculative decoding beats plain batched decoding — the memory savings only carry through to the speedup ratio if the draft model is quantized too.

5. Sequence D's cache grows by exactly one token: the resampled correction token at position 101, the only slot that was actually verified and accepted (as a correction) before the round ended for this sequence. The key/value vectors the draft model computed for its proposed positions 102 through 105 (and the bonus slot) must be discarded, not written into Sequence D's persistent cache: they were computed by extending a chain that started with the very token the target model just rejected, so committing them would let a token the target model disagreed with keep influencing attention in every future step, corrupting the exact-distribution guarantee the whole accept/reject construction exists to preserve.

6. No. At B = 128 with γ = 4, the round is already well past the 30.6 crossover, so the verification pass is compute-bound; going to γ = 8 was already shown in this chapter's misconception section to drop the speedup further, from 0.666× to 0.444×, because E[T] is bounded above by 1/(1 − α) and rises only slowly with γ while compute-bound latency rises linearly with γ and has no such ceiling. The colleague should do the opposite — shrink γ to pull the crossover back above the batch size actually being served, improve α with a better-aligned draft model, or disable speculation above whatever batch size this GPU's own B* turns out to be — not widen the draft window further.

Think About It

Think about this: How would you explain inference optimization: kv cache, speculative decoding, and batching strategies to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where inference optimization: kv cache, speculative decoding, and batching strategies is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting inference optimization: kv cache, speculative decoding, and batching strategies to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind inference optimization: kv cache, speculative decoding, and batching strategies, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.