Counting FLOPs and Bytes by Hand

Why this matters

Being able to derive FLOPs and bytes moved from raw tensor shapes — no profiler required — is a load-bearing skill you'll lean on in every module that follows. A matmul of an (M,K) tensor by a (K,N) tensor does 2MN*K FLOPs (one multiply and one add per output element, summed K times) and, in the naive case, reads A and B once and writes C once. Once you can do this arithmetic in your head, you can predict whether a shape is efficient before you run it, and you can catch a profiler lying to you — trace tools occasionally misattribute time or double count, and the fastest way to notice is to already know roughly what the number should be. This also sets up the memory-budget task next: parameter count, activation size, and optimizer state are all just this same shape arithmetic applied to a whole model instead of one op.

Demo

The demo computes FLOPs and bytes for several matmul shapes, including the 'thin' shape typical of single-token autoregressive decoding and the 'fat' shape typical of a training batch, and shows why one is memory-bound and the other isn't just from the numbers.

Try it yourself

1. Change dtype_bytes from 2 (fp16) to 4 (fp32) and to 1 (int8) for the decode step, and note how much the arithmetic intensity moves versus the training-batch shape.
2. Compute FLOPs and bytes by hand for your own model's largest linear layer shape and check it against what the function prints.
3. Explain in a comment why the decode-step shape (M=1) has such low arithmetic intensity compared to the training-batch shape, in terms of how much of B gets reused per byte loaded.
4. Add a fourth shape representing a batch-8 decode step (M=8) and observe how much arithmetic intensity improves versus M=1 — this is the entire argument for batching inference requests.

References

• EleutherAI — Transformer Math 101
• kipply — Transformer Inference Arithmetic