Prefill, Decode, KV Caches, and Context Cost

“The model took three seconds” collapses several different queues and compute shapes:

admission queue
  → template / media preprocessing / tokenization
  → prefill the prompt
  → first token decode + selection
  → repeated decode + stream
  → terminal validation / persistence

A long prompt can dominate time to first token. A long answer can dominate total latency. Two requests with the same total tokens can put different pressure on compute and KV-cache memory. A model may advertise 128k context while the service cannot safely run more than a few such sequences concurrently—or quality may degrade before the maximum.

This chapter adds a checked planner to mosaic-ml. It counts hidden context components, reserves output space, estimates dense decoder K+V bytes from model geometry, applies device headroom, calculates a conservative memory-only concurrency ceiling, and separates time-to-first-token from decode cadence. The planner is useful because it is explicit; it is not a substitute for measuring the exact runtime.

Learning objectives

By the end you will be able to:

distinguish prefill from decode by work shape and user-visible latency;
explain why autoregressive runtimes cache keys and values but still compute a new query;
derive dense decoder KV bytes from layers, KV heads, head dimension, dtype, tokens and sequences;
distinguish multi-head, multi-query and grouped-query cache geometry;
count system, tools, retrieval, history, user and output reserve in one context budget;
explain why advertised context, admitted context and useful context are different;
distinguish dynamic, static, paged, sliding, quantized, offloaded and prefix caches;
calculate a memory-only concurrency ceiling with fixed memory and headroom;
state which allocator, architecture and runtime costs that estimate omits;
measure queue, preprocessing, prefill, first decode, inter-token latency and end-to-end latency;
explain continuous batching, chunked prefill and prefill/decode contention;
recognize that throughput-maximizing schedules can harm individual latency;
design tenant-safe cache keys/lifetimes and avoid cross-request leakage;
debug OOM, latency spikes, cache misses and context-overflow disagreements systematically.

Dependencies and completion artifact

This chapter depends on:

tokenization and exact rendered token counts;
causal self-attention shapes;
sampling and streaming;
checked arithmetic, backpressure and capacity planning;
model/runtime release identity.

Run:

cd rust/rust-ai-engineering
cargo run -q -p mosaic-ml --example prefill_decode_plan

The example reports:

{
  "context": {
    "input_tokens": 2048,
    "reserved_output_tokens": 512,
    "reserved_total_tokens": 2560,
    "unreserved_tokens": 1536
  },
  "kv_cache": {
    "bytes_per_token_per_sequence": 131072,
    "mib_per_reserved_sequence": 320,
    "requested_sequences": 8,
    "maximum_sequences_by_memory": 30,
    "admitted": true
  },
  "latency": {
    "time_to_first_token_ms": 450.0,
    "prefill_tokens_per_second": 5120.0,
    "output_tokens_per_second": 50.0,
    "mean_inter_token_latency_ms": 20.0,
    "end_to_end_ms": 3010.0
  }
}

The “30 sequences” value is only a memory arithmetic ceiling under the declared assumptions. It does not promise that 30 sequences meet latency, throughput, power, thermal, bandwidth or quality SLOs.

Mental model: two inference phases

Prefill

Prefill processes the input sequence and constructs the per-layer state needed for generation. Within causal attention, all prompt positions are known, so the runtime can process prompt tokens in parallel using matrix-shaped work, subject to the causal mask.

Application-visible properties:

work grows with prompt length and architecture;
long prompts can delay first output;
the runtime creates KV state for prompt positions;
prompt batches can use accelerator compute efficiently;
a large prefill can interfere with ongoing decodes unless scheduled/chunked carefully.

“Prefill tokens per second” is meaningful only with batch, length distribution, runtime/device and measurement boundaries attached.

Decode

After prefill, autoregressive generation repeatedly handles one new position per active sequence:

new token embedding
  → one-position forward through layers
  → attend to cached prior K/V
  → logits
  → processors + selection
  → append new K/V
  → stream token

Decode often has small matrix dimensions per sequence and repeatedly reads model weights/cache, making bandwidth, batching and scheduling important. More concurrent sequences can improve device utilization, but each sequence consumes cache memory and latency rises when the scheduler is overloaded.

The exact boundary around “first decode” varies by engine. This book defines time to first token as:

queue + preprocessing + prefill + first_decode_and_selection

Document your production definition, especially whether network serialization/flush is included.

Why cache keys and values?

For a single attention layer:

Q = XWq
K = XWk
V = XWv

At decode step t, earlier token representations do not need their K/V projections recomputed. Cache them per layer and append the new position’s K/V. The new query is still computed because it asks what the current position should attend to.

