Inference Systems · Accelerator Scheduling
How Key-Value Cache Bytes Set the Batch Ceiling
The byte arithmetic that decides how many sequences fit beside the weights, why the ceiling moves with context length, and what remains after it is reached.
A serving system’s maximum batch size is usually not a number anyone chose. It is the quotient of two quantities: the device memory left over after weights, activations, and working buffers are placed, divided by the per-sequence residency cost. For a long-context model the second quantity grows with every token a sequence generates, so the ceiling falls while the sequence is still running. Understanding that quotient explains most of the behaviour that otherwise looks like arbitrary throttling.
The residency of one sequence
For a decoder-only transformer with $L$ layers, $H_k$ key heads, head dimension $d$, and $b$ bytes per stored element, a sequence holding $T$ tokens in its cache occupies:
$$L \times 2 \times H_k \times d \times b \times T$$
The factor of two covers keys and values. With grouped-query attention, $H_k$ is the number of key heads, which is smaller than the number of query heads; with multi-head attention it equals the query head count. A representative shape uses 32 layers, 8 key heads, a head dimension of 128, and two-byte elements, which gives $32 \times 2 \times 8 \times 128 \times 2 = 131{,}072$ bytes, or 128 KiB, per token.
That figure is the one worth memorising, because everything else is multiplication. A sequence with a 4,096-token context occupies 512 MiB by that arithmetic. Four such sequences occupy 2 GiB. A model with 70 billion parameters at two bytes each occupies 140 GB of weight memory before any cache exists at all, which is why weight residency and cache residency compete for the same pool and why the competition becomes binding at long context.
The per-token constant is not a law of nature. Doubling the layer count doubles it; halving key heads halves it; quantising the cache to one byte per element halves it again. The last of these is a scheduling lever as much as a numerical one, since a cache stored in eight-bit elements roughly doubles the number of resident sequences for the same memory, at whatever accuracy cost the quantization scheme carries.
Dividing the pool
Total device memory is partitioned before the first request arrives. The useful decomposition is:
- Weights, fixed for the life of the process. For a model sharded across $k$ devices, each device holds $1/k$ of the weights.
- Activation working set, which scales with the maximum batch and maximum chunk of tokens processed in one step, not with context length.
- Kernel workspace, fixed but non-trivial for some attention and communication implementations.
- Cache pool, whatever remains.
- Fragmentation and allocator overhead, which is a percentage of the pool and, in naive implementations, can be much larger than a percentage.
If the pool is $M$ bytes and a sequence’s eventual residency is $T$ tokens times $c$ bytes per token, then the number of sequences that can be admitted at full context is $M/(Tc)$. Serving systems rarely admit at full context, because the eventual length is unknown. They admit against the requested maximum length, which is a declared parameter, or against a growing allocation that is extended as tokens are produced.
The distinction between those two policies is the difference between two very different ceilings. Admitting against the requested maximum is conservative and wastes the difference whenever a sequence stops early, which is most of the time. Admitting against actual growth risks running out of memory mid-generation, which forces preemption and either a swap to host memory or a recomputation of the cache.
Why the ceiling moves during a step
Three effects make the ceiling a function of time rather than a constant.
Growth. Every decode step appends one token’s keys and values to each running sequence. A batch of $B$ sequences consumes $Bc$ additional bytes per step. With the 128 KiB constant above, 32 sequences consume 4 MiB per step. Over a thousand steps that is 4 GiB, enough to exhaust a pool that was comfortable at admission.
Uneven lengths. Sequences in the same batch have different context lengths, so their residency differs by whatever ratio their lengths do. The pool is consumed by the sum, and the sum is dominated by the longest sequences. A single request at 100,000 tokens occupies more cache than a hundred requests at 1,000 tokens, and the scheduler cannot swap the one out cheaply because its cache is the most expensive thing in the system.
Allocator behaviour. A pool carved into contiguous per-sequence extents cannot be reused until the sequence ends. If the pool has $M$ bytes and the largest free extent is smaller than the next request’s need, admission fails even when total free bytes are sufficient. Fragmentation of this kind is bounded by the ratio of the largest allocation to the pool size, and it is the reason cache pools are managed in fixed-size blocks rather than per sequence. The block-table mechanism is covered in paged key-value residency and block tables.
The arithmetic in practice
A scheduler that wants to know whether it can admit one more request needs an estimate of that request’s cost, and the estimate has to be an upper bound or the system will overcommit. The calculation is mechanical:
def kv_bytes_per_token(layers, kv_heads, head_dim, elem_bytes):
return layers * 2 * kv_heads * head_dim * elem_bytes
def admissions(pool_bytes, per_token, max_len, block_bytes):
seq_bytes = per_token * max_len
blocks = -(-seq_bytes // block_bytes) # ceiling division
return pool_bytes // (blocks * block_bytes)
With a 128 KiB per-token constant, a 2,048-token maximum, 16 MiB blocks, and a 40 GiB pool, each sequence needs 256 MiB, which is 16 blocks, so the pool admits 160 sequences by this bound. If the same pool is asked to serve 8,192-token maxima, each sequence needs 1 GiB and the count falls to 40. Nothing about the compute units changed between those two cases; the ceiling is entirely a memory statement.
Two refinements matter. First, block rounding wastes on average half a block per sequence, which is negligible for long sequences and material for short ones. Second, the maximum length a request declares is often far above what it uses, so a scheduler that admits against the declaration will run a batch of 40 sequences when 160 would have fit. Production systems therefore combine a conservative admission bound with the ability to extend a sequence’s allocation as it grows, and rely on preemption when the extension fails.
What is left when memory is the limit
Once the cache pool is the binding constraint, additional compute capacity on the device does nothing for throughput. This has a counterintuitive consequence for the batch-size-versus-latency trade. In the bandwidth-bound regime, doubling the batch roughly halves the per-token cost because the weight read is amortised over twice as many sequences. In the memory-bound regime, doubling the batch is simply impossible, and the per-token cost is fixed by the arithmetic of the quotient.
The escape routes are all memory-side. Quantising the cache reduces $c$ and raises the ceiling proportionally. Sharing cache between requests with a common prefix removes duplicated residency entirely, which is what prefix sharing and cache reuse exploits. Reducing the number of key heads through grouped-query attention lowers the constant at training time. Shortening the maximum context the scheduler will admit bounds the worst case at the cost of refusing work.
There is also a coupling with the weights that is easy to miss. Tensor parallelism splits weights across devices, so each device’s pool has more room, but it also splits key heads, so the per-device cache cost per token falls by the same factor. Pipeline parallelism splits layers, so each device holds only some layers’ caches and the per-token cost falls with the stage count. Both raise the per-device ceiling, and both change the cost model in ways that interact with the interconnect; the residency side of that trade belongs with tensor-parallel residency and interconnect pressure.
Edge cases that break the estimate
Cache quantisation asymmetry. Storing keys in one format and values in another doubles the bookkeeping without halving the cost, and a scheduler that assumes a single element size will mispredict the ceiling.
Sliding-window attention. A model that attends only to a bounded window can evict old keys, so its residency is capped at the window length rather than growing with context. The per-token constant still applies, but $T$ is bounded, and the ceiling stops falling.
Speculative decoding. Proposing $k$ tokens per step requires cache space for the proposals before they are accepted, so the effective residency is the accepted length plus up to $k$ extra positions per sequence. The pool must be sized for the transient, not the steady state.
Multi-tenant reservation. A per-tenant memory quota is a scheduling constraint that binds before the physical pool does, and the resulting ceiling is a policy number rather than a hardware one.
The general lesson is that the batch ceiling should be computed from the pool that actually remains, with block granularity, against a length bound the scheduler is prepared to enforce. Systems that treat batch size as a tunable parameter independent of context length end up with a parameter that works for one workload shape and fails for the next.
References
- Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP 2023
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al., NeurIPS 2022
- CUDA C++ Best Practices Guide — NVIDIA