The Ultimate Guide to LLM Inference Optimization- Part 2

The second part of a guide to LLM inference optimization focuses on attention mechanisms and KV cache management, explaining the compute-bound prefill phase and memory-bound decode phase. It introduces Multi-Query Attention as a technique to reduce memory overhead by sharing key and value vectors across attention heads. The article references frameworks like vLLM and SGLang that handle cache eviction automatically.

In the previous blog, we explored techniques to make LLM inference faster & efficient by working at the model level. Quantization reduces the precision of model weights, distillation transfers capabilities into smaller models, pruning removes unnecessary parts of a model, and speculative decoding reduces the number of expensive autoregressive steps. If you haven’t read Part 1 yet, I’d recommend starting there. It’ll make the ideas in this section much easier to follow. The Ultimate Guide to LLM Inference Optimization- Part 1 https://pub.towardsai.net/the-ultimate-guide-to-llm-inference-optimization-part-1-71ec501c189d All of these approaches attack the problem from different angles, but they largely leave the Transformer’s core architecture intact. At the heart of every large language model is attention . The mechanism that allows a token to decide which other tokens are related to it. You can refer to the below article to understand Transformer architecture and attention in depth. No Libraries, No Shortcuts: LLM from Scratch with PyTorch https://pub.towardsai.net/no-libraries-no-shortcuts-llm-from-scratch-with-pytorch-664c557997ee Attention mechanism is also one of the biggest sources of computational and memory cost, especially as context windows grow from thousands to millions of tokens. To understand why that is, let us understand some foundational and really important concepts. Once an LLM inference server receives a request, it doesn’t start generating text right away. For each token that came into the context window, embeddings have to be generated, key, value, and query vectors constructed, and attention scores calculated. This initial process occupies the maximum number of GPU cores, as it requires a lot of parallel computation. The prefill phase is thus called “ compute-bound ”. Performance is limited by the number of arithmetic operations the hardware is capable of executing. What we measure here is Once all required computations are done, it's time to generate the next token sequentially, one at a time. This involves loading the actual LLM weights into memory as required, for the forward pass of each token. Performance depends on the hardware’s capability to move large matrices or model layers to and fro from memory, also known as memory bandwidth. The decode phase is called “ memory-bound ”. What we measure here is In the attention mechanism, every new token needs to attend to all the tokens that came before it, but the keys and values for those previous tokens do not change. If you recompute key and value vectors for the same tokens again while decoding, it’s a straight-away inefficiency. So in the decoding step, previous key and value vectors are cached and reused instead of recalculating them. As the conversation grows, more tokens mean more key and value vectors, so the KV cache keeps growing and can consume a significant amount of GPU memory. Since context windows are in the order of millions today, cache eviction should also be implemented. Your favourite frameworks like VLLM or SGLang do these already for you. Since you have covered the foundation, let’s hop back into model optimization. Multi-Query Attention In transformer models, there usually are multiple self-attention heads within a single transformer block. Each attention head creates its own key, value, and query vectors and later aggregates attention scores, as you can see in the code below. python class MultiHeadAttention torch.nn.Module : def init self, num heads, embed dim, attention dim, dropout=0.1 : super . init self.head size = attention dim//num heads self.heads = torch.nn.ModuleList for i in range num heads : self.heads.append SelfAttention embed dim=embed dim, attention dim=self.head size,dropout=dropout def forward self,x : head outputs = for head in self.heads: head outputs.append head x B x T x A//num heads concatenated = torch.cat head outputs, dim = 2 return concatenated This becomes a real bottleneck because there are now a lot of key and value vectors that has to be managed in memory now. Fast Transformer Decoding https://arxiv.org/abs/1911.02150 1 proposed Multi-Query Attention , which shared keys and values across attention heads. A single projection computes the key and value, while each attention head has its own query projection. This reduced memory utilization, especially KV cache. But it still came with a potential accuracy drop due to the components removed. Grouped Query Attention Later, Grouped Query Attention 2 As shown in the figure, instead of a single key-value projection like MQA, there are multiple key-value projections, each of which focuses on a specific group of queries. This means num heads % num kv groups == 0. You can read and understand the code implementation of GQA in the previous articles. No Libraries No Shortcuts: Reasoning LLMs from Scratch with PyTorch — Part 1 https://pub.towardsai.net/no-libraries-no-shortcuts-reasoning-models-from-scratch-with-pytorch-part-1-bdc5bcb42042 Mixture Of Experts In vanilla transformers, every token passes through the entire model, regardless of the context. Since we are trying to model intelligence, a possible design is to structure it like the human brain. Only specific parts of the brain activate for specific cognitive tasks, not the entire brain at once. MoE is thus a collection of “ experts” within the neural network, where only the top relevant experts get activated for a given token, decided by a gating or router. This allows a trillion-parameter model to work at the same quality, while only a small fraction of the parameters actually run at any moment, giving a huge boost in efficiency and performance. This is exactly why you see models named as gemma-4–26B-A4B, which means that out of total 26B parameters, only 4B are active at a time. The idea is not new since it has been around since the 90s in the paper Adaptive Mixtures of Local Experts https://www.cs.toronto.edu/~fritz/absps/jjnh91.pdf . But the application in transformer architectures came in GShard https://arxiv.org/abs/2006.16668 2020 and Switch Transformer https://arxiv.org/abs/2101.03961 2021 , with Switch being the first trillion-parameter sparse MoE Transformer. Mixtral https://arxiv.org/abs/2401.04088 5 later demonstrated how sparse MoE transformers can be made practical at scale, delivering strong performance while maintaining efficient inference through token-level expert routing. Let’s understand MoE in detail. Just read through and understand the concept. Don’t overwhelm yourself with the code if that’s fine. Consider the following configuration: Tokens: T0, T1, T2, T3Experts: E0, E1, E2Top-k K : 2 experts per token The configuration is shown for simplicity. This data actually occurs in “batches of sequences” and not as individual tokens. The experts for each token are decided by a gate, say gate x . The output of gate x becomes, According to the configuration, only the top 2 experts are required per token. topk scores, topk expert indices = torch.topk scores, self.num experts per tok, dim=-1 topk probs = torch.softmax topk scores, dim=-1 where,topk scores : b, seq len, k topk expert indices : b, seq len, k Calculating topk scoresand applying softmaxgives topk probs. If we loop over tokens next, it would require calling each expert many times, making it slow. Instead, it is processed expert-by-expert. topk expert indices flat = topk expert indices.reshape -1, K topk probs flat = topk probs.reshape -1, K Flattening turns the “batches of sequences” into a “list of tokens” which is easier to process here. From the table, experts E0 , E1, and E2 are active. mask selected = mask selected idx slot indices = mask selected.int .argmax dim=-1 selected probs = torch.gather topk probs flat selected idx , dim=-1, index=slot indices.unsqueeze -1 For expert E0 , this gives where slot is simply the index given in topk indices in the previous table. Expert E0 produces an output for the token, which is then multiplied by its assigned gating probability. E0 T0 → y0E0 T2 → y2E0 T3 → y3output T0 += 0.88 · y0output T2 += 0.27 · y2output T3 += 0.62 · y3 The same process is repeated for each expert, and the final output for each token is a weighted sum that includes all selected experts. output T0 = 0.88·E0 T0 + 0.12·E1 T0 output T1 = 0.82·E1 T1 + 0.18·E2 T1 output T2 = 0.73·E2 T2 + 0.27·E0 T2 output T3 = 0.62·E0 T3 + 0.38·E1 T3 Paged Attention The scenario we are going to discuss now will already be familiar to you but just not in the LLM context. Imagine you have served an LLM for a chatbot on a GPU. You have now learned that every deployment needs memory reserved not only for the model but also for the KV cache. But for a chatbot, you don’t know what the maximum length of the user input could be. So you reserve a max sequence length But the problem here is that you have reserved 60–80% of your GPU memory for KV cache, but it remains unused. Your requests get rejected due to low memory errors, but then you go and check nvidia-smi to see your actual utilization as 20–40%. Sounds familiar? This is the same problem your OS solved 50 years ago by paging . Apply the same concept here. Let the KV cache be stored anywhere in your GPU memory as fixed-size blocks pages . Each block could hold KV data for, say 16 tokens. A block table maps each request’s logical sequence to physical GPU addresses, thus keeping track. The keyword is non-contiguous , which means blocks can be anywhere in physical memory. When new tokens come in, allocate a new block only when needed. Also free unused blocks once the request is finished 4 . Is there anything here besides solving fragmented memory? Absolutely. In production systems, if 50 requests use the same system prompt, the KV cache for that prefix is stored once and reused across all 50 requests. It is duplicated only when a request diverges. To understand the massive savings, see the numbers. UC Berkley termed this KV cache issue a systems problem instead of an ML problem and introduced VLLM https://arxiv.org/pdf/2309.06180 in June 2023. Paged attention is exactly what made VLLM the default inference framework for everyone. KV Cache Compression Since we already saw model weights Result: ~6× KV cache reduction, no retraining, and near full-precision behavior. You can see the code implementation of this in llama.cpp https://github.com/ggml-org/llama.cpp/issues/20977 . It is also fully functional in VLLM. https://github.com/vllm-project/vllm/pull/38479 Another approach called RotorQuant https://www.scrya.com/rotorquant.pdf further enhanced this. It uses block-diagonal Clifford algebra rotors instead of heavy dense matrix multiplications in TurboQuant, achieving a 5x to 10x memory reduction with massive speed boosts. From Mac mini to gigawatt-scale data centers, hardware is the hottest commodity in the AI boom. Although Nvidia is the centre of attention, there are a lot of companies providing accelerators for AI, including AMD, Intel, Cerebras, Groq, Triton OpenAI , TPU Google , Trianium & Inferentia AWS , and much more. All of them have different architectures. Even within companies, for example Nvidia, RTX 5090 is very different from an H200 in architecture and capacity. Running models on these chips differs too. The same model does not guarantee the same throughput and quality on all silicon. Thus, it is important to optimize the specific model for the target hardware whenever possible. Anacceleratoris a specialized piece of hardware designed to perform a particular type of computation much faster and/or more efficiently than a general-purpose CPU. For example, a GPU is an AI accelerator because it can perform thousands of mathematical operations in parallel, making it well suited for operations such as the matrix multiplications used by neural networks. This is where the big brains in AI are in demand. They write low-level code called kernels that perform operations such as matrix multiplication and are carefully optimized for the target hardware. Different hardware uses different languages or frameworks for this, which have distinct levels of low-level thread and memory access. You may have heard of CUDA, NVIDIA’s programming platform for running and optimizing workloads on NVIDIA GPUs. As different models and GPUs are released, they could have different architectures, and custom kernels are released for them. You already know that a neural network is actually a bunch of mathematical operations. At a high level, kernels are written to speed up these operations by combining or dividing them and also executing them in parallel wherever possible. Writing efficient kernels can be a topic for a future blog. Here is some high-level intuition for common approaches used in writing kernels. Compilers take higher-level code or computational graphs and transform them into lower-level operations that can efficiently execute on a target hardware platform. For a large language model, compilers convert all the operations involved in it into efficient kernels for the particular hardware. If you have worked with PyTorch, you may have come across torch.compile. In eager mode, which is the normal usual one, PyTorch executes operations as they are encountered, which makes development and debugging straightforward but can leave opportunities for optimization on the table. torch.compile instead captures the model's operations and uses a compiler to optimize the resulting computation, for example by fusing operations and generating more efficient kernels for the target hardware. In practice, you can start with eager mode for simplicity and switch to torch.compile when you want the compiler to optimize the model for faster execution. Coming to GPUs and chips, NVIDIA offers a suite of tools through TensorRT that optimize neural networks for NVIDIA GPUs. TensorRT-LLM https://developer.nvidia.com/blog/optimizing-inference-on-llms-with-tensorrt-llm-now-publicly-available/ is the library for large language models, which creates an executable graph of operations with the model weights embedded. A usual production method is having an LLM in ONNX format and then running custom scripts to convert it into .engine formats, which is then deployed on GPUs using Triton Inference Server https://www.nvidia.com/en-us/ai-data-science/products/triton-inference-server/ . You can check out some ready-made scripts for most LLMs in the GitHub repo https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core . Similarly, Qualcomm has MLC LLM https://github.com/mlc-ai/mlc-llm Machine Learning Compilation for Large Language Models , which compiles and optimizes LLMs for Snapdragon chipsets GPUs & NPUs . Apple has provided extensive support for iPhone’s A series and Mac’s M series chipsets through the Apple MLX framework https://machinelearning.apple.com/research/exploring-llms-mlx-m5 . You can compile models easily using existing or custom scripts for various open-source LLMs on HuggingFace. pip install mlx-lmmlx lm.convert \ --hf-path mistralai/Mistral-7B-Instruct-v0.3 \ -q \ --upload-repo mlx-community/Mistral-7B-Instruct-v0.3-4bit This method is still at an early stage. Released a few days back by Google Deepmind, recirculation https://arxiv.org/html/2608.17981 involves only inference time implementations rather than retraining the model. During generation, activations from deeper layers are fed back into the shallower layers of the model which in fact becomes a second pass over its own internal representations. No extra tokens are involved and this is purely meant as a plug in implementation over existing inference pipelines. You have now learned to run your LLMs efficiently on your hardware using model-level adaptations. But you still can’t say it will respond the same when people start using it. This is where we should discuss what can be done at the software level. Let’s meet again at Part 3 . 1 Noam Shazeer. 2019 . Fast Transformer Decoding: One Write-Head is All You Need https://arxiv.org/abs/1911.02150 2 Joshua Ainslie, James Lee-Thorp, et al. 2023 . GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints https://arxiv.org/abs/2305.13245 3 Amir Zandieh, Majid Daliri, Majid Hadian, et al. 2025 . TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate https://arxiv.org/abs/2504.19874 4 Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, et al. 2023 . Efficient Memory Management for Large Language Model Serving with PagedAttention https://arxiv.org/abs/2309.06180 5 Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, et al 2024 . Mixtral of Experts https://arxiv.org/abs/2401.04088 . arXiv:2401.04088 6 Sebastian Raschka https://medium.com/u/3d00fdc9f1e7?source=post page---user mention--bdc5bcb42042--------------------------------------- .LLMs From Scratch 7 Danny Sawyer, Shoaib Ahmed Siddiqui, et al. 2026 . Recirculation https://arxiv.org/html/2608.17981 If not otherwise stated, all images are created by the author. The Ultimate Guide to LLM Inference Optimization- Part 2 https://pub.towardsai.net/the-ultimate-guide-to-llm-inference-optimization-part-2-dcfcc960120a was originally published in Towards AI https://pub.towardsai.net on Medium, where people are continuing the conversation by highlighting and responding to this story.