Chapter 1: A Brief Introduction to Roofline Analysis

Scaling Book Exercises – Chapter 1

Chapter: 1 – All About Rooflines
Source: handwritten Onyx Boox A4 PDF
Note: Transcribed from handwritten solutions; book markdown used only for exercise statements and notation.

Exercise 1 – int8 matmul

Exercise statement

Say we want to do the matmul 𝑋[𝐵,𝐷]⋅𝐷𝑌[𝐷,𝐹]→𝑍[𝐵,𝐹] in int8 precision (1 byte per parameter) instead of bfloat16 (2 bytes per parameter) since TPUs/GPUs can do matmuls faster in lower precision.

  1. How many bytes need to be loaded from memory? How many need to be written back to memory?
  2. How many total OPs are performed?
  3. What is the arithmetic intensity?
  4. What is a roofline estimate for 𝑇math and 𝑇comms? What are reasonable upper and lower bounds for the runtime of the whole operation?

Assume our HBM bandwidth is 8.2⋅1011bytes/s and our int8 peak OPs/s is 3.94⋅1014 (about 2x bfloat16).

Solution

Scan pages: 1-2

Assume the operation is 𝑋[𝐵,𝐷]⋅𝐷𝑌[𝐷,𝐹]→𝑍[𝐵,𝐹]. Then, bytes in int8:

1 ⋅ ( 𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹 ) = 𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹 .

There are 𝐵𝐹 dot products. A dot product has 𝐷 multiplications and computes

𝑎 1 𝑏 1 + 𝑎 2 𝑏 2 + … + 𝑎 𝐷 𝑏 𝐷

hence 𝐷−1 sums. Total OPs is then:

OPs = 𝐵 𝐹 ( 𝐷 + 𝐷 − 1 ) = 𝐵 𝐹 ( 2 𝐷 − 1 ) .

Arithmetic Intensity:

AI = 𝐵 𝐹 ( 2 𝐷 − 1 ) 𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹 .

Roofline estimates:

𝑇 math = 𝐵 𝐹 ( 2 𝐷 − 1 ) 3.94 ⋅ 10 14 FLOPs/s ≈ 𝐵 𝐹 ( 𝐷 − 1 ) 2 ⋅ 10 14 s . 𝑇 comms = 𝐵 𝐷 + 𝐷 𝐹 + 𝐵 𝐹 8.2 ⋅ 10 11 Bytes/s .

Bounds of the whole operation:

max ( 𝑇 comms , 𝑇 math ) ≤ 𝑇 operation ≤ 𝑇 comms + 𝑇 math .

Exercise 2 – int8 weights and bfloat16 activations

Exercise statement

In practice we often do different weight vs. activation quantization, so we might store our weights in very low precision but keep activations (and compute) in a higher precision. Say we want to quantize our weights in int8 but keep activations (and compute) in bfloat16. At what batch size do we become compute bound? Assume 1.97⋅1014 bfloat16 FLOPs/s.

Hint: this means specifically bf16[B, D] * int8[D, F] -> bf16[B, F] where 𝐵 is the “batch size”.

Solution

Scan pages: 3-4

Let the accelerator FLOPs/s for bf16 be:

1.97 ⋅ 10 14 ≈ 2 ⋅ 10 14 .

Assume the operation is:

bf16 [ 𝐵 , 𝐷 ] ⋅ 𝐷 int8 [ 𝐷 , 𝐹 ] → bf16 [ 𝐵 , 𝐹 ]

Assuming no conversion needed to cast int8 to bf16, we get:

AI = 𝐵 𝐹 ( 2 𝐷 − 1 ) 2 𝐵 𝐷 + 𝐷 𝐹 + 2 𝐵 𝐹 ≥ 1.97 ⋅ 10 14 FLOPs/s 8.2 ⋅ 10 11 Bytes/s ≈ 2 8 ⋅ 10 14 10 11 = 250 FLOPs/Byte .

Since:

𝐷 ≫ 𝐵 ∧ 𝐹 ≫ 𝐵 ∧ 𝐷 𝐹 ≫ 2 𝐵 𝐷 ∧ 𝐷 𝐹 ≫ 2 𝐵 𝐹 ⇒ 𝐷 𝐹 ≫ 2 𝐵 𝐷 + 2 𝐵 𝐹 .

we can approximate:

AI ≈ 2 𝐵 𝐹 𝐷 𝐷 𝐹 .

Hence:

2 𝐵 𝐹 𝐷 𝐷 𝐹 > 250 ⇒ 𝐵 > 125 .

“Per second I do more starts of the algorithm since less bytes to be transferred, so then per-algo FLOPs can be smaller.”

Exercise 3 – roofline plot for two matrix sizes

Exercise statement

Taking the setup from Question 2, make a roofline plot of peak FLOPs/s vs. 𝐵 for 𝐹=𝐷=4096 and 𝐹=𝐷=1024. Use the exact number of bytes loaded, not an approximation.

Solution

Scan pages: 5-8

Let:

  • Accelerator bfloat16 = 1.97⋅1014.
  • Accelerator HBM = 8.2⋅1011.

Bandwidth limited performance can be defined as AI⋅BW. Therefore the operation from Question 2 has:

Achievable FLOPs/s = min ( AI ⋅ BW , Accelerator bfloat16 ) = min ( 2 𝐵 𝐷 𝐹 2 𝐵 𝐷 + 𝐷 𝐹 + 2 𝐵 𝐹 ⋅ BW , 1.97 ⋅ 10 14 ) .

Assuming 𝐷=𝐹 the AI becomes:

AI = 2 𝐵 𝐷 2 2 𝐵 𝐷 + 𝐷 2 + 2 𝐵 𝐷 = 2 𝐵 𝐷 2 4 𝐵 𝐷 + 𝐷 2 = 2 𝐵 𝐷 2 𝐷 ( 4 𝐵 + 𝐷 ) = 2 𝐵 𝐷 4 𝐵 + 𝐷 .

Let 𝐷=2𝑘:

AI = 2 𝐵 2 𝑘 4 𝐵 + 2 𝑘 = 2 𝐵 2 𝑘 2 𝑘 ( 𝐵 2 𝑘 − 2 + 1 ) = 2 𝐵 𝐵 2 𝑘 − 2 + 1 .

Note that 𝑘=12 for 𝐷=4096 and 𝑘=10 for 𝐷=1024. Taking logarithm we get:

log ( AI ⋅ BW ) = log ( 2 ⋅ 𝐵 ) − log ( 𝐵 2 𝑘 − 2 + 1 ) + log ( BW ) ≈ log ( 2 ⋅ 𝐵 ) − 𝐵 2 𝑘 − 2 + log ( BW ) .

For 𝐷=4096:

2 𝐵 ⋅ 4096 4 𝐵 + 4096 > 250 FLOPs/token ⇔ 𝐵 ≈ 142 .

For 𝐷=1024:

2 𝐵 ⋅ 1024 4 𝐵 + 1024 > 250 ⇔ 𝐵 ≈ 244 .

Exercise 4 – different matrix for each batch element

Exercise statement

What if we wanted to perform int8[𝐵,𝐷]⋅𝐷int8[𝐵,𝐷,𝐹]→int8[𝐵,𝐹] where we imagine having a different matrix for each batch element. What is the arithmetic intensity of this operation?

Solution

Scan page: 9

Assume the operation is:

𝐵 times int8 [ 𝐷 ] ⋅ 𝐷 int8 [ 𝐷 , 𝐹 ] → int8 [ 𝐹 ]

Hence:

FLOPs = 𝐵 ( 2 𝐹 𝐷 ) .

Therefore:

AI = 𝐵 ⋅ 2 𝐹 𝐷 𝐵 ( 𝐷 + 𝐷 𝐹 + 𝐹 ) = 2 𝐹 𝐷 𝐷 + 𝐷 𝐹 + 𝐹 .

Exercise 5 – memory roofline for H100 SXM

Exercise statement

Using the spec sheet provided by NVIDIA for the H100 SXM, calculate the batch size at which a bfloat16 matrix multiplication will become compute-bound. Note that the Tensor Core FLOPs numbers are twice the true value since they are only achievable with structured sparsity.

Solution

Scan pages: 9-10

Let Accelerator FLOPs = 1979⋅10122≈989.5⋅1012≈9.89⋅1014. To be computed bound we need:

AI = 2 𝐵 𝐹 𝐷 2 𝐵 𝐷 + 2 𝐷 𝐹 + 2 𝐵 𝐹 ≥ 9.89 ⋅ 10 14 3.35 ⋅ 10 12 .

Assume 𝐷≫𝐵 and 𝐹≫𝐵. Hence:

2 𝐵 𝐹 𝐷 2 𝐷 𝐹 ≥ 295 ⇔ 𝐵 ≥ 295 .