Lab 03 — Triton Fused Kernel
Phase: 04 — torch.compile, FX Graph & Custom Backends | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours
Write a GPU kernel in Triton that fuses multiple operations into a single pass over memory.
What you build
- Triton kernel for fused
LayerNorm + Linearin a single kernel - Tile-based computation for cache efficiency
- Memory bandwidth analysis: fused vs unfused
- Benchmarking with
triton.testing.Benchmark
Key concepts
| Concept | What to understand |
|---|---|
| Triton programming model | Blocks of threads operate on tiles |
tl.load / tl.store | Coalesced memory access patterns |
| Fusion benefit | Avoid extra HBM round-trips for intermediate results |
@triton.jit | JIT-compiles kernel at first call |
Files
| File | Purpose |
|---|---|
lab.py | Stubs with TODO markers |
solution.py | Complete working implementation |
test_lab.py | pytest suite (CUDA-only) |
requirements.txt | Dependencies |
Run
pip install -r requirements.txt # requires CUDA GPU
python lab.py
pytest test_lab.py -v