Conceptually at a layer:

cached K [past_positions, Hkv, Dhead]
cached V [past_positions, Hkv, Dvalue]

new token
  → q_new, k_new, v_new
  → append k_new/v_new
  → attention(q_new, all cached K, all cached V)

The cache contains model-derived representations, not ordinary text. It remains sensitive: representations can depend on prompts, retrieved data, tools and media. Treat cache isolation and erasure as data-security concerns.

KV caching trades memory for avoided recomputation. Disabling it can reduce persistent cache memory but makes generation repeatedly process prior context and is usually much slower for ordinary decoder generation.

Step 1: count the complete context

Raw user text is only one component. The companion requires:

pub struct ContextBreakdown {
    pub system_tokens: u64,
    pub tool_schema_tokens: u64,
    pub retrieval_tokens: u64,
    pub history_tokens: u64,
    pub user_tokens: u64,
    pub maximum_output_tokens: u64,
}

For the runnable case:

system          128
tool schemas    256
retrieval       512
history         640
user            512
-------------------
input          2048
output reserve  512
-------------------
reserved       2560
model maximum  4096
unreserved     1536

Every addition uses checked_add. Input and output reserve must be positive. The planner rejects reserved total above the model context before any cache-byte multiplication.

Why reserve output?

If the prompt consumes the whole context, there is nowhere to place generated tokens. Decide:

maximum output;
minimum useful output;
stop/EOS behavior;
what input components may be trimmed;
truncation priority and evidence/citation consequences.

Do not let a tokenizer silently truncate the end and then claim the model considered it. Return the actual retained windows and omissions in trace/provenance.

Advertised, admitted and useful context

Advertised context is a model/runtime limit under some configuration.

Admitted context is what the service accepts after output reserve, cache memory, batch, deadline, cost and policy.

Useful context is the length/composition that improves task quality on evaluation. More tokens can distract, dilute retrieval, change positions or exceed training distribution even when the runtime accepts them.

Measure all three separately.

Step 2: calculate dense decoder KV bytes

For a common decoder layout with equal key/value head dimensions:

bytes per token per sequence
  = 2                 K and V
  × layers
  × key_value_heads
  × head_dimension
  × bytes_per_element

Then:

cache bytes
  = bytes_per_token_per_sequence
  × cached_tokens_per_sequence
  × active_sequences

The example uses:

2 × 32 layers × 8 KV heads × 128 dimensions × 2 bytes
= 131,072 bytes/token/sequence
= 128 KiB/token/sequence

At 2,560 reserved tokens:

131,072 × 2,560
= 335,544,320 bytes
= 320 MiB per reserved sequence

Eight requested sequences reserve an estimated 2.5 GiB of dense K/V.

The Rust implementation:

[
    2,
    self.layers,
    self.key_value_heads,
    self.head_dimension,
    self.element_bytes.get(),
]
.into_iter()
.try_fold(1_u64, |total, value| total.checked_mul(value))
.ok_or(ServingPlanError::ArithmeticOverflow)

This is checked integer arithmetic, not an allocation. It rejects zero geometry and overflow.

MHA, MQA and GQA

With ordinary multi-head attention, K/V head count may equal query head count. Multi-query attention shares one K/V head across query heads. Grouped-query attention uses fewer K/V heads than query heads.

For identical layers, head dimension and dtype, 32 KV heads require four times the dense KV bytes of eight KV heads. The companion test proves the ratio. This is one reason architecture config—not parameter count alone—matters to serving capacity.

Formula limitations

The estimate deliberately omits:

allocator/page/block rounding and metadata;
fragmentation and reserved-but-unused blocks;
tensor alignment;
quantization scales/metadata;
cross-attention caches;
sliding/chunked/local-attention limits;
state-space/recurrent architecture state;
latent/compressed cache representations;
beam/parallel candidate sharing or duplication;
prefix sharing;
host offload/staging buffers and transfer workspaces;
pipeline/tensor-parallel partition details;
speculative draft/target cache behavior;
runtime-specific temporary buffers.

Keys and values can also have architecture-specific dimensions. Inspect the exact model/runtime. The formula is a first plan and a cross-check against measured allocated/resident memory.

Step 3: turn bytes into an admission ceiling

The example device contract:

device capacity       16.0 GiB
headroom                10%
usable                  14.4 GiB
model resident           4.0 GiB
runtime workspace        1.0 GiB
available for KV         9.4 GiB
KV per reserved seq      0.3125 GiB
memory-only maximum     floor(9.4 / 0.3125) = 30

