2024-05-31

High Performance LLM Serving

How to optimize LLM inference and create cost-effective, robust and custom solution

These are my notes on the GPU optimization workshop organized by Chip Huyen.

Sharan Chetlur talks about inference optimization techniques and how to achieve high performance while serving LLMs.

1. Majority of the apps are real-time(online systems) and they require acceptable latencies, should be highly accurate(helpful) and should be cost effective as deploying these large models can be expensive

2. These models will keep getting bigger and better. see chinchilla scaling laws.

3. So to optimize for these challenges we need a high-performant, robust and customizable solution

4. Quantization came up again here as its becoming necessary to do it(its an all-round win w/ just a bit of effort) as long as you maintain accuracy

5. So you can train llms in the BF16 dtype but inference can be done in lower precision dtypes like int8, int4 or even lower. this is post-training quantization and usually just needs calibration.

6. Post-training quantization can make compute faster, communications between GPUs can happen with high throughput

7. There are other optimization techniques as well like quantization aware training (retrain model on small dataset), sparsity

8. At inference-time an LLM request has 2 phases

9. Prefill - processes the prompt, generate 1st token and initialize the KV cache(this cache stores intermediate activations). this phase is compute-heavy.

10. Generate - generate next token using last generated token and the KV cache and update the cache. this phase is memory bound.

11. Compute in the attention block depends on the kind of implementation (can be compute or memory bound)

12. These are compute and memory bounded phases and serving solution should have custom CUDA kernels for high perf (e.g. flash attention). remember we need to keep our CUDA cores busy.

13. Static batching - traditionally, requests are accumulated over some time window, and executed as a batch until completion. this can work if we have to do fixed amount of work but in the space of llms outputs differs massively in length(e.g. "the tallest ferry wheel" vs "explain the multiverse theory")

14. In-flight batching - in llm inference theres multiple forward passes per request which is kinda unknown and unbounded(we don't know how many tokens need to be generated). so in-flight batching treats an llm iteration as a single request.

    • so imagine there are 3 requests. req-1 and req-2 is in some state of generation. req-3 is a new request.
    • we check for active requests that have met the end-of-sequence token (or some end condition) and then evict that request (e.g. req-1)
    • replace it with the new one (req-3)
    • run the next iteration

15. Sometimes in a single iteration there are requests that are in generation phase whereas others are in prefill phase. these tokens in different phases can then be concatenated for higher throughput.

16. Paged KV cache - traditionally these caches were contiguous in memory and this leads to wastage of memory since its based on max sequence length. paged KV cache however are partitioned into blocks and these do not need to be contiguous. this way memory waste only happens in the last block.

17. KV cache reuse - now this representation allows memory sharing. for example the system prompt which might be common for 2 requests can now be cached and reused for the 2nd request. see pagedAttention by vllm

18. Speculative decoding - makes educated guesses about future tokens while generating the current token, all within a single forward pass. see hitchhikers guide to speculative decoding by pytorch