• Company B, because its cost is cheaper

• Probably not, because flash attention is essential for large scale model deployment

1 - Hardware: inference on GPUs

1.1 - Preliminary

1.1.1 - GPU architecture

• Basics: DRAM, L2 cache, and SM

• Comparison to CPU
• SM is simialr to CPU cores, but with significantly larger level of parallelism
• L2 cache and DRAM is simialr to CPU L2 and DRAM
• In FlashAttention paper, L2 L1 is called SRAM

• A100 80G SXM
• 108 SM, DRAM 80G, 40M L2 cache

• L1 cache: instruction and data

• Tensor core: where matrix multiplication happens. Recall that neural network computation is basically gigantic batchs of matrix multiplication

1.1.2 - GPU programming basics

• Memory access:
• load model weights from HBM to L2 cache to SM

• Compute:
• perform a matrix multiplication in SM, SM asks tensor core to do it
• Which operation takes more time? Actually it is the memory access.
• We will show that when running transformer inference, most of the time is spent on waiting to move model parameters / activation from / to the memory, rather than the actual computation

• A100:
• 108 SM, DRAM 80G, 40M L2 cache
• bf16 tensor core: 312 trillion float point operation per second (tflops)
• DRAM memory bandwidth 2039 GB/sec = 2.039T /sec

• If the model is large, we split it into multiple GPUs, say two, where the two GPUs are connected by NVLink
• NVLink 300GB/sec = 0.3T / sec

• We roughly observe the speed hierarchy. Although they are not directly comparible, their orders of magnitude difference is the major places we need to optimize:
• 312T (SM compute) > 2.03T (DRAM memory access) > 0.3T = 300G (NVLink cross-device communication) > 60G (PCIe cross-device communication)

• This means, if we we want things to be fast, we should try our best to
• fully utilize the SMs,
• reducing memory access in one GPU (because its much slower than compute),
• and reducing communication between GPUs (because its even slower than memory access)

1.1.3 - Compute v.s. memory bound

• define GPU operations per byte = flop / memory bandwidth
• A100 = 312 / 2.039

• define Arithmetic intensity = compute / memory access

• if arithmetic intensity ~ ops per byte, then compute bound, if smaller then memory bound. below is the arithmetic intensity of typical neural network layers, as noted by Nvidia blog:



• Increasing batch size will change the behavior from compute memory bound to memory compute bounded

• kernel fushion: reduce memory access ops, because we fuse multiple operations into one

GPU Performance Background User's Guide

GPUs accelerate machine learning operations by performing calculations in parallel. Many operations, especially those representable as matrix multipliers will see good acceleration right out of the box. Even better performance can be achieved by tweaking operation parameters to efficiently use GPU resources. The performance documents present the tips that we think are most widely useful.

https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html

1.2 - Transformer inference basics

1.2.1 - Prefilling and decoding

• Prefilling:
• compute kv cache for prompt.
• Compute bound, because we compute a sequence of tokens in parallel

• decoding:
• sample next tokens autoregressively.
• Memory bound, because we only compute one token, not fully utilizing SMs

1.2.2 - Transformer inference is memory bound

• because decoding only sample one token at one pass

• Increase batch size improves hardware efficiency.
• because we compute multiple tokens at one time
• large batch changes from memory bound to flop bound.

• However, we may not do so large batch, because GPU memory is not large, currently maximum 80G.

1.2.3 - Memory layout

• we cannot have too large batch (though we want the batch size to be large to improve efficiency),

• or too long sequence (though we want to serve 100k length for sure)

1.2.4 - Online and offline inference, throughput v.s. latency

• Offline: throughput optimization
• we care about this scenario because we may want to offline evaluate the model, say, run an intermediate pretrained checkpoint on 100 benchmarks to verify our pretraining is healthy
• Increase batch size can help, but remember currently 80G memory is maximum for single device.

• Online: latency and throughput tradeoff
• When batch size large (assume still fit in memory), we become compute bound, then latency increases
• The latency should not be slower than human read speed, or the user will complain, or switch to your competitor’s model
• But again, we do want the batch size to be large, to improve efficiency

1.2.5 - Context length scaling

• Prefilling,
• this time, prefilling takes sigfinicantly longer time, because input lengths is so long
• in this case, latency to first token generation is important, because the user does not want to wait for 10s to see the model speak

• large KV cache
• simply because the context length is long

• currently it seems not so much of work improving inference in this space

2 - MLSys: Flash attention and vLLM

2.1 - vLLM and Paged attention

• Basically constructing a memory management system similar to CPU memory management, to reduce fraction and fully utilize memory throughput

• Now the goto places for transformer inference

2.2 - Flash attention

• Instead of storing the full attention matrix in the HBM, do blockwise computation of the dot product, such that all the computation is performed in the L2 cache

• significantly reduced memory usage, such that you can put in 100k context length using brutal force — yes, there is no fancy algorithm for 100k, just brutal force
• in the original paper the authors only test up to 16k, but 100k is totally doable

• significantly improved throughput, particularly for small models where a large portion of the flop is in the dot-product operation

2.3 - Flash Decoding

3 - Modeling: architecture and decoding algorithm

3.1 - Classical methods

3.1.1 - Distillation and sparse attention

• Distillation: to finetune a small model using larger models’ outputs / logits
• Finetune on outputs: nowadays everybody distill from GPT and you guys are very good at it, so I just skip this part
• Finetune on logits / distribution is a field that is less explored, some results show faster convergence speed and better quality

• Sparse and local attention, particularly important for long-context
• In the age of small models, this part is very well studied (c.f. Yi Tay’s awesome survey), but not sure whether these results hold for larger scale
• In the age of large models, there is little work in this space except Mistral’s sliding window attention

3.1.2 - Quantization

• Most foundamental method, quantize the model weight to int 8.

• Quantization is nowadays a must to deploy a large model. The good news is that it does not really harm performance

• Yi-34B chat 4 bit quantization. Only requires 17G memory. Performance on benchmarks almost the same.

• Actually it’s quite fast, you can try it on huggingface

3.1.3 - Multi-query attention and group-query attention

• In genereal, multi-query attention significantly speed up training and inference by simultanously reducing memory and compute.

• Multi-query attention is also a great example where differences in small models do not exist in large models: for small models like 7B, multi-query attettion is worse than full attention, but when models become as large as 70B, multi-query attention has basically the same performance as full attention. LLaMA2 7B uses full attention, and 70B uses GQA.

• Current state of art large models by default use multi-query attention

3.2 - Advanced techniques

3.2.1 - Mixture of experts

• Say we have 7B activation, 34B params in total.

• Can we achieve?
• performance similar to 34B
• throughput better than 34B
• latency similar to 7B

• The above goal is made possible with the recently released Mistral MoE, see discussion on X

• Performance: the 50B MoE Mistral with 7B dense part is somewhere near a 34B model Yi and 67B DeepSeek

• Inference efficiency: the dense part of MoE is 7B, with top 2 activation

• When concurrency is low, most of the time is spent on loading the two activated expert to the memory, which is smaller than a dense layer, thus requiring less memory access time, thus lower latency
• in other words, for a single query, MoE has lower latency than dense because we need to read less parameters from the memory

• When concurrency is large, we enter the flop bound regime, but since MoE has less activation than dense, the throughput is higher
• in other words, for many concurrent queries, MoE has higher throughput because the

• MoEfication: decompose a dense model to an MoE model, making it
• as efficient as a small model,
• but as strong as a large dense model.

• Ref:
• Zhang e.t al. 2021. MoEfication: Transformer Feed-forward Layers are Mixtures of Experts
• Zhang et. al. 2023. Emergent Modularity in Pre-trained Transformers

3.2.2 - Early exit

• for all tokens, use a gate to determine whether early exit or not

• Schuster et. al. 2022. Confident Adaptive Language Modeling

• Bae et. al. 2023. Fast and Robust Early-Exiting Framework for Autoregressive Language Models with Synchronized Parallel Decoding

3.2.3 - Blockwise decoding

• Leviathan et. al. 2022. Fast Inference from Transformers via Speculative Decoding

• Chen et. al. 2023. Accelerating Large Language Model Decoding with Speculative Sampling

• Liu et. al. 2023. Online Speculative Decoding

• it is weak, so rejection rate may be high in some challening domains

• we should consider the accuracy and overhead tradeoff: a smaller draft model is fast but inaccurate, a larger draft model is accurate and slow

• we need to put in two models in your GPU, but remember we have already running out of GPU memories

• It is strong, so rejection rate can be reduced

• we only need to keep one model in the memory

• use multiple heads to decode multiple tokens at a time

• use the large model itself as the draft model

4 - What we have not yet covered

• Deep water Hardware and MLSys
• I am new to MLSys. I have to consistently ask my friends whether my take on MLSys is correct or not. So inevidently my take in this space, although I try to cover the most important foundamentals, is merely scratching the surface
• There are many more awesome articles covering hardware the MLSys which I will add to the reference at the end of this article

• Distributed Inference of super large models. c.f. Pope et. al. Efficiently Scaling Transformer Inference
• key tech: model sharding, pipeline and tensor parallelism
• Strongly suggest reading the Nvidia blog on this part

• Continuous batching and libraries like DeepSpeed MII

• Linear architectures like Mamba

• Distillation and quantization for mixture of experts
• It seems that MoE’s experts in general have a higher level of intrinsic sparsity than dense FFNs, thus may take more radical quantization
• It is not so clear about the mistillation from / between MoE and dense, like how to distill from a gigantic MoE to a small dense

• More advanced blockwise decoding
• Lookahead decoding
• Retrieval-Based Speculative Decoding
• EAGLE

5 - Personal pick

Conclusion

An ice latte at Manner, Shanghai

A cappuccino at Blue Bottle, New York

References

Mastering LLM Techniques: Inference Optimization | NVIDIA Technical Blog

Stacking transformer layers to create large models results in better accuracies, few-shot learning capabilities, and even near-human emergent abilities on a wide range of language tasks.

https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization

• Nvidia official blog on transformer inference. This blog covers model parallelism that we do not discuss

Accelerating Generative AI with PyTorch II: GPT, Fast

This post is the second part of a multi-series blog focused on how to accelerate generative AI models with pure, native PyTorch. We are excited to share a breadth of newly released PyTorch performance features alongside practical examples to see how far we can push PyTorch native performance. In part one, we showed how to accelerate Segment Anything over 8x using only pure, native PyTorch. In this blog we’ll focus on LLM optimization.

https://pytorch.org/blog/accelerating-generative-ai-2/

• Pytorch official support for transformer inference

Full Stack Optimization of Transformer Inference: a Survey

Recent advances in state-of-the-art DNN architecture design have been moving toward Transformer models. These models achieve superior accuracy across a wide range of applications. This trend has...

https://arxiv.org/abs/2302.14017

• GPT-2 level inference, not so large though

Efficiently Scaling Transformer Inference

We study the problem of efficient generative inference for Transformer models, in one of its most challenging settings: large deep models, with tight latency targets and long sequence lengths....

https://arxiv.org/abs/2211.05102

• Deployment of PaLM 540B

• Deployment of Claude 52B

GPU Performance Background User's Guide

GPUs accelerate machine learning operations by performing calculations in parallel. Many operations, especially those representable as matrix multipliers will see good acceleration right out of the box. Even better performance can be achieved by tweaking operation parameters to efficiently use GPU resources. The performance documents present the tips that we think are most widely useful.

https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html

• Awesome GPU guide

A guide to LLM inference and performance

To attain the full power of a GPU during LLM inference, you have to know if the inference is compute bound or memory bound. Learn how to better utilize GPU resources.

https://www.baseten.co/blog/llm-transformer-inference-guide/

• Single device inference. Compute and memory bound. Arithmetic indensity.

Why GPT-3.5 is (mostly) cheaper than Llama 2

Llama-2 is more expensive than you'd think. In this post, we explore why it's often more expensive than gpt-3.5-turbo.

https://cursor.sh/blog/llama-inference

• Detailed analysis on the math of transformer inference

Large Transformer Model Inference Optimization

[Updated on 2023-01-24: add a small section on Distillation.] Large transformer models are mainstream nowadays, creating SoTA results for a variety of tasks. They are powerful but very expensive to train and use. The extremely high inference cost, in both time and memory, is a big bottleneck for adopting a powerful transformer for solving real-world tasks at scale. Why is it hard to run inference for large transformer models? Besides the increasing size of SoTA models, there are two main factors contributing to the inference challenge (Pope et al.

https://lilianweng.github.io/posts/2023-01-10-inference-optimization

• Inference optimization algorithms

Dissecting the Runtime Performance of the Training, Fine-tuning,...

Large Language Models (LLMs) have seen great advance in both academia and industry, and their popularity results in numerous open-source frameworks and techniques in accelerating LLM pre-training,...

https://arxiv.org/abs/2311.03687

• Runtime of 7B / 13B model on 8 * 80G A100 with different paralelism strategies