Serving and Training a Transformer

David I. Inouye

Dynamic Content · Last Updated: November 2, 2026

Serving a model means managing user latency and request memory

Prefill and decode → why K/V can be reused → memory → context growth.

Then: data movement, throughput and inference diagnosis.

First serve the trained model; afterward, consider training it.

Prefill reads the prompt; decode extends it one token at a time

Each model pass has inputs below and selected outputs above; passes proceed left to right. A prompt of three tokens is processed together during prefill, yielding report; the gray Select label is separate from the bold token. The report token is processed in the next decode step yielding today; today is then processed yielding EOS. Arrows show the sequential dependency.

Work	Inputs already known	User-facing measure
Prefill	All prompt tokens	Time to first token
Decode	One new input token per step	Time per subsequent token

Prefill saves per-layer keys and values; each decode step reuses and extends them.

Earlier tokens stay fixed; attention needs only their keys and values

At one transformer layer, old tokens are not recomputed. Their saved keys and values are read by the new token's attention. The new token computes Q K V; new K V are appended to the layer cache and new Q attends to that cache. Repeat at every layer.

Causal masking: appending a token cannot change earlier token states.
Attention mixes tokens; MLPs, normalization and residual paths act within each token.

Keep old K/V at every layer; discard old queries and other intermediate states.

Five minutes. This explanation applies to the standard causal transformer already taught, with token-wise normalization and MLPs, fixed weights, unchanged prefix/positions and deterministic inference. At a layer, the only input-dependent route from earlier tokens into the new token is attention, through their keys and values. Earlier queries were used for earlier outputs; they are not needed for the new row. Earlier MLP outputs were needed to construct later-layer caches during prefill, but need not all be stored once the per-layer K/V are available. Never imply only the first layer is cached. The new token still executes the whole stack, including MLPs; caching avoids rerunning the old tokens, not reading their K/V. At one head: a_l = softmax(K_1:l q_l / sqrt(d_k)), o_l = V_1:l^T a_l, with column vectors and transposed rows in K/V. Append the new token’s K/V before its causal attention so it can attend to itself. RoPE keys retain their original position rotation. Edits to the prefix, reassigned positions or changed weights can invalidate cache entries. Training backward and nonstandard cross-token operators are outside this reuse claim. [Sources] source:huggingface-kv-cache-4-57-1; source:vaswani-attention-2017.

Inference keeps weights and K/V, plus temporary workspace

Memory term	Why it stays or grows
Model weights	Shared across requests; size depends on parameters and precision
K/V at every layer	Retained for later tokens; grows with prefix length and active requests
Temporary activations/workspace	Needed for current computations; reused or freed afterward

We cache the history attention will read—not every activation produced for that history.

A longer context grows both the cache and the attention work

Quantity being counted	If context length doubles: \(L\rightarrow 2L\)
Compute: dense attention over a full prompt	About 4× as many query–key pairs
Compute: cached attention for one new token	About 2× as many query–key pairs
Memory: persistent K/V cache	2× as many stored K/V entries

Fixed model, precision and request count. Operation counts are not measured latency.

Conventional multi-head cache bytes \(\approx 2NBLdb\).

Synthetic example: \(N=12\) layers, \(B=1\), \(L=1024\), \(d=512\), \(b=2\) bytes → 24 MiB.

Fewer memory transfers can make the same attention faster

Conventional attention sends full score and weight arrays from on-chip compute to GPU main memory and reads them back. FlashAttention instead loads blocks, computes score tiles and running softmax/output in fast on-chip SRAM, and writes the final output without materializing the full attention array in main memory.

Same dense attention; fewer large memory round trips.

Tiling reduces intermediate storage and traffic; the pairwise arithmetic remains quadratic.

Three minutes. Connect to the previous context-scaling table: quadratic pairwise work does not require storing every pair score at once. FlashAttention tiles Q/K/V and maintains the normalization needed to combine blocks into the same mathematical attention result. Floating-point evaluation order can change rounding; exact does not promise bitwise equality. The original dense algorithm still performs quadratic pairwise arithmetic at fixed width. Its block-sparse extension is a separate approximation, not the claim here. Focus here on forward attention during prompt prefill. The same algorithm also supports training; its backward recomputation belongs to the later training discussion. Do not quote a universal speedup or assume every transformer phase is memory-bound. [Sources] source:dao-flashattention-2022, sections 1–3.

