PyTorch-AI-Compiler.pdf

Transcript

1. From PyTorch to Hardware An Introduction to LLM Compilers Tommy Han, 2026
2. About me • Not working with ML/LLM Compiler on daily basis. • One of the Compiler Meetup Host in Hong Kong • Software Engineer
3. Before we begin…
4. Before we begin… A few things to mention This is not a “how to tune/train your model better” talk -> Leave to Mathematicians! Expect you to know some ML Basics (Matrix Multiplication), Tensor Basics Expect you to know how a pipeline-based compiler works. Feel free to ask questions, and point out the mistakes I made!
5. Agenda What are we going to talk about… Level 0: Background Build Up! Level 1: What actually is ML/LLM Compiler Level 2: Why ML/LLM Compiler is important
6. Level 0: Background Build Up!
7. Level 0 How a LLM is born?
8. Level 0 Mostly starts with PyTorch… • Start with PyTorch(or sth else), build up datasets and training • Export the training data to graph (Done by ATenGraph/StableHLO) • Graph will getting to the compiler for optimization -> Output • Inferencing • Runtime (e.g. vLLM, ONNX, llama.cpp, etc)
9. Level 0 When the Model is Compiled and ready… • Runtime running on Hardwares • LLM Applications (The stu we ff fl mostly using.) e.g. generation, RAG, Inferencing API, chatbots, agents, AI work ows.
10. Level 0 What are the hardwares? • GPU (Most Common) • SoC Chips (e.g. Apple M-Chips) • TPUs • Custom Made NPUs
11. Level 0 Why there are new types of hardware? • Accelerate Graph Computation • Build Speci c Instruction Sets to handle the maths calculations easier. fi (will talk about that later)
12. Level 0 How is the current open-source ecosystem?
13. Level 0 Where is the compiler? Between runtime and Hardware. It is similar to normal compilers. ff But the details are di erent.
14. Level 1: ML/LLM Compiler
15. Level 1 Normal Programming Language Compiler
16. But you can’t scan a neural-network layer right?
17. Level 1 Then, what Actually is a ML/LLM Compiler? ffi An ML/LLM Compiler transforms high-level tensor computations into e cient programs for the target hardware.
18. Level 1 What is a tensor operation? A tensor operation describes what to compute. fi The compiler gures out how to compute it.
19. Level 1 ffi But the loops are pretty ine cient, right?
20. Level 1 Lets zoom back out to the general architecture…
21. Demo
22. Level 2: Why they are important?
23. Level 2 The Math is Simple, execution is expensive Model says: Y = GELU(X @ W + b) But execution involves: • Huge tensor operations • Large memory movement • Many GPU kernels ff • Di erent hardware architectures
24. Level 2 Compute is not only the bottleneck
25. Level 2 Hardware is highly specialized CPU → vectorization → cache locality GPU → thousands of threads → shared memory → Tensor Cores NPU → specialized matrix units fi → speci c data layouts
26. Level 2 LLMs make the problem much bigger
27. Level 2 ML Compiler = Automatic Performance Engineering
28. Level 2 Why not optimize everything?
29. Last words And the beginning… • Knowing what’s under the hood is great! ff • Try out di erent frameworks for helping you building LLMs!
30. How to get started It seems pretty challenging… https://www.cmu.edu/ ame/research/2024/machine-learning-compilation-forlarge-language-models.html fl https://aramesh10.github.io/ml-compilers-primer/index.html
31. How to get started It seems pretty challenging… https://github.com/apache/tvm/ https://github.com/triton-lang/triton https://github.com/openxla/xla
32. Thank You!
33. Appendix Mostly include references