Headroom covers measurement error, allocator behavior, transient buffers and operational safety; 10% is an example, not a universal value. Profile high-water marks under realistic worst cases.

The companion separates fixed and variable memory:

pub struct CapacityInputs {
    pub requested_sequences: u64,
    pub model_resident_bytes: u64,
    pub runtime_workspace_bytes: u64,
    pub device_capacity_bytes: u64,
    pub headroom_basis_points: u64,
}

It scales capacity by basis points without overflow-prone capacity * 10_000, requires fixed memory to fit usable memory, and calculates:

maximum_sequences_by_memory
  = floor((usable - fixed) / kv_bytes_per_reserved_sequence)

CapacityPlan::ensure_admitted turns an over-capacity plan into:

requested 31 sequences,
but the conservative KV-memory plan admits at most 30

Dynamic vs static reservation

The example multiplies by reserved_total_tokens (2,560), appropriate for an admission reservation that can allocate cache as tokens grow. A static cache may allocate the full configured shape—here 4,096—even when the request intends fewer tokens. Substitute the actual allocated length for the cache strategy.

Dynamic growth reduces upfront waste but can introduce allocation/fragmentation/compile-shape trade-offs. Static shapes can help compilation/kernel predictability while consuming more memory.

Memory ceiling is not an SLO ceiling

At 30 sequences the device may fit yet decode cadence may be unacceptable. Admission takes the minimum of:

memory limit
compute/throughput limit
latency-SLO limit
queue/deadline limit
tenant quota
cost/power policy
runtime hard limit

Measure at increasing concurrency with real prompt/output distributions. Find saturation before production finds it for you.

Cache strategies and trade-offs

Dynamic cache

Grows as positions are added. It can better match actual length but may interact with allocations, fragmentation and compilation. Current Transformers documentation uses a dynamic cache as the default for many models.

Static cache

Preallocates a maximum shape. It can enable predictable addresses/shapes and some compilation paths, but reserves capacity for unused positions. Batch/context maxima must be chosen deliberately.

Paged/block-managed cache

Allocates token blocks/pages and maps logical sequences to physical blocks. The PagedAttention work targets wasted reservation/fragmentation and supports sharing patterns such as common prefixes or parallel candidates. Page size creates its own internal waste and metadata/kernel complexity.

“Paged” is not magic free memory. Capacity still includes allocated blocks, partial final blocks, metadata, copy-on-write, scheduler reservations and runtime workspaces.

Sliding/chunked/local attention

Some layers attend only a bounded recent window/chunk, so cache growth may stop for those layers. This changes model behavior/architecture and cache formula. It is not a generic trick to evict old tokens from a full-attention model without quality consequences.

Quantized cache

Stores K/V at lower precision with scales/metadata. It reduces memory and may enable longer or more concurrent generation, at possible conversion cost, kernel limitations and quality change. Evaluate the complete path.

Offloaded cache

Moves some cache to host memory and transfers it as layers execute. It trades scarce device memory for host capacity and interconnect traffic/latency. Hugging Face documents offloaded cache as a memory-saving option with possible throughput degradation.

Prefix cache

Reuses K/V for an identical token prefix under the same complete model/execution contract. Cache key must bind:

model/adapters/quantization/runtime semantics;
exact token IDs and position behavior;
template/tools/system/policy;
modality processor outputs;
tenant/authorization/data-policy domain;
cache algorithm version.

A hash match on prompt text alone is unsafe. Cross-tenant prefix sharing can leak presence/timing or serve unauthorized derived state. Prefer tenant/policy isolation unless a reviewed threat model and content classification permit broader sharing.

Step 4: measure latency by phase

The companion observation:

pub struct LatencyObservation {
    pub prompt_tokens: u64,
    pub generated_tokens: u64,
    pub queue_ms: f64,
    pub preprocessing_ms: f64,
    pub prefill_ms: f64,
    pub first_decode_ms: f64,
    pub remaining_decode_ms: f64,
}

It requires finite nonnegative durations, positive prompt/output counts, positive prefill and first decode, and consistent remaining-decode semantics.

For the example:

queue            20 ms
preprocess       10 ms
prefill         400 ms
first decode     20 ms
----------------------
TTFT            450 ms

remaining 128 token intervals = 2,560 ms
end to end                         3,010 ms

Metrics:

prefill throughput = 2048 / 0.400 = 5120 prompt tokens/s

output throughput = 129 / (0.020 + 2.560)
                  = 50 output tokens/s

