Fusion byte ledger
Count round trips for an eager elementwise chain, predict the fused traffic, then explain a speedup in bytes rather than in adjectives.
- Module
- PyTorch and the memory wall
- Objective
- Predict a fusion speedup from traffic ratios, and notice when launch overhead destroys the prediction.
def chain(x):
return torch.relu(x * 2.0 + 1.0) * 0.5
x is fp16, shape (B, T, H).
Part A — Eager traffic
Eager PyTorch typically launches several pointwise kernels. A conservative model: each kernel reads its input and writes its output, intermediates live in HBM.
Count bytes for one call, as a multiple of numel(x) × 2.
List a plausible kernel sequence (mul, add, relu, mul). For each: bytes in, bytes out. Total / |x|.
Fused: one read of x, one write of y. Total / |x|.
Predicted speedup if you are purely memory-bound and hit the same GB/s: eager/fused.
Part B — Three sizes
You measure (made-up but consistent) on a 2.0 TB/s card:
| numel | eager µs | fused µs |
|---|---|---|
| 4,096 | 14 | 12 |
| 1e6 | 95 | 28 |
| 1.34e8 (128 MiB fp16) | 210 | 72 |
- Predicted ratio from Part A vs measured ratio at each size.
- Which row is overhead-bound? Evidence?
- Which row is the traffic story landing?
- Write the one-paragraph explanation in bytes for the large row. No "optimized." No "faster."
Part C — torch.compile discipline
List four ways to time torch.compile and get a nonsense number. For each, the one-line fix.
Part D — When fusion is the wrong week
A 4096 GEMM is 70% of peak TFLOP/s. Someone wants to fuse the following ReLU into the GEMM epilogue. Using regimes: is that a week, or a Tuesday afternoon epilogue, or neither because it will not move the trace?
Acceptance
- Fused traffic is 2×|x|; eager is several times that (you showed the sum).
- Tiny numel: ratio collapses toward 1. Large numel: ratio approaches traffic ratio.
- Compile first-call, no warmup, no sync, CPU timer — all named as invalid.
- Fusing ReLU into a compute-bound GEMM is epilogue candy, not a 3×.
Check
Eager example: 4 kernels × (read+write) ≈ 8×|x| if every intermediate is materialised; fused 2×; naive 4×. Real eager may fuse some pointwise already — if measured << predicted, say so.
4,096 elements: 8 KiB. At 2 TB/s that is 4 ns of traffic; 14 µs is ~launch. 128 MiB: 64 ms at 2 TB/s for a 2× round trip... wait: 128 MiB = 1.34e8×2 = 256 MiB of |x|. Fused 2×|x| = 512 MiB ≈ 0.26 ms at 2 TB/s. Measured 72 µs is faster than that floor? Then either numel was elements not bytes, or the table's 1.34e8 is elements: |x|=256 MiB, fused 512 MiB, 256 µs floor — 72 µs would then mean you over-counted traffic or hit cache. The drill is to notice the contradiction, not to bless the table. Write that paragraph.
(Use |x| bytes = numel×2. Large row: 256 MiB tensor, fused ~512 MiB traffic, floor 0.26 ms. Measured 72 µs is below floor → table is illustrative; on a real GPU you would remeasure. A passing answer flags this.)
Debrief
Fusion is a traffic delete. If time does not scale with the deleted bytes, you were never in the memory regime.