TPU Deep Dive
Course
Review
Slides
Exams
Sources
Theme: system
Sources
SB4
All the Transformer Math You Need to Know
Counting dots
SB4
All the Transformer Math You Need to Know
Global FLOPs and params
SB4
All the Transformer Math You Need to Know
Attention share
SB4
All the Transformer Math You Need to Know
MoE
SB4
All the Transformer Math You Need to Know
Remat
SB4
All the Transformer Math You Need to Know
KV caching
SB4
All the Transformer Math You Need to Know
Problems
SB4
All the Transformer Math You Need to Know
Flash attention
SB5
How to Parallelize a Transformer for Training
LLaMA-2-13B problems
SB6
Training LLaMA 3 on TPUs
LLaMA-3-70B counts
SB7
All About Transformer Inference
Attention intensity
SB7
All About Transformer Inference
Worked problems
SB10
Programming TPUs in JAX
MoE with ragged_dot
R-tpuinf
tpu-inference (being renamed vLLM JAX)
megablox, fused MoE kernels
D-7x
TPU7x (Ironwood)
G-8dd
Gupta & Mugazambi, "TPU 8t and TPU 8i technical deep dive"
Chapter 4. Day 2, second block, 55 minutes
Transformer accounting on TPU
Read
Quiz
12
Flashcards
20
Teach-back
7
Question 1 of 12
Derivation: enter a number
How many FLOPs does the contraction
A
[
I
,
J
,
K
,
L
]
⋅
B
[
I
,
J
,
M
,
N
,
O
]
→
C
[
K
,
L
,
M
,
N
,
O
]
A[I,J,K,L] \cdot B[I,J,M,N,O] \to C[K,L,M,N,O]
A
[
I
,
J
,
K
,
L
]
⋅
B
[
I
,
J
,
M
,
N
,
O
]
→
C
[
K
,
L
,
M
,
N
,
O
]
take, for
I
=
J
=
K
=
L
=
M
=
N
=
O
=
4
I = J = K = L = M = N = O = 4
I
=
J
=
K
=
L
=
M
=
N
=
O
=
4
?
FLOPs
Check answer
Enter