GPU Lab CHAPTER 12
What limits an inference step?
How do a growing cache, the invocation phase and batch size change an inference budget?
Start with the operation
PyTorch
import torch
import torch.nn.functional as F
# Three cached tokens, each with four components.
k_cache = torch.arange(12, dtype=torch.float32, device="cuda").reshape(3, 4) / 10
v_cache = torch.arange(12, dtype=torch.float32, device="cuda").reshape(3, 4)
q_new = torch.tensor([[1., 0., -1., 0.]], dtype=torch.float32, device="cuda")
k_new = torch.tensor([[0.2, 0.3, 0.4, 0.5]], dtype=torch.float32, device="cuda")
v_new = torch.tensor([[12., 13., 14., 15.]], dtype=torch.float32, device="cuda")
k_cache = torch.cat((k_cache, k_new), dim=-2)
v_cache = torch.cat((v_cache, v_new), dim=-2)
out = F.scaled_dot_product_attention(
q_new, k_cache, v_cache, dropout_p=0.0, is_causal=False)
print(k_cache.shape, out)A decoder reuses keys and values from earlier tokens. This complete CUDA-device PyTorch example starts with three cached tokens, appends one new key and value, and computes one new query's weighted result over all four tokens. Each row is one token and each column is one component, as in the previous chapter. reshape(3, 4) arranges twelve numbers into three rows of four. dim=-2 selects the first dimension of these two-dimensional tensors, so concatenation appends rows.
F is the short name given to PyTorch’s functional operations by the import. Its scaled_dot_product_attention function performs the score, softmax and weighted-value calculation from the previous chapter. dropout_p=0.0 keeps every attention weight. is_causal=False is intentional: the sole query is the newest token, so all four supplied keys are permitted. For a multi-query prompt, a causal mask is needed to hide future tokens. This example receives Q/K/V directly; it contains no model weights, projections or token generator.
torch.cat creates new tensors and copies existing entries. Production caches can instead reserve storage and write new positions. The small code illustrates cache contents, not efficient cache management. PyTorch handles allocation and implementation selection. The browser budgets below use a separate, explicitly synthetic decoder; they do not estimate this four-component PyTorch example.
Run the examples below in your browser with no installation or GPU required. The browser simulates the calculation; it does not execute PyTorch or GPU code. To run the PyTorch snippet itself, use a Python environment with PyTorch and an NVIDIA GPU.
Idea 1 of 3
Reserve space for past tokens
Definitions
- KV cache
- Stored keys and values from earlier tokens, reused by later attention queries.
- Decode
- Processing one new token per sequence while reusing earlier cache entries.
- Batch
- Independent sequences processed together.
- Invocation
- One execution of the work being counted: a whole prompt or one decode step for each sequence.
Start with the cache just before one decode invocation. First count storage, without asking about time. The cache payload is 2 × layers × batch × cached tokens × KV width × bytes/value. The first factor of two stores both keys and values.
The synthetic decoder stores keys and values for 32 layers, with 1,024 components in each key or value and two bytes per component. A layer is one stage of the model. These are teaching dimensions; the opening code uses one set of four-component values.
Here batch stays at one. The selected token count S is the cache length before a decode invocation; the invocation appends one entry per sequence. Payload capacity excludes any reserved unused slots and allocator metadata.
Try it
Predict the cache payload when the cached length changes from 512 to 2,048 tokens. Run at 512, change only the length, and explain the factor. Distinguish the existing payload from the new bytes appended.
Predict the result, then press Run.
512 cached tokens per sequence
1 token per sequence
513 cached tokens per sequence
- Existing KV payload
- 67,108,864 bytes
- New KV bytes written
- 131,072 bytes
- KV payload after append
- 67,239,936 bytes
2 × 32 layers × 1 sequence × 512 tokens × 1024 KV components × 2 bytes = 67,108,864 existing bytes. This is payload capacity, excluding spare allocation.
Capacity is a storage question; the separate phase experiment introduces time bounds.
Idea 2 of 3
Compare two invocation phases
Definitions
- Prefill
- Processing the prompt tokens and building their cache entries.
- Lower bound
- An ideal minimum implied by the counted work and assumed rate; it is not a prediction of measured elapsed time.
- FLOP
- One floating-point arithmetic operation in the work estimate.
Now keep batch = 1 and S = 512. For prefill, S means the number of prompt tokens processed from an empty cache. For decode, S means the number already cached, followed by one new token. These invocations do different amounts of work.
This synthetic budget assumes 7 billion learned weight values (parameters), 32 layers, hidden width 4,096, KV width 1,024 and two bytes per stored value. Hidden width sizes the model’s internal vectors in the operation estimate; KV width is the number of cached key or value components per token per layer. It is not the shape of a named model. One invocation reads the 14 GB weight payload once, with ideal reuse across its batch. Activations, communication, intermediate tensors, allocation, cache-copy costs and other overheads are excluded.
Decode reads the S-token cache and writes one new entry per sequence. Prefill writes the full S-token cache; this model counts no previous-cache read. The prefill operation estimate uses full square attention, without causal triangular savings. Dense-weight work is approximately 2 × parameters × tokens; attention adds 4 × layers × batch × S × hidden width for decode, and another factor S for prefill.
The selected ideal rates are 1,000 decimal GB/s (1012 bytes/s) and 100 TFLOP/s (1014 floating-point operations/s). These are teaching inputs, not measured specifications. Divide counted bytes by bandwidth and counted operations by compute rate; take the larger time assuming ideal overlap. Real elapsed time needs a measurement with a matching scope.
Try it
Predict which bound is larger for decode, then run. Change only phase to prefill. Explain why the compute work changes even though the model reads the same weight payload per invocation.
Predict the result, then press Run.
decode
1
512 → 513
- Memory lower bound
- 14.06724 ms
- Compute lower bound
- 0.142684 ms
- Invocation lower bound
- 14.06724 ms
- Larger ideal bound
- memory
| Counted work per invocation | Amount |
|---|---|
| Weight reads | 1.40000e+10 bytes |
| Previous KV reads | 67,108,864 bytes |
| New KV writes | 131,072 bytes |
| Total counted transfers | 1.40672e+10 bytes |
| Arithmetic work | 1.42684e+10 FLOPs |
Selected accounting step: Read weights, 14 ms from 1.40000e+10 counted bytes. Component memory times share the same bandwidth; the final bound takes the larger of total memory time and compute time.
At the stated ideal rates, max(14.06724, 0.142684) = 14.06724 ms. This is an analytical lower bound, with no measured time supplied.
The larger lower bound identifies the tighter of these two ideal budgets. It does not diagnose an actual GPU bottleneck. This page supplies no timing measurements; uncounted work and implementation details need separate investigation.
Idea 3 of 3
Share a weight read across a batch
Return to decode with S = 512 cached tokens per sequence. Increasing batch produces more tokens per invocation. The ideal budget reuses each weight read across the whole batch, while each sequence needs its own cache.
This synthetic budget assumes 7 billion learned weight values (parameters), 32 layers, hidden width 4,096, KV width 1,024 and two bytes per stored value. Hidden width sizes the model’s internal vectors in the operation estimate; KV width is the number of cached key or value components per token per layer. It is not the shape of a named model. One invocation reads the 14 GB weight payload once, with ideal reuse across its batch. Activations, communication, intermediate tensors, allocation, cache-copy costs and other overheads are excluded.
The selected ideal rates are 1,000 decimal GB/s (1012 bytes/s) and 100 TFLOP/s (1014 floating-point operations/s). These are teaching inputs, not measured specifications. Divide counted bytes by bandwidth and counted operations by compute rate; take the larger time assuming ideal overlap. Real elapsed time needs a measurement with a matching scope.
Compare both milliseconds per invocation and the bound divided by the number of new tokens. The latter is a budget per produced token, not a user's time to receive the next token. More batch members also need more cache capacity.
Try it
Predict which grows with batch: weight payload read, KV cache, or new tokens. Run at batch 1, change only batch to 4, and compare the invocation bound with the bound per new token. Explain the ideal reuse assumption.
Predict the result, then press Run.
1
1
67,108,864 bytes
14.06724 ms/token
- Memory lower bound
- 14.06724 ms
- Compute lower bound
- 0.142684 ms
- Invocation lower bound
- 14.06724 ms
- Larger ideal bound
- memory
| Counted work per invocation | Amount |
|---|---|
| Weight reads | 1.40000e+10 bytes |
| Previous KV reads | 67,108,864 bytes |
| New KV writes | 131,072 bytes |
| Total counted transfers | 1.40672e+10 bytes |
| Arithmetic work | 1.42684e+10 FLOPs |
Weight bytes stay fixed under ideal batch reuse; cache bytes and arithmetic work scale with batch. The per-token bound is not an individual request’s next-token latency or a measured throughput.
Optional: plan a real timing comparison
Use a real model and document its shapes, datatype, batch, context, attention implementation and cache allocation strategy. Choose whether timing covers one invocation or the whole request. Warm up, account for asynchronous CUDA execution, synchronise the chosen interval, and repeat measurements. A result below this page's bound would first require checking units, counted work, assumed rates and timing scope. The browser exercises run no PyTorch or GPU benchmark.
Try a new case
Apply it
Compare a cache for batch 4 with 128 tokens per sequence against batch 1 with 512 tokens per sequence, keeping layers, width and value size fixed. Which needs more cache payload? Does equal cache payload imply equal per-token work or measured latency?
Compare your answer
Run the example on your desktop
Download this chapter’s opening PyTorch example. See the complete example index and setup appendix for CPU PyTorch, the supported GPGPU-Sim programs and optional NVIDIA GPU execution.
The browser still requires no installation. Downloads are generated from the displayed source; CPU checks and simulator checks do not establish GPU performance.
References
Use these official references to check the concepts and code on this page.
- PyTorch: torch.cat (new tab)
Check how the opening example appends one token along the cache’s token dimension.
- PyTorch: Tensor.element_size (new tab)
Check the byte unit used in tensor-payload estimates.
- PyTorch: asynchronous CUDA execution (new tab)
For optional measurements, find why timing must account for queued GPU work.
These browser exercises do not run an installed PyTorch or CUDA toolkit. For optional local execution, choose documentation matching your installed tools. Links open online in a new tab; the questions and answer comparisons below work offline.
Use the documentation
Use the references to explain each answer, then apply it to this page’s example. No installation is needed.