Higher throughput need not mean a faster response for each user

Measure	Main dependencies—and why
Time to first token	Queue/load + prompt length: wait for service, then prefill the prompt
Time per subsequent token	Model size + context + active batch: each step reads weights/K/V and computes the new token
Aggregate output tokens/sec	Batch size + hardware utilization: process several requests together to share weight access and parallel work

Larger batches can raise throughput while adding queueing, cache memory and per-step work.

More output tokens mean more sequential decode steps—and a longer total response time.

Three minutes. Batching lets multiple requests share a model pass, potentially improving hardware utilization and amortizing weight access. Each request still has its own prefix/cache; the cache ledger scales with active batch size. Waiting to assemble a batch can add queueing delay. These are possible tradeoffs, not a claim that every batching implementation makes every latency measure worse. A service can improve aggregate tokens/sec without improving one user’s time to first token or inter-token delay. Hardware, prompt lengths, output lengths and offered load must be comparable; no benchmark or measured speedup is asserted here. Low-batch decode is often constrained by memory traffic, but bottlenecks vary. [Sources] source:huggingface-kv-cache-4-57-1; source:dao-flashattention-2022.

Which memory grows when the prompt grows?

Fixed weights and request count; cached inference now runs out of memory on longer prompts.

What should we inspect first?

A. Parameter gradients and optimizer state

B. KV cache and temporary activations

C. The number of learned weights

B. Prefix-dependent memory grows even though the model’s weights do not.

Compare task quality, user latency, throughput and peak memory before choosing a fix.

Three minutes including voting. A recalls gradient-based learning from Arc 5, but normal inference does not maintain parameter gradients or optimizer state. C confuses a fixed model with variable request state. B names the likely growing terms without claiming the cache is always the sole cause: measure the allocations. Reducing context or concurrency changes the service offered; an exact memory-efficient attention implementation changes intermediate storage, but does not eliminate persistent KV entries. Tie the diagnosis back to choosing an optimization for a measured bottleneck. Close the inference section by distinguishing architectural access, held-out prediction and practical task utility; the existing RNN/bigram experiment is measured, while this transformer has no trained comparison yet. [Sources] source:huggingface-kv-cache-4-57-1; source:dao-flashattention-2022; source:pytorch-activation-checkpointing-2025.

Training adds backward computation and parameter updates

Memory ledger → precision and recomputation → sharing state across GPUs.

Serving keeps weights fixed. Training must also compute gradients and update them.

Full fine-tuning can need far more memory than serving

Qwen3-8B: approximately 8.2 billion parameters. Illustrative full-parameter AdamW setup:

Stored arrays	Bytes / parameter	Approximate memory
16-bit weights used for computation	2	16.4 GB
FP32 master weights for updates	4	32.8 GB
FP32 gradients	4	32.8 GB
Two FP32 Adam moment arrays	8	65.6 GB
Training state subtotal	18	147.6 GB

Inference: 16.4 GB weights + K/V + workspace.
Training: 147.6 GB state + saved activations + workspace.

Four minutes. Decimal GB (10^9 bytes), not GiB. The official model card reports 8.2B total parameters; this is approximate parameter-count arithmetic, not a measured run. The inference comparison assumes weights stored directly at 16 bits and no FP32 master copy. The illustrative mixed-precision full-parameter training setup retains 16-bit compute weights, FP32 master weights, FP32 gradients and two FP32 Adam moments. That is 18 bytes per parameter, or 147.6 GB (about 137.5 GiB); compute weights alone are 16.4 GB (about 15.3 GiB). Training state here is nine times inference weight storage, not nine times total inference peak memory. Other implementations may use different parameter/gradient precision or avoid a separate copy; 16-bit gradients would reduce this example to 131.2 GB. Activations depend on sequence length, microbatch, checkpointing and implementation, so no universal number is asserted. Without sharding/offload this subtotal alone exceeds an 80-GB GPU. LoRA/QLoRA train a smaller parameter subset and have a different ledger; do not present 147.6 GB as their requirement. Normal teacher-forced training uses saved activations, not a persistent autoregressive generation cache. This motivates precision, checkpointing and ZeRO next. [Sources] source:qwen3-8b-model-card, Model Overview; source:hf-model-memory-4-57-1, Anatomy of Model’s Memory.

Training can save memory by changing precision or recomputing activations

Choice	What changes	What to check
Mixed precision	Use lower precision for suitable operations; retain higher precision where needed	Numerical stability, quality and runtime
Activation checkpointing	Save fewer intermediates; recompute them during backward	Peak memory versus extra computation

Neither removes weights, gradients, optimizer state, and activations all at once.

Choose the technique for the memory term that is limiting the run.

Four minutes. Recall the training memory ledger. Mixed precision is not a blanket conversion of every tensor to a smaller format: operations and stored states can use different precisions; some schemes retain full-precision parameter or optimizer copies. Reduced precision can alter numerical results; compare convergence and held-out quality. Gradient scaling is useful for some low-precision formats but is not needed by every format. Activation checkpointing saves boundary inputs and reruns selected forward regions when backward needs their intermediates; it does not mean saving a model file to disk. Recomputed regions must preserve the required behavior, including random-state handling. More aggressive low-bit inference quantization is a separate choice and also needs quality evaluation. Keep API and distributed-training details out of this conceptual comparison. [Sources] source:pytorch-amp-recipe-2026; source:pytorch-activation-checkpointing-2025.

ZeRO reduces duplicated training state across GPUs

Ordinary data parallelism trains on different examples with a model replica on each GPU.

ZeRO (Zero Redundancy Optimizer) partitions state among those GPUs:

Stage	State partitioned across GPUs, cumulatively
1	Optimizer state
2	Optimizer state + gradients
3	Optimizer state + gradients + parameters

Less state stored per GPU; communication supplies the pieces when needed.

Three minutes, recognition-level only. Data-parallel workers contribute to training the same model using different examples; they must synchronize information for consistent updates. ZeRO avoids keeping every training-state array fully replicated on every worker. At stage 3 parameters are gathered as needed for computation and then repartitioned. This is not activation checkpointing and does not by itself eliminate activation/workspace memory. Communication, temporary gathered tensors and hardware affect the practical tradeoff; do not claim total memory divides exactly by GPU count or promise a speedup. No API/configuration, offload taxonomy or distributed-training derivation. The stages are cumulative. Close with three different levers: precision changes representation, checkpointing recomputes activations, ZeRO partitions duplicated state. [Sources] source:deepspeed-zero-tutorial-2026, ZeRO Overview.

Serving and training optimize different system outcomes

	Serving an LLM	Training / fine-tuning an LLM
Goal	Useful responses at acceptable latency, throughput and cost	Reach useful model quality within a compute/time budget
Memory	Weights + per-request K/V + workspace	Weights + gradients + optimizer state + saved activations
Levers	Cache K/V; manage context and batching; reduce attention traffic	Mixed precision; recompute activations; shard state with ZeRO

Identify the limiting resource, change the relevant term, then recheck quality and performance.

Two minutes. Summarize the whole systems section before moving to Arc 7. These are emphases, not exclusive technique lists: efficient attention and lower precision can help both serving and training. Caching reuses fixed history; checkpointing recomputes intermediates for backward; ZeRO shares state across workers. Each shifts memory, compute, communication or numerical behavior. Measure comparable workloads, including prompt/output length and offered load for serving, and learning progress as well as steps/second for training. A smaller allocation or faster kernel alone does not prove a better end-to-end system. [Sources] source:huggingface-kv-cache-4-57-1; source:dao-flashattention-2022; source:pytorch-amp-recipe-2026; source:pytorch-activation-checkpointing-2025; source:deepspeed-zero-tutorial-2026.