AI Inference & GPU System Optimization
AI inference is the process of taking a trained model and using it to generate predictions or tokens.
For a modern LLM, inference is not simply:
1Model 2 ↓ 3GPU 4 ↓ 5Answer
A production inference system must coordinate:
1 LLM Inference System 2 │ 3 ┌──────────────────┼──────────────────┐ 4 │ │ │ 5 Compute Memory Scheduling 6 │ │ │ 7 CUDA Kernels Weights Batching 8 Tensor Cores KV Cache Queues 9 Fusion Activations Streams 10 │ │ │ 11 └──────────────────┼──────────────────┘ 12 ↓ 13 Generated Tokens
This course connects CUDA programming, GPU memory, kernels, and distributed systems to practical AI inference.
1. GPU Inference Architecture
A simplified LLM inference pipeline looks like:
1User Prompt 2 ↓ 3Tokenizer 4 ↓ 5Input Tokens 6 ↓ 7Embedding 8 ↓ 9Transformer Layers 10 ↓ 11CUDA Kernels 12 ↓ 13Logits 14 ↓ 15Sampling 16 ↓ 17Next Token 18 ↓ 19Repeat
For example:
1"Hello AI" 2 ↓ 3[15496, 9552] 4 ↓ 5Embedding 6 ↓ 7Transformer 8 ↓ 9logits 10 ↓ 11token = 220 12 ↓ 13"!"
The inference engine repeatedly executes the transformer until a stopping condition is reached.
2. Loading Model Weights
A neural network contains a large number of parameters.
A simplified PyTorch model:
1import torch 2import torch.nn as nn 3 4 5class TinyModel(nn.Module): 6 def __init__(self): 7 super().__init__() 8 9 self.embedding = nn.Embedding( 10 10000, 11 768 12 ) 13 14 self.projection = nn.Linear( 15 768, 16 10000 17 ) 18 19 20model = TinyModel() 21 22model = model.to("cuda") 23model.eval()
The important operation is:
1model.to("cuda")
which moves model parameters to GPU memory.
You can inspect memory:
1print( 2 torch.cuda.memory_allocated() / 1024**2, 3 "MB" 4)
A production inference engine must carefully control this memory because GPU VRAM is finite.
3. Weight Memory
Suppose a model has:
17 billion parameters
Using FP16:
12 bytes / parameter
Approximate weight memory:
17,000,000,000 × 2 2= 314 GB
Using INT4:
10.5 bytes / parameter
the theoretical raw weight storage becomes approximately:
17,000,000,000 × 0.5 2= 33.5 GB
Real implementations require additional metadata and scaling information, so actual memory usage is higher.
This demonstrates why precision and quantization are critical for inference.
4. FP16 and BF16
FP16 uses 16 bits:
116 bits 2 ↓ 32 bytes
BF16 also uses 16 bits but has a different floating-point representation.
Conceptually:
1FP32 2 │ 3 ├── Higher precision 4 └── Larger memory 5 6FP16 / BF16 7 │ 8 ├── Lower storage 9 ├── High GPU throughput 10 └── Common AI inference formats
Example:
1model = model.half()
or:
1model = model.to( 2 dtype=torch.bfloat16 3)
The correct datatype depends on the GPU architecture and workload.
5. FP8
FP8 uses 8 bits per value.
1FP32 → 32 bits 2FP16 → 16 bits 3FP8 → 8 bits
Reducing precision can significantly reduce memory traffic and storage.
However, lower precision requires careful handling of:
- Scaling
- Numerical range
- Accuracy
- Accumulation
- Hardware support
Modern NVIDIA GPUs provide hardware designed to accelerate low-precision AI workloads.
6. Quantization
Quantization converts model values into lower-precision representations.
A simplified example:
1import torch 2 3x = torch.tensor( 4 [-2.0, -1.0, 0.0, 1.0, 2.0] 5) 6 7scale = x.abs().max() / 127 8 9q = torch.round(x / scale) 10 11print(q)
Conceptually:
1FP32 weights 2 ↓ 3Quantization 4 ↓ 5INT8 / INT4 6 ↓ 7Smaller memory footprint 8 ↓ 9Potentially faster inference
Quantization is not simply "convert everything to integers."
Real inference systems need:
1Weights 2 + 3Scales 4 + 5Zero points / quantization parameters 6 + 7Efficient kernels
7. INT8
INT8 uses 8 bits per value.
Compared with FP16:
1FP16 → 16 bits 2INT8 → 8 bits
This can reduce weight storage substantially.
A simplified quantization operation:
1def quantize_int8(x): 2 scale = x.abs().max() / 127.0 3 4 q = torch.clamp( 5 torch.round(x / scale), 6 -128, 7 127 8 ) 9 10 return q.to(torch.int8), scale
Dequantization:
1def dequantize_int8(q, scale): 2 return q.float() * scale
Real production systems use optimized quantization schemes and GPU kernels rather than performing these operations with ordinary Python code.
8. INT4
INT4 uses only four bits.
1FP16 → 16 bits 2INT8 → 8 bits 3INT4 → 4 bits
The theoretical storage reduction is significant.
But there is an important distinction:
Smaller weights do not automatically mean faster inference.
The GPU also needs efficient kernels capable of processing the compressed representation.
Therefore:
1Quantization 2 + 3Efficient kernel 4 + 5Hardware support 6 = 7Useful inference optimization
9. KV Cache
The KV cache is one of the most important concepts in LLM inference.
During transformer attention, the model generates:
1Query (Q) 2Key (K) 3Value (V)
For autoregressive generation, previously calculated keys and values can be reused.
Without caching:
1Token 1 → calculate 2Token 2 → recalculate Token 1 + Token 2 3Token 3 → recalculate Token 1 + Token 2 + Token 3 4...
With KV cache:
1Token 1 → K/V stored 2Token 2 → reuse Token 1 K/V 3Token 3 → reuse Token 1 + Token 2 K/V 4...
Conceptually:
1 KV Cache 2 │ 3 ┌───────────┼───────────┐ 4 ↓ ↓ ↓ 5 Layer 1 Layer 2 Layer N 6 │ │ │ 7 K/V K/V K/V
This dramatically changes the memory requirements of long-context inference.
10. A Simplified KV Cache
A conceptual PyTorch implementation:
1class KVCache: 2 def __init__( 3 self, 4 num_layers, 5 batch_size, 6 num_heads, 7 max_seq_len, 8 head_dim, 9 device="cuda" 10 ): 11 self.key = torch.zeros( 12 num_layers, 13 batch_size, 14 num_heads, 15 max_seq_len, 16 head_dim, 17 device=device, 18 dtype=torch.float16 19 ) 20 21 self.value = torch.zeros_like( 22 self.key 23 )
The real layout used by inference engines can be significantly more sophisticated.
The key idea is:
1Previously generated tokens 2 ↓ 3 K/V storage 4 ↓ 5Reuse during decode
11. Prefill
Prefill processes the initial prompt.
Suppose:
1Prompt: 2 3"Explain GPU memory optimization"
The tokenizer produces:
1[token1, token2, token3, token4]
The model processes these tokens together.
1Prompt 2 ↓ 3Tokenizer 4 ↓ 5Many input tokens 6 ↓ 7Transformer 8 ↓ 9KV Cache populated
Prefill is generally highly parallel.
Therefore it tends to utilize GPU compute efficiently.
12. Decode
After prefill, the model generates tokens one at a time.
1Prompt 2 ↓ 3Prefill 4 ↓ 5Token 1 6 ↓ 7Token 2 8 ↓ 9Token 3 10 ↓ 11Token 4 12 ↓ 13...
Each decode iteration generally processes the newly generated token while reusing cached K/V values.
This creates a major performance difference:
1Prefill 2→ parallel processing 3 4Decode 5→ iterative token generation
13. Prefill vs Decode
| Stage | Main Characteristic |
|---|---|
| Prefill | Processes prompt tokens |
| Decode | Generates tokens iteratively |
| Prefill | More parallel |
| Decode | More sequential |
| Prefill | Often compute-heavy |
| Decode | Often memory-bandwidth sensitive |
This distinction is fundamental to LLM serving.
14. Simple Inference Loop
A simplified inference loop might look like:
1tokens = tokenizer(prompt) 2 3past_key_values = None 4 5for _ in range(max_new_tokens): 6 7 outputs = model( 8 tokens, 9 past_key_values=past_key_values, 10 use_cache=True 11 ) 12 13 logits = outputs.logits[:, -1, :] 14 15 past_key_values = ( 16 outputs.past_key_values 17 ) 18 19 next_token = torch.argmax( 20 logits, 21 dim=-1 22 ) 23 24 tokens = next_token.unsqueeze(1)
This is educational rather than production-grade.
A production engine must handle batching, scheduling, memory management, sampling, synchronization, and many other concerns.
15. Batch Processing
Suppose three users send requests:
1Request A 2Request B 3Request C
Instead of processing them independently:
1GPU 2 ├── Request A 3 ├── Request B 4 └── Request C
we can form a batch:
1 Batch 2 ┌─────┼─────┐ 3 ↓ ↓ ↓ 4 A B C 5 └─────┼─────┘ 6 ↓ 7 GPU
PyTorch example:
1input_ids = torch.tensor( 2 [ 3 [10, 20, 30], 4 [11, 21, 31], 5 [12, 22, 32] 6 ], 7 device="cuda" 8) 9 10with torch.inference_mode(): 11 output = model( 12 input_ids 13 )
Batching improves GPU utilization when workloads are suitable for it.
16. Continuous Batching
Static batching waits for a group of requests.
Continuous batching dynamically adds and removes requests as generation progresses.
Conceptually:
1Time ─────────────────────────────► 2 3Batch: 4A A A A 5 B B B B B 6 C C C 7 D D D D
Instead of waiting for every request to finish:
1Request A finishes 2 ↓ 3Slot becomes available 4 ↓ 5Request D enters
This is extremely useful for LLM serving systems.
17. Memory Fragmentation
Inference workloads allocate memory dynamically.
Suppose GPU memory looks like:
1┌────────┬──────┬──────────┬─────┬────────┐ 2│ Used │Free │ Used │Free │ Used │ 3└────────┴──────┴──────────┴─────┴────────┘
There may be enough total free memory but not enough contiguous usable space for a particular allocation.
This is memory fragmentation.
A production inference engine therefore needs careful memory management.
18. Paged KV Cache Concept
A useful strategy is to divide KV cache memory into fixed-size blocks.
1KV Memory Pool 2 3┌─────┬─────┬─────┬─────┬─────┬─────┐ 4│ B0 │ B1 │ B2 │ B3 │ B4 │ B5 │ 5└─────┴─────┴─────┴─────┴─────┴─────┘
A sequence can reference blocks:
1Sequence A 2 ↓ 3B0 → B3 → B5 4 5Sequence B 6 ↓ 7B1 → B2 8 9Sequence C 10 ↓ 11B4
This avoids requiring one large contiguous allocation for every sequence.
It also enables more flexible memory sharing and scheduling.
19. Kernel Fusion
Suppose a workload performs:
1Kernel A 2 ↓ 3Kernel B 4 ↓ 5Kernel C
Each kernel may require:
1Launch 2 ↓ 3Read 4 ↓ 5Compute 6 ↓ 7Write
Separate kernels can cause additional memory traffic and launch overhead.
Fusion combines operations:
1Kernel A + Kernel B + Kernel C 2 ↓ 3 Fused Kernel
For example, instead of:
1x = x * scale 2x = x + bias 3x = torch.relu(x)
a fused implementation can conceptually perform:
1x = ReLU(x * scale + bias)
inside one optimized kernel.
The benefit depends on the workload and implementation.
20. CUDA Graphs for Inference
LLM inference often repeats similar execution patterns.
A simplified sequence:
1Embedding 2 ↓ 3Attention 4 ↓ 5MLP 6 ↓ 7Normalization 8 ↓ 9Sampling
If the execution structure is stable, CUDA Graphs can reduce repeated CPU-side kernel-launch overhead.
Conceptually:
1CPU 2 │ 3 └── Launch Graph 4 │ 5 ▼ 6GPU: 7Embedding → Attention → MLP → Norm
This is particularly useful for repeated workloads where graph capture requirements can be satisfied.
21. Multi-Stream Execution
Multiple CUDA streams can allow independent work to overlap.
1Stream 0: 2Copy → Kernel A → Copy 3 4Stream 1: 5Copy → Kernel B → Copy
Conceptually:
1Time ──────────────────────► 2 3Stream 0 █████████████ 4Stream 1 █████████████
However, streams do not magically make a workload faster.
The GPU must have enough independent work and available resources.
22. Tensor Cores
Tensor Cores are specialized hardware for matrix operations.
Transformer workloads contain many matrix multiplications:
1Q = XWq 2K = XWk 3V = XWv
and:
1Output = Attention × W
Tensor Cores can accelerate supported low-precision matrix operations.
Conceptually:
1Transformer 2 ↓ 3GEMM 4 ↓ 5Tensor Core 6 ↓ 7High-throughput matrix computation
This is one reason choosing the correct datatype matters.
23. Inference Memory Model
A simplified inference memory budget is:
1GPU VRAM 2 │ 3 ├── Model Weights 4 ├── KV Cache 5 ├── Activations 6 ├── Temporary Buffers 7 ├── CUDA Runtime 8 └── Memory Allocator
Therefore:
1Available VRAM 2= 3Total VRAM 4- 5Weights 6- 7KV Cache 8- 9Activations 10- 11Runtime Buffers 12- 13Other allocations
This is why a model that technically fits into VRAM may still fail during inference.
24. Estimating KV Cache Memory
A simplified KV-cache memory calculation is:
1KV Memory ≈ 22 × Layers × Tokens × KV Heads × Head Dimension × Bytes
The 2 represents:
1K + V
For example, conceptually:
1kv_bytes = ( 2 2 3 * num_layers 4 * sequence_length 5 * num_kv_heads 6 * head_dim 7 * bytes_per_element 8)
For batched inference:
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)
This simplified equation is extremely useful for understanding why long context and large batches consume substantial VRAM.
25. Grouped-Query Attention
Modern LLMs may use fewer K/V heads than query heads.
For example:
1Query heads = 32 2KV heads = 8
Instead of storing:
132 K heads 232 V heads
the model stores:
18 K heads 28 V heads
This can substantially reduce KV-cache memory.
Conceptually:
1Q heads 232 ───────────────┐ 3 │ 4 ▼ 5 Attention 6 ▲ 7 │ 8KV heads 98 ───────────────┘
This is one architectural technique for reducing inference memory requirements.
26. Sampling
After the model produces logits:
1logits 2 ↓ 3probabilities 4 ↓ 5sampling 6 ↓ 7next token
Greedy sampling:
1next_token = torch.argmax( 2 logits, 3 dim=-1 4)
Temperature scaling:
1temperature = 0.7 2 3scaled_logits = ( 4 logits / temperature 5)
Then:
1probabilities = torch.softmax( 2 scaled_logits, 3 dim=-1 4)
Sample:
1next_token = torch.multinomial( 2 probabilities, 3 num_samples=1 4)
Lower temperature generally makes the distribution sharper; higher temperature generally makes it flatter.
27. A Minimal LLM Inference Engine
A conceptual architecture:
1 ┌──────────────┐ 2 │ Request │ 3 └──────┬───────┘ 4 ↓ 5 ┌──────────────┐ 6 │ Tokenizer │ 7 └──────┬───────┘ 8 ↓ 9 ┌──────────────┐ 10 │ Scheduler │ 11 └──────┬───────┘ 12 ↓ 13 ┌──────────────┐ 14 │ Model │ 15 └──────┬───────┘ 16 ↓ 17 ┌──────────────────┐ 18 │ CUDA Kernels │ 19 └────────┬─────────┘ 20 ↓ 21 ┌──────────────┐ 22 │ KV Cache │ 23 └──────┬───────┘ 24 ↓ 25 ┌──────────────┐ 26 │ Sampling │ 27 └──────┬───────┘ 28 ↓ 29 Generated Token
A simplified Python skeleton:
1class InferenceEngine: 2 3 def __init__(self, model, tokenizer): 4 self.model = model.cuda() 5 self.model.eval() 6 7 self.tokenizer = tokenizer 8 9 @torch.inference_mode() 10 def generate( 11 self, 12 prompt, 13 max_tokens=50 14 ): 15 input_ids = self.tokenizer( 16 prompt, 17 return_tensors="pt" 18 ).input_ids.cuda() 19 20 for _ in range(max_tokens): 21 22 output = self.model( 23 input_ids, 24 use_cache=True 25 ) 26 27 logits = output.logits[:, -1, :] 28 29 next_token = torch.argmax( 30 logits, 31 dim=-1, 32 keepdim=True 33 ) 34 35 input_ids = torch.cat( 36 [input_ids, next_token], 37 dim=1 38 ) 39 40 return self.tokenizer.decode( 41 input_ids[0] 42 )
This is intentionally simplified. A production engine would separate:
1Request Management 2Scheduler 3Batch Manager 4KV Cache Manager 5Model Executor 6CUDA Runtime 7Sampler 8Output Manager
28. Multi-GPU Inference
A model may be too large for a single GPU.
Suppose:
1Model 2 ↓ 340 GB
but:
1GPU 0 = 24 GB 2GPU 1 = 24 GB
The model can potentially be distributed across GPUs.
Conceptually:
1 Model 2 │ 3 ┌───────┴───────┐ 4 ↓ ↓ 5 GPU 0 GPU 1 6 Layers Layers 7 0–N N–M
This is a form of model parallelism.
Communication between GPUs becomes important.
29. NCCL
NCCL is NVIDIA's collective communication library.
It provides primitives for communication across GPUs.
Examples include:
1AllReduce 2Broadcast 3Reduce 4AllGather 5ReduceScatter
Conceptually:
1GPU 0 ─────┐ 2 │ 3GPU 1 ─────┼──► Collective Communication 4 │ 5GPU 2 ─────┤ 6 │ 7GPU 3 ─────┘
For distributed AI systems:
1Computation 2 + 3Communication 4 = 5Overall Performance
A fast kernel does not help much if GPUs spend most of their time waiting for communication.
30. Multi-GPU Bottleneck
Consider:
1GPU 0 2 │ 3 │ computation 4 ▼ 5Communication 6 │ 7 ▼ 8GPU 1 9 │ 10 │ computation 11 ▼ 12Communication
If communication dominates:
1Compute █████ 2Communication █████████████
the system may be communication-bound.
Therefore multi-GPU optimization requires understanding:
- GPU topology
- Communication volume
- Synchronization
- Collective operations
- Computation/communication overlap
31. Inference Profiling
A useful profiling workflow:
1Run workload 2 ↓ 3Measure latency 4 ↓ 5Measure throughput 6 ↓ 7Profile GPU 8 ↓ 9Inspect kernels 10 ↓ 11Inspect memory 12 ↓ 13Inspect communication 14 ↓ 15Optimize 16 ↓ 17Benchmark again
Important metrics include:
Time to First Token
TTFT measures how long the user waits before receiving the first generated token.
1Request 2 ↓ 3Prefill 4 ↓ 5First token
Time Per Output Token
This measures generation speed after the first token.
1Token 1 2 ↓ 3Token 2 4 ↓ 5Token 3 6 ↓ 7Token 4
Tokens Per Second
A common throughput metric:
1tokens / second
GPU Utilization
Useful, but do not interpret it alone.
High utilization does not necessarily mean optimal performance.
32. Latency vs Throughput
These are different goals.
Latency
How quickly one request completes.
1Request 2 ↓ 3Response
Throughput
How much work the system processes over time.
1Requests 2████████████████ 3 ↓ 4tokens/sec
A larger batch may improve throughput:
1Batch 1 → 100 tokens/s 2Batch 8 → 700 tokens/s
but could increase individual request latency.
A production inference system must balance both.
33. Inference Optimization Stack
Think of optimization from the bottom upward:
1Application 2 ↓ 3Scheduler 4 ↓ 5Batching 6 ↓ 7KV Cache 8 ↓ 9Model Graph 10 ↓ 11Kernel Fusion 12 ↓ 13CUDA Kernels 14 ↓ 15Tensor Cores 16 ↓ 17GPU Memory 18 ↓ 19Multi-GPU Communication
Every layer can become a bottleneck.
34. Complete Optimization Strategy
When an inference engine is slow, do not immediately rewrite CUDA kernels.
Use:
11. Verify correctness 2 ↓ 32. Establish baseline 4 ↓ 53. Measure TTFT 6 ↓ 74. Measure token latency 8 ↓ 95. Measure throughput 10 ↓ 116. Check GPU memory 12 ↓ 137. Profile kernels 14 ↓ 158. Check KV cache 16 ↓ 179. Check batching 18 ↓ 1910. Check communication 20 ↓ 2111. Optimize bottleneck 22 ↓ 2312. Benchmark again
This prevents optimization based on assumptions.
35. Production Inference Architecture
A more complete system looks like:
1 Clients 2 │ 3 ▼ 4 ┌─────────────────┐ 5 │ API Gateway │ 6 └────────┬────────┘ 7 ↓ 8 ┌─────────────────┐ 9 │ Scheduler │ 10 └────────┬────────┘ 11 ↓ 12 ┌──────────────────┐ 13 │ Continuous Batch │ 14 └────────┬─────────┘ 15 ↓ 16 ┌──────────────────────┐ 17 │ Model Executor │ 18 └──────────┬───────────┘ 19 ↓ 20 ┌────────────────────────┐ 21 │ CUDA / Tensor Cores │ 22 └───────────┬────────────┘ 23 ↓ 24 KV Cache 25 │ 26 ↓ 27 Sampling 28 │ 29 ▼ 30 Tokens
For multiple GPUs:
1 Scheduler 2 │ 3 ┌──────────┼──────────┐ 4 ↓ ↓ ↓ 5 GPU 0 GPU 1 GPU 2 6 │ │ │ 7 └──────────┼──────────┘ 8 ↓ 9 NCCL
36. Final Project — Build an LLM Inference Engine
The course project should implement:
1Tokenizer 2 ↓ 3Embedding 4 ↓ 5Transformer 6 ↓ 7CUDA Kernels 8 ↓ 9KV Cache 10 ↓ 11Sampling 12 ↓ 13Generated Tokens
Then progressively add:
1Phase 1 2Basic inference 3 ↓ 4Phase 2 5GPU model loading 6 ↓ 7Phase 3 8KV cache 9 ↓ 10Phase 4 11Batch processing 12 ↓ 13Phase 5 14Continuous batching 15 ↓ 16Phase 6 17Quantization 18 ↓ 19Phase 7 20Kernel fusion 21 ↓ 22Phase 8 23CUDA Graphs 24 ↓ 25Phase 9 26Multi-stream execution 27 ↓ 28Phase 10 29Multi-GPU + NCCL 30 ↓ 31Phase 11 32Profiling
37. Project Benchmark
Record at least:
| Configuration | TTFT | Token Latency | Tokens/sec | VRAM |
|---|---|---|---|---|
| FP16 | — | — | — | — |
| BF16 | — | — | — | — |
| INT8 | — | — | — | — |
| INT4 | — | — | — | — |
| Batched | — | — | — | — |
| Continuous Batch | — | — | — | — |
| CUDA Graph | — | — | — | — |
| Multi-GPU | — | — | — | — |
The goal is not simply to achieve the highest tokens/second.
You should be able to explain why performance changed.
38. What You Should Be Able to Do
After completing this course, you should understand:
1✓ GPU inference architecture 2✓ Model loading 3✓ Weight memory 4✓ FP16 5✓ BF16 6✓ FP8 7✓ INT8 8✓ INT4 9✓ KV cache 10✓ Prefill 11✓ Decode 12✓ Batch processing 13✓ Continuous batching 14✓ Memory fragmentation 15✓ Paged KV-cache concepts 16✓ Kernel fusion 17✓ CUDA Graphs 18✓ CUDA streams 19✓ Tensor Cores 20✓ Multi-GPU inference 21✓ NCCL 22✓ Inference profiling 23✓ TTFT 24✓ Token latency 25✓ Tokens/sec 26✓ GPU memory analysis
Most importantly, you should be able to reason about an inference engine as a complete GPU system, rather than viewing the neural network as only a collection of Python layers.
1Model 2 ↓ 3Memory 4 ↓ 5Kernels 6 ↓ 7GPU Scheduling 8 ↓ 9KV Cache 10 ↓ 11Batching 12 ↓ 13Communication 14 ↓ 15Inference Performance
The central engineering principle is:
Measure the complete inference pipeline, identify the actual bottleneck, and optimize the bottleneck rather than optimizing components in isolation.