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
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
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 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.
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 [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
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. But the application in transformer architectures came in GShard (2020) and Switch Transformer (2021), with Switch being the first trillion-parameter sparse MoE Transformer. Mixtral [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 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. It is also fully functional in VLLM. Another approach called RotorQuant 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 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. You can check out some ready-made scripts for most LLMs in the GitHub repo.
Similarly, Qualcomm has 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. 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 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
[2] Joshua Ainslie, James Lee-Thorp, et al.(2023).GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
[3] Amir Zandieh, Majid Daliri, Majid Hadian, et al.(2025).TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate
[4] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, et al.(2023).Efficient Memory Management for Large Language Model Serving with PagedAttention
[5] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, et al(2024).Mixtral of Experts. arXiv:2401.04088
[6] Sebastian Raschka.LLMs From Scratch
[7] Danny Sawyer, Shoaib Ahmed Siddiqui, et al.(2026).Recirculation
If not otherwise stated, all images are created by the author.
The Ultimate Guide to LLM Inference Optimization- Part 2 was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.