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 + Linear in a single kernel
  • Tile-based computation for cache efficiency
  • Memory bandwidth analysis: fused vs unfused
  • Benchmarking with triton.testing.Benchmark

Key concepts

ConceptWhat to understand
Triton programming modelBlocks of threads operate on tiles
tl.load / tl.storeCoalesced memory access patterns
Fusion benefitAvoid extra HBM round-trips for intermediate results
@triton.jitJIT-compiles kernel at first call

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyComplete working implementation
test_lab.pypytest suite (CUDA-only)
requirements.txtDependencies

Run

pip install -r requirements.txt  # requires CUDA GPU
python lab.py
pytest test_lab.py -v