mean inter-token latency = 2560 / (129 - 1)
                         = 20 ms

The first token has no prior streamed-token interval, so mean ITL uses generated_tokens - 1.

Avoid metric ambiguity

State whether:

queue begins at gateway receipt or model scheduler;
preprocessing includes retrieval/media/tokenization;
TTFT ends at token selection, serialization, server flush or client receipt;
output token rate includes first token and sampling;
end-to-end includes terminal schema validation/tool loop;
cached-prefix tokens count as prefill work or avoided work;
failed/cancelled requests are included.

Keep both service-side spans and client-observed latency for network/stream buffering diagnosis.

Scheduling prefill and decode

Static batching

Runs a group together until completion. Short requests can wait for a long one; finished slots waste capacity.

Continuous batching

At each scheduling step, finished sequences leave and waiting work can enter when token/cache budget permits. Current Transformers documentation describes admission as requiring both token budget and cache space. This improves utilization but makes scheduler policy central to latency and replay.

Prefill/decode contention

A large prefill can occupy compute while active decodes wait, producing a visible streaming pause. Always prioritizing decode can starve new requests and inflate TTFT.

Policies include:

separate prefill/decode token budgets;
chunk long prefills and interleave decode;
deadline/age/fairness priority;
length buckets;
separate workers/devices for phases;
limit maximum prefill per admission step.

Evaluate:

TTFT and ITL percentiles by input/output-length bucket;
queue age/starvation;
tokens scheduled per phase;
cache utilization/fragmentation/eviction;
cancellations and wasted prefill;
tenant fairness;
throughput at SLO, not unconstrained maximum.

The scheduler is part of the harness/runtime release. Two schedulers over the same weights can produce very different product experience and cost.

Context editing and iterative conversations

Reusing cache across chat turns is valid only if new token IDs extend the cached prefix exactly. Re-rendering history can change:

whitespace/control tokens;
tool schemas;
system policy;
prior assistant reasoning/hidden tokens;
media placeholder expansion;
truncation;
position IDs.

If the new input is not a byte/token prefix, reuse must stop at the verified common prefix or the cache must rebuild. Never splice K/V based on message IDs alone.

When context overflows, options include:

remove low-value history under explicit policy;
retrieve relevant prior facts from a versioned memory store;
summarize with provenance and evaluate information loss;
trim tool schemas/routes;
use a longer-context model/runtime;
start a new session with a checked handoff artifact.

“Summarize everything” can erase constraints or inject incorrect memory. Treat context compaction as a model transformation with evals.

Failure investigations

Failure 1: OOM below advertised context

Collect:

model/runtime/device identity;
actual rendered prompt and output reservation;
batch/active sequence count;
K/V heads, dimensions, layers and cache dtype;
cache strategy/static allocation length/block size;
model/workspace/temporary/device memory high-water marks;
fragmentation/offload/speculative/beam settings.

Compare formula, runtime allocation report and measured peak. Common causes are static maximum allocation, concurrent sequences, full MHA rather than assumed GQA, different KV dtype, workspace spikes or allocator blocks.

Failure 2: TTFT spikes while total throughput looks good

Split queue, preprocessing and prefill. Inspect long prompts, batch token totals, scheduler delay and large prefills interrupting decode. Throughput aggregation can hide tail latency. Bound/chunk prefill, length-bucket metrics and enforce queue deadlines/fairness.

Failure 3: streamed output pauses every time a long request arrives

Likely prefill/decode contention. Correlate active decode sequences with admitted prefill tokens and ITL histograms. Test chunked prefill or phase budgets. Avoid “always decode first” without starvation tests.

Failure 4: prefix cache hit returns wrong behavior

The key omitted tokenizer/template/tools/model adapter/tenant or position semantics, or a hash was computed over text rather than exact IDs. Quarantine cache, compare complete identity, add negative key tests and isolate authorization domains.

Failure 5: context API says 4,096 but runtime rejects 3,800

The API count may exclude template/tool/special/media tokens or output reserve; runtime maximum may differ from catalog; positions/cache config may be lower. Count exact final IDs with the pinned processor and expose a structured breakdown like the companion.

Failure 6: cache offload prevents OOM but latency collapses

Host-device transfer exceeds overlap/interconnect capacity. Measure bytes transferred per layer, bandwidth, stalls, pinned memory and ITL. Reduce concurrency/context, use quantized/paged/device cache, or select hardware/model that fits. Offload is a trade, not free capacity.

Failure 7: one cancelled request leaves memory unavailable

Cache blocks/sequence state were not released on cancellation/error/disconnect, or release waits for a lost owner. Use RAII ownership, structured cancellation and terminal cleanup metrics. Inject cancellation at prefill and every decode stage; assert available capacity returns to baseline.

Security and privacy

KV cache is derived from sensitive inputs; isolate, zero/release according to threat model, and never expose it through debug endpoints.
Prefix-cache hit timing can reveal whether a prefix exists.
Tenant, authorization and policy identity belong in cache-domain decisions.
Enforce byte/media/token limits before expensive prefill.
Long prompts and maximum-output reservations are resource-exhaustion vectors.
A request should not reserve unlimited cache while waiting indefinitely in a queue.
Cancellation/disconnect must stop compute according to durable-work policy and release cache.
Avoid high-cardinality prompt/cache keys in metrics.
Detailed token/cache traces require protected storage and retention limits.
Memory reuse must not return uninitialized prior-tenant data to kernels or outputs.

Exercises

Exercise 1: architecture-specific geometry

Parse a pinned model config and derive query heads, KV heads, head dimension, layers and cache dtype. Compare the planner estimate with the runtime’s allocation report for 1, 512 and 4,096 tokens. Explain every difference.

Exercise 2: block rounding

Extend the planner with page/block size. Each sequence reserves ceil(tokens/block_tokens) blocks. Report internal waste and show distributions where a smaller page saves memory but increases metadata/scheduling work.

Exercise 3: dynamic versus static

Plan 100 requests with real length samples under dynamic reserved length and static maximum length. Compare admitted concurrency, unused bytes, allocation events and compile/kernel performance. Do not assume the memory winner is the latency winner.

Exercise 4: scheduler simulator

Simulate pending, prefilling, decoding and done states with per-step prefill/decode token budgets. Compare FIFO, decode-priority and chunked-prefill fairness. Report TTFT/ITL p50/p95 and starvation.

Exercise 5: cancellation recovery

Use a fake cache allocator with RAII sequence leases. Inject failure/cancellation after admission, mid-prefill, first token and later decode. Available blocks and active sequence count must return to baseline exactly once.

Exercise 6: safe prefix reuse

Create a key over model bundle, runtime semantics, tokenizer/template/tools/policy, tenant domain and exact prefix token IDs. Add negative tests changing one field at a time. Implement longest verified common prefix for an edited conversation.

Solution guidance

Treat parsed configuration as untrusted, validate cross-field dimensions, and compare actual runtime metadata. Differences may include alignment, scales, blocks or architecture-specific state.
Use checked ceiling division: (tokens + block - 1) / block only after proving addition cannot overflow, or tokens / block + (tokens % block != 0). Waste is allocated minus used token slots.
Replay one fixed length distribution and runtime release. Static cost uses allocated maximum, not intended output. Measure compilation/warm-up separately.
Use manual time and deterministic event ordering. A step admits only when both cache blocks and phase token budget fit. Report per-request, not only aggregate throughput.
A lease owns blocks and returns them on Drop; explicit completion consumes/marks it. Add idempotent release and stale-generation protection for distributed ownership.
Use length-prefixed canonical encoding/version and a cryptographic digest. Authorization/domain is checked before lookup and again before use; hash equality is not authorization.

Check your understanding

What work occurs during prefill that is reused during decode?
Why is the query for the new position not read from the old cache?
What does the leading 2 in the dense KV formula mean?
Why can GQA reduce cache memory?
Why must output tokens be reserved before admission?
Does a 128k advertised window imply one request should use 128k?
Why can 30 sequences fit memory but fail the latency SLO?
How can static cache use more memory than the planner’s intended-length estimate?
What causes a long prefill to harm already-streaming requests?
Why must tenant/policy identity participate in prefix-cache scope?

Answers:

Prompt K/V (and other architecture state) is constructed through the model layers.
It represents the current position and must be computed from the new hidden state.
Separate key and value tensors.
Fewer KV heads are stored/shared across more query heads.
Generated positions also consume context/cache; otherwise a full prompt leaves no generation capacity.
No. Admission and useful-quality limits may be lower.
Memory is only one constraint; compute/bandwidth/scheduling determine cadence and queues.
It may allocate the full configured maximum rather than actual/reserved total.
Prefill competes for accelerator compute with decode unless scheduling/chunking isolates it.
Derived cache state can contain unauthorized sensitive context and hit timing can leak presence.

Completion checklist

References

We can now account for model bytes, logits, selection and inference phases. Next we separate model capability, memorized/trained knowledge, instruction following and calibration so the harness routes and verifies behavior instead of treating fluent output as evidence of competence.