Part 1 — AI Inference Fundamentals
Course: AI Inference & GPU System Optimization Level: Advanced Focus: GPU inference architecture, model loading, weight memory, inference execution, prefill, and decode
1. Introduction to AI Inference
AI inference is the process of using a trained machine learning model to produce an output from a new input.
For a large language model (LLM), inference means taking a sequence of input tokens and generating new tokens one at a time.
A simplified inference pipeline looks like this:
1User Request 2 ↓ 3Tokenizer 4 ↓ 5Input Tokens 6 ↓ 7Model 8 ↓ 9GPU Kernels 10 ↓ 11Logits 12 ↓ 13Sampling 14 ↓ 15Next Token 16 ↓ 17Repeat
For example, suppose the user asks:
1What is CUDA?
The inference system does not directly send the English sentence to the GPU.
Instead, the request passes through several stages:
1Text 2 ↓ 3Tokenizer 4 ↓ 5Tokens 6 ↓ 7Token IDs 8 ↓ 9Transformer Model 10 ↓ 11GPU Computation 12 ↓ 13Logits 14 ↓ 15Probability Distribution 16 ↓ 17Selected Token
The selected token is then added to the sequence and the process continues until the model reaches a stopping condition.
2. What Happens Inside an LLM During Inference?
A transformer model contains many computational layers.
A simplified view is:
1Input Tokens 2 ↓ 3Embedding 4 ↓ 5Transformer Layer 6 ↓ 7Transformer Layer 8 ↓ 9Transformer Layer 10 ↓ 11... 12 ↓ 13Final Layer 14 ↓ 15Logits 16 ↓ 17Sampling 18 ↓ 19Generated Token
Each transformer layer performs operations such as:
1Input 2 ↓ 3Normalization 4 ↓ 5Attention 6 ↓ 7Normalization 8 ↓ 9MLP 10 ↓ 11Output
The GPU performs most of the expensive numerical computation.
For example:
1CPU 2 │ 3 │ prepares inputs 4 ↓ 5GPU 6 │ 7 ├── Matrix multiplication 8 ├── Attention 9 ├── Normalization 10 ├── Activation functions 11 └── Other CUDA kernels 12 │ 13 ↓ 14Logits
This is why GPU performance is extremely important for modern LLM inference.
3. GPU Inference Architecture
A production inference system contains more components than simply loading a model onto a GPU.
A simplified architecture is:
1 User 2 │ 3 ↓ 4 API Server 5 │ 6 ↓ 7 Tokenizer 8 │ 9 ↓ 10 Scheduler 11 │ 12 ↓ 13 Inference Engine 14 │ 15 ┌───────┴───────┐ 16 ↓ ↓ 17 CPU Work GPU Work 18 │ 19 ↓ 20 Transformer 21 │ 22 ↓ 23 CUDA Kernels 24 │ 25 ↓ 26 Tensor Cores 27 │ 28 ↓ 29 Logits 30 │ 31 ↓ 32 Sampling 33 │ 34 ↓ 35 Generated Token
Each component has a different responsibility.
API Server
Receives requests from applications.
1Client 2 ↓ 3POST /generate 4 ↓ 5Inference Server
Tokenizer
Converts text into token IDs.
1"Hello CUDA" 2 3 ↓ 4 5[15496, 12345]
Scheduler
Controls which requests should be processed by the GPU.
For multiple users:
1Request A ─┐ 2Request B ─┼──→ Scheduler ──→ GPU 3Request C ─┤ 4Request D ─┘
Inference Engine
Executes the model.
1Scheduler 2 ↓ 3Inference Engine 4 ↓ 5Transformer 6 ↓ 7CUDA Kernels
GPU
The GPU executes highly parallel numerical operations.
1Transformer 2 ↓ 3Matrix Operations 4 ↓ 5CUDA Kernels 6 ↓ 7Tensor Cores / CUDA Cores
4. CPU vs GPU During Inference
Understanding CPU and GPU responsibilities is important for optimization.
A simplified architecture is:
1 CPU 2 │ 3 ┌─────────┼─────────┐ 4 ↓ ↓ ↓ 5 Request Tokenizer Scheduler 6 Handling 7 │ 8 ↓ 9 GPU 10 │ 11 ┌─────────┼─────────┐ 12 ↓ ↓ ↓ 13 Attention GEMM MLP 14 │ │ │ 15 └─────────┼─────────┘ 16 ↓ 17 Output
The CPU commonly handles:
- Request management
- Tokenization
- Scheduling
- Memory management
- Kernel launches
- Communication with applications
The GPU commonly handles:
- Matrix multiplication
- Attention
- Neural-network layers
- Tensor operations
- Large-scale numerical computation
The goal of inference optimization is not simply:
1"Make the GPU faster."
Instead, the goal is:
1Reduce unnecessary work 2 + 3Keep GPU efficiently utilized 4 + 5Reduce memory overhead 6 + 7Reduce synchronization 8 + 9Improve execution efficiency
5. Model Loading
Before inference can begin, the model must be loaded.
A simplified process is:
1Model Files 2 ↓ 3Load Weights 4 ↓ 5Allocate GPU Memory 6 ↓ 7Copy Weights to GPU 8 ↓ 9Initialize Model 10 ↓ 11Ready for Inference
For example:
1import torch 2from transformers import AutoModelForCausalLM, AutoTokenizer 3 4model_name = "your-model" 5 6tokenizer = AutoTokenizer.from_pretrained(model_name) 7 8model = AutoModelForCausalLM.from_pretrained( 9 model_name, 10 torch_dtype=torch.float16 11) 12 13model = model.to("cuda")
The important operation is:
1model = model.to("cuda")
This transfers model parameters from CPU memory to GPU memory.
You can inspect GPU availability:
1import torch 2 3print(torch.cuda.is_available()) 4 5if torch.cuda.is_available(): 6 print(torch.cuda.get_device_name(0))
Example output:
1True 2NVIDIA GeForce RTX 4060 Laptop GPU
6. What Does Model Loading Actually Consume?
A common beginner mistake is assuming:
1Model size = GPU memory required
This is not always true.
GPU memory can contain:
1GPU VRAM 2 │ 3 ├── Model Weights 4 ├── KV Cache 5 ├── Activations 6 ├── Temporary Buffers 7 ├── CUDA Runtime 8 └── Other Allocations
Therefore:
1Required VRAM 2≈ 3Weights 4+ 5KV Cache 6+ 7Activations 8+ 9Temporary Memory 10+ 11Runtime Overhead
This becomes especially important when serving multiple users.
7. Understanding Weight Memory
A neural-network model contains millions or billions of parameters.
Each parameter requires memory.
The approximate weight memory can be calculated as:
1Weight Memory 2= 3Number of Parameters 4× 5Bytes Per Parameter
For example:
1parameters = 7_000_000_000 2bytes_per_parameter = 2 3 4memory = parameters * bytes_per_parameter 5 6print(memory / 1024**3, "GB")
The result is approximately:
17 billion parameters 2× 32 bytes 4= 514 GB
This is a simplified theoretical calculation.
It does not include:
1KV Cache 2Activations 3Temporary Buffers 4CUDA Runtime 5Memory Fragmentation
Therefore, a model advertised as a "7B model" does not necessarily mean that a GPU with exactly 14 GB of VRAM is sufficient for practical inference.
8. Parameter Precision and Memory
The number of bytes used by each parameter depends on its data type.
A simplified comparison:
1FP32 24 bytes 3 ↓ 4FP16 / BF16 52 bytes 6 ↓ 7FP8 81 byte 9 ↓ 10INT8 111 byte 12 ↓ 13INT4 140.5 byte
For a 7B parameter model, a simplified theoretical comparison is:
1FP32 27B × 4 bytes 3≈ 28 GB 4 5FP16 67B × 2 bytes 7≈ 14 GB 8 9INT8 107B × 1 byte 11≈ 7 GB 12 13INT4 147B × 0.5 byte 15≈ 3.5 GB
These numbers represent approximate raw parameter storage, not complete inference memory requirements.
9. Inference Execution Flow
Once the model is loaded, the inference process begins.
Consider:
1"What is CUDA?"
The complete flow can be represented as:
1User Text 2 ↓ 3Tokenizer 4 ↓ 5Token IDs 6 ↓ 7Embedding 8 ↓ 9Transformer Layers 10 ↓ 11Attention 12 ↓ 13MLP 14 ↓ 15Final Projection 16 ↓ 17Logits 18 ↓ 19Sampling 20 ↓ 21Next Token
The model does not generate an entire paragraph in one simple operation.
Instead, autoregressive generation repeatedly predicts the next token.
For example:
1Input: 2 3"What is CUDA?" 4 5 ↓ 6 7Model predicts: 8 9"CUDA"
Then:
1"What is CUDA? CUDA"
The model predicts another token.
Then:
1"What is CUDA? CUDA is"
Then:
1"What is CUDA? CUDA is a"
The process continues until generation finishes.
10. What Are Logits?
The final layer of an LLM produces logits.
Conceptually:
1Transformer 2 ↓ 3Final Hidden State 4 ↓ 5Logits
Suppose the vocabulary contains thousands of possible tokens.
The model produces a score for each possible next token:
1Token Logit 2 3"GPU" 8.7 4"CPU" 5.2 5"CUDA" 9.4 6"Memory" 6.8 7"Python" 3.1 8...
These logits can then be transformed into probabilities.
Conceptually:
1Logits 2 ↓ 3Softmax 4 ↓ 5Probabilities 6 ↓ 7Sampling / Selection 8 ↓ 9Next Token
The exact sampling strategy can depend on the inference configuration.
11. Prefill
One of the most important concepts in LLM inference is prefill.
Prefill processes the input prompt before the model starts generating new tokens.
For example:
1Prompt: 2 3"Explain GPU memory optimization for LLM inference."
The tokenizer converts this into tokens:
1[Explain, GPU, memory, optimization, for, LLM, inference]
The model processes these input tokens during prefill.
Conceptually:
1Prompt 2 ↓ 3Tokenizer 4 ↓ 5Input Tokens 6 ↓ 7 PREFILL 8 ↓ 9Transformer 10 ↓ 11Attention 12 ↓ 13KV Cache
The important characteristic of prefill is that the prompt tokens can generally be processed in a highly parallel manner.
1Token 1 ─┐ 2Token 2 ─┤ 3Token 3 ─┤ 4Token 4 ─┼──→ GPU 5Token 5 ─┤ 6Token 6 ─┘
This is very different from the token-by-token generation stage.
12. Decode
After prefill, the model begins generating output tokens.
This stage is called decode.
A simplified process is:
1Prefill 2 ↓ 3Token 1 4 ↓ 5Token 2 6 ↓ 7Token 3 8 ↓ 9Token 4 10 ↓ 11...
Each new token depends on the previously generated sequence.
Therefore, decode is inherently iterative.
For example:
1Prompt 2 ↓ 3Token 1 4 ↓ 5Token 2 6 ↓ 7Token 3 8 ↓ 9Token 4 10 ↓ 11Token 5
The inference engine repeatedly executes the model to produce the next token.
13. Prefill vs Decode
The distinction is fundamental to inference performance.
1 LLM Inference 2 │ 3 ┌───────────┴───────────┐ 4 ↓ ↓ 5 Prefill Decode 6 │ │ 7 Process prompt Generate tokens 8 │ │ 9 Highly parallel Iterative 10 │ │ 11 Prompt processing Token generation
A useful mental model is:
1PREFILL 2"Read the question." 3 4 ↓ 5 6DECODE 7"Write the answer."
For example:
1Prompt: 2"Explain CUDA streams." 3 4 ↓ 5 6 PREFILL 7 8 ↓ 9 10"CUDA" 11 ↓ 12"CUDA streams" 13 ↓ 14"CUDA streams allow" 15 ↓ 16"CUDA streams allow asynchronous" 17 ↓ 18...
14. Why Prefill Can Be Parallel
Suppose the prompt contains:
1Token 1 2Token 2 3Token 3 4Token 4 5Token 5 6Token 6 7Token 7 8Token 8
During prefill, the GPU can perform substantial parallel computation across the prompt.
Conceptually:
1Token 1 ─┐ 2Token 2 ─┤ 3Token 3 ─┤ 4Token 4 ─┤ 5Token 5 ─┼──→ GPU 6Token 6 ─┤ 7Token 7 ─┤ 8Token 8 ─┘
This makes GPUs extremely effective at processing large prompts.
15. Why Decode Is Different
During decode, the model generates one new token at a time.
1Step 1 → Token 1 2Step 2 → Token 2 3Step 3 → Token 3 4Step 4 → Token 4
The next step depends on the previous generated token.
Therefore:
1Token 1 2 ↓ 3Token 2 4 ↓ 5Token 3 6 ↓ 7Token 4
This dependency makes decode fundamentally different from simply processing a large batch of independent inputs.
16. KV Cache
During transformer attention, the model produces key and value tensors.
These are commonly stored in a KV cache during autoregressive generation.
Conceptually:
1Transformer 2 ↓ 3Attention 4 ↓ 5 ┌───┴───┐ 6 ↓ ↓ 7 K V 8 └───┬───┘ 9 ↓ 10 KV Cache
For generated tokens:
1Token 1 → K/V 2Token 2 → K/V 3Token 3 → K/V 4Token 4 → K/V
The cache allows the inference engine to reuse previously computed key/value information instead of recomputing all previous attention states from scratch for every generated token.
A simplified representation is:
1KV Cache 2 3┌───────────────┐ 4│ Token 1 K/V │ 5├───────────────┤ 6│ Token 2 K/V │ 7├───────────────┤ 8│ Token 3 K/V │ 9├───────────────┤ 10│ Token 4 K/V │ 11└───────────────┘
This is one of the reasons KV-cache memory becomes a major concern in production inference systems.
17. Simplified KV Cache Memory Calculation
A simplified KV-cache memory equation is:
1kv_bytes = ( 2 2 3 * num_layers 4 * batch_size 5 * sequence_length 6 * num_kv_heads 7 * head_dim 8 * bytes_per_element 9)
The factor:
12
represents:
1K + V
The major factors are:
1Number of Layers 2 × 3Batch Size 4 × 5Sequence Length 6 × 7Number of KV Heads 8 × 9Head Dimension 10 × 11Bytes Per Element
This explains why increasing context length or batch size can dramatically increase GPU memory consumption.
18. A Complete Inference Picture
At this point, combine everything learned in Part 1:
1 User 2 │ 3 ↓ 4 Request 5 │ 6 ↓ 7 Tokenizer 8 │ 9 ↓ 10 Input Tokens 11 │ 12 ↓ 13 ┌─────────┐ 14 │ Prefill │ 15 └────┬────┘ 16 ↓ 17 Transformer 18 │ 19 ↓ 20 Attention 21 │ 22 ↓ 23 KV Cache 24 │ 25 ↓ 26 First Token 27 │ 28 ↓ 29 ┌─────────┐ 30 │ Decode │ 31 └────┬────┘ 32 ↓ 33 Next Token 34 │ 35 ↓ 36 Next Token 37 │ 38 ↓ 39 ... 40 │ 41 ↓ 42 Sampling 43 │ 44 ↓ 45 Final Text
The most important mental model is:
1Inference 2 │ 3 ├── Model Loading 4 │ 5 ├── Weight Memory 6 │ 7 ├── Prefill 8 │ 9 ├── KV Cache 10 │ 11 └── Decode
19. Practical GPU Inspection
You can inspect GPU usage using:
1nvidia-smi
Typical information includes:
1GPU 2Memory Usage 3GPU Utilization 4Temperature 5Power 6Processes
From Python:
1import torch 2 3if torch.cuda.is_available(): 4 device = torch.cuda.current_device() 5 6 print("GPU:", torch.cuda.get_device_name(device)) 7 print( 8 "Allocated:", 9 torch.cuda.memory_allocated(device) / 1024**3, 10 "GB" 11 ) 12 print( 13 "Reserved:", 14 torch.cuda.memory_reserved(device) / 1024**3, 15 "GB" 16 )
This helps distinguish between memory currently allocated by tensors and memory reserved by PyTorch's caching allocator.
20. Mini Project — Inspect an LLM's Inference Memory
Create a small experiment that loads a model and observes GPU memory.
1import torch 2from transformers import AutoTokenizer, AutoModelForCausalLM 3 4MODEL_NAME = "your-model" 5 6device = "cuda" if torch.cuda.is_available() else "cpu" 7 8tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) 9 10model = AutoModelForCausalLM.from_pretrained( 11 MODEL_NAME, 12 torch_dtype=torch.float16 13) 14 15model = model.to(device) 16model.eval() 17 18print("Model loaded on:", device) 19 20if device == "cuda": 21 print( 22 "Allocated:", 23 torch.cuda.memory_allocated() / 1024**3, 24 "GB" 25 ) 26 27 print( 28 "Reserved:", 29 torch.cuda.memory_reserved() / 1024**3, 30 "GB" 31 )
Then perform inference:
1prompt = "Explain GPU inference." 2 3inputs = tokenizer( 4 prompt, 5 return_tensors="pt" 6).to(device) 7 8with torch.inference_mode(): 9 outputs = model.generate( 10 **inputs, 11 max_new_tokens=50 12 ) 13 14result = tokenizer.decode( 15 outputs[0], 16 skip_special_tokens=True 17) 18 19print(result)
After generation, inspect memory again:
1if device == "cuda": 2 print( 3 "Allocated after inference:", 4 torch.cuda.memory_allocated() / 1024**3, 5 "GB" 6 )
This experiment demonstrates an important principle:
1Model Loaded 2 ↓ 3GPU Memory Used 4 ↓ 5Inference 6 ↓ 7Additional Runtime Memory 8 ↓ 9Generation 10 ↓ 11KV Cache / Temporary Memory
21. What You Should Understand Before Part 2
Before moving to Part 2 — Model Memory & KV Cache, you should be able to explain:
11. What is AI inference? 22. Why are GPUs used for LLM inference? 33. What happens during model loading? 44. What are model weights? 55. How does precision affect weight memory? 66. What are logits? 77. What is prefill? 88. What is decode? 99. Why is decode iterative? 1010. What is the KV cache? 1111. Why does KV cache consume GPU memory? 1212. Why is inference memory larger than model weight memory?
The key progression is:
1User Request 2 ↓ 3Tokenization 4 ↓ 5Model Loading 6 ↓ 7GPU Execution 8 ↓ 9Prefill 10 ↓ 11KV Cache 12 ↓ 13Decode 14 ↓ 15Generated Tokens
Part 1 establishes the foundation for the rest of the course:
1Part 1 2AI Inference Fundamentals 3 ↓ 4Part 2 5Model Memory & KV Cache 6 ↓ 7Part 3 8Batching & Scheduling 9 ↓ 10Part 4 11CUDA Kernel Optimization 12 ↓ 13Part 5 14Asynchronous GPU Execution 15 ↓ 16Part 6 17Inference Performance 18 ↓ 19Part 7 20Multi-GPU Inference 21 ↓ 22Part 8 23Inference Profiling 24 ↓ 25Part 9 26Complete LLM Inference Engine