TPU Deep Dive
Course
Review
Slides
Exams
Sources
Theme: system
Sources
PL-tpu
Writing TPU kernels
PL-grid
Grids and BlockSpecs
PL-pipe
TPU pipelining
PL-mm
TPU matmul tutorial
PL-sparse
Scalar prefetch and block-sparse
PL-dist
Distributed Pallas
PL-core
Core-specific programming
PL-gpu
Pallas Mosaic GPU
K-tokamax
Tokamax kernel library
SB9
How to Profile TPU Programs
The JAX profiler
SB9
How to Profile TPU Programs
Reading an XLA op
SB9
How to Profile TPU Programs
Example profile
SB9
How to Profile TPU Programs
Problems
XP-home
XProf
XP-trace
Trace Viewer
XP-opprof
HLO Op Profile
XP-roof
Roofline Model
XP-mem
Memory Viewer
XP-jax
Profiling JAX programs
R-tpuinf-prof
tpu-inference profiling guide
X-tiled
Tiled layout
X-async
Async HLO ops
A-nki
NKI programming model
SB1
All About Rooflines
Tile intensity
Chapter 10. Day 5, second block, 55 minutes
Pallas, XProf and a design review
Read
Quiz
12
Flashcards
21
Teach-back
7
Question 1 of 12
Recall
How does a Pallas TPU grid execute?
1
Sequentially in lexicographic order on a TensorCore; parallel axes can be split across a megacore's two cores.
2
Concurrently, one grid point per thread block, like CUDA.
3
In random order chosen by the compiler.
4
Once per MXU.
Check answer
Enter
Press 1–4 to choose.