Output Latency Optimization

To begin optimizing your output latency, first understand the different time measurements in a response:

Time To First Token (TTFT): influenced by model loading, tokenization, prefill, and scheduling. Model math has not been used here at all for now.
Time Per Output Token (TPOT): time between consecutive tokens once streaming has begun.
End-to-end latency: $\text{TTFT} + (\text{No. of output tokens} \times \text{TPOT})$

Two more measurements are also worth considering:

Throughput: total tokens per second across all users simultaneously.
P95/P99 latency: the 95th and 99th percentile latency.

LLMs generate autoregressively: generate one token, then use it as feedback (that is, as input) to generate the next. Do the same for all future tokens. This serialization is what prevents parallelization.

Output token 57 depends on 56, no matter the amount of GPUs you allocate. Specifically, it depends on all the previous 56 K and V vectors, so to avoid recomputing them, cache them.

The size of that cache, per token, is given by:

\text{Per token cache cost} = 2 \cdot H \cdot D \cdot B \cdot L

where:

$2$ comes from the 2 vectors (K and V)
$H$ = number of attention heads
$D$ = dimension of each head
$B$ = bytes per number (2 for fp16, 4 for fp32 in PyTorch)
$L$ = number of transformer layers stacked

Check this example out — Llama-2-7B: 32 layers, 32 heads, dimension of each head = 128, fp16.

\text{Per token cost in cache} = 2 \times 32 \times 32 \times 128 \times 2 \text{ bytes} = 0.5 \text{ MB}

\text{Scaled to a 4000 token context} = 0.5 \times 4000 = 2 \text{ GB}

\text{128K context window costs} = 0.5 \times 128000 = 64 \text{ GB!}

At a 128K context window, the cache itself becomes larger than the model.

Now, there are 2 operations that must happen: the prefill and the decode.

Prefill is done in one single forward pass, because all tokens are present, so K and V pairs are calculated parallelly. This is a compute bound operation. After prefill, the first token generated takes time = TTFT.

Decode is the main problematic operation, and it is memory bound; memory bandwidth bound, more like. Model weights, and K and V, are loaded in the GPU's VRAM, but VRAM is not where compute happens, so the GPU must move bytes from VRAM to chip for calculations. This bandwidth is why decode is slow.

Take a concrete example: Llama 7b, fp16, on an A100 (batch size 1, one decode step).

The model's weights MUST move:

7 \text{ billion} \times 2 \text{ bytes} = 14 \text{ GB}

\text{A100 memory bandwidth} = 2 \text{ TB/s} \longrightarrow \frac{14}{2000} = \mathbf{7 \text{ ms}} \text{ to move the weights}

The arithmetic time required is 2 FLOPs per parameter per token, so the total is 14 GFLOPs:

\text{A100 fp16 throughput} = 312 \text{ TFLOPS} \longrightarrow \frac{14}{312{,}000} = \mathbf{0.045 \text{ ms}} \text{ to make the calculation}

∴ 7 ms to move, and just 0.045 ms to actually compute. The GPU is 99.4% idle, and TPOT works out to 140 tokens/s.

No matter how good your code is, this is hardware bound. VRAM bandwidth dictates how fast tokens stream out, not FLOPS.

∴ TTFT is a compute problem and TPOT is memory bound, and the two require different fixes.

Optimization Techniques

Batching: Utilize the idle GPU. Instead of computing one request per weight move, compute a batch of size $n$ (let $n = 32$ ):

32 \times 0.045 = 1.45 \text{ ms} + 7 \text{ ms to move} = \sim 8.5 \text{ ms total}

If you would not batch, it would be at 224 ms.

Static Batching: take N requests, process one by one. The problem here is that a small request may be stuck behind a big one, which gives bad tail-latency.

Chunked Pre-fill: Prefill is usually a long task, and it is also dependent on context size. If the context is huge, prefill becomes long and decode is stalled.

Fix: chunk the prefill tasks, and sprinkle them between decode steps.

∴ TTFT for one person is slightly bad, but TPOT for everyone else stays healthy.

Paged Attention: Borrowed from operating systems, so as to not waste cache memory. You allocate 4096 MB to a response, but the response only took 50 MB, so you wasted 4096 - 50 MB. That memory could be used elsewhere for more bandwidth.

Fix: page the VRAM. Chop it into KV blocks, and maintain a block table to map logical addresses to physical ones. Output needs more memory? Refer to the block table and use it. The program thinks it is using one continuous stretch of memory, but it is not.

PagedAttention is lossless, gives 2-4x higher throughput, and its only overhead is the memory mapping.

FlashAttention: A GPU has 2 memories, and they serve different purposes:

SRAM: small, fast, computes.
HBM (VRAM): large, does not compute.

Matrices are moved between SRAM and HBM because a GPU kernel can only leave its results in the VRAM. SRAM dies between kernel launches.

The steps to be performed are:

S = QK^T \quad \text{— raw score}

S = S/\sqrt{d} \quad \text{— scale}

P = \text{softmax}(S) \quad \text{— normalize}

O = PV \quad \text{— weighted sum of values}

Each of these steps is a separate kernel, so the intermediates make a round trip to memory: $S$ is written to HBM and read back for the softmax, then $P$ is written to HBM and read back for the final multiplication.

Flash Attention restructures the computation into tiles (small chunks in SRAM), fuses the steps, and eliminates the trips to HBM. This takes memory complexity from $O(N^2)$ to $O(N)$ .

Other optimization techniques:

Quantization: reduce the size of one number by reducing its bits. Cons: precision loss.

Eviction: keep track of each token's accumulated attention scores. If the cache is full, remove the tokens with the least score. Evicted tokens are lost forever. Cons: tokens that have been in memory for a long time dominate.

Hybrid memory: Move the cache to CPU RAM or SSD, which is very very slow. Since GPU compute is mostly serial, ship the data just in time, as the GPU will need it. LayerKV does this, keeping offload time ≤ prefill time, and it uses an SLO-aware scheduler.

All optimization techniques seen so far can be used if the model is yours. If you are using APIs, you can use the following optimization techniques instead:

Streaming: display the token as it becomes available, do not wait for the full output. This reduces perceived latency.

Prefix Caching: Cache the prefill of previous context so that prefill is not computed from the beginning again and again.

Semantic caching: Use a semantic cache to answer already answered queries. (Check out my project on semantic caching here.)

Model Routing: Check and send queries to models. Easy queries go to small models, hard queries go to large ones.

Shorten context.

Shorten reasoning: Models overthink, this is costly, so reduce it.

Parallelize tool calls: Independent calls happen independently.

Guide to optimization

The whole decision comes down to one question: is it your own model?

Own Model? YES
- Paged Attention, Batching, Flash Attention (use vLLM to get all three of these essentially for free)
- Chunked Prefill
- Once you implement the above two, only then implement eviction, quantization, and hybrid memory
Own Model? NO
- Streaming, Prefix Caching, Semantic layer
- Model routing, shorten context and reasoning
- Parallelize tool calls