Performance-Model a Sharded Matrix Multiply: FLOPs, Memory, Network and Roofline

This round is a performance-modeling exercise for machine learning systems, done by hand. Consider the matrix multiplication `C = A × B`, where `A` has shape `M × K`, `B` has shape `K × N`, and every element takes `s` bytes (for example `s = 2` for 16-bit formats). An accelerator has a peak compute rate of `P` floating-point operations per second (FLOP/s) and a memory bandwidth of `W` bytes per second. Devices are connected by links that move `L` bytes per second per device. Work from first principles: count floating-point operations, memory traffic and network traffic, and use a roofline argument to decide what limits the performance of each configuration. ### Clarifying Questions - Are `A` and `B` both activations, or is one of them a weight matrix that stays resident on the devices? - Must the result `C` end up complete on every device, or may it stay split across devices? - Can communication overlap with computation, or does each step wait for the previous one to finish? - Which interconnect is assumed: a ring of point-to-point links, a switch, or something else? - Should the estimate assume ideal on-chip data reuse, or account for limited on-chip memory? ### Part 1 — One device Derive the number of floating-point operations and the minimum number of bytes that must move between memory and the compute units. Compute the arithmetic intensity, and use the roofline model to decide when the multiplication is compute-bound and when it is memory-bound. Apply this to two cases: a large square multiplication, and a multiplication in which `M` is small (such as a small batch of tokens multiplied by a large weight matrix during generation). ```hint Intensity versus the ridge Compare operations per byte of the kernel with the hardware's operations-per-byte balance point, and watch which matrix dominates the byte count when one dimension is small. ``` #### What This Part Should Cover - Correct operation and byte counts, and the arithmetic intensity formula - The roofline ridge point and the resulting classification - How intensity changes with the shape, especially when `M` is small - The gap between the ideal byte count and real memory traffic ### Part 2 — Sharding across D devices Now split the multiplication across `D` devices. For each of the three ways to shard it (along `M`, along `N`, or along the shared dimension `K`), work out the per-device computation, the per-device memory traffic, and the network traffic needed to produce the result. Decide in each case whether the step is limited by compute, memory or network, and derive the condition under which sharding along `K` becomes network-bound. ```hint Where do partial results live For each split, ask whether any device ends up holding only part of a sum that must be combined with other devices' parts before `C` is correct. ``` #### What This Part Should Cover - Per-device operation, memory and network counts for each sharding choice - Collective operations needed (such as all-gather or reduce-scatter) and their cost - A network-versus-compute comparison and the threshold it produces - How overlapping communication with computation changes the estimate ### Part 3 — Put numbers on it Pick an illustrative accelerator and interconnect, state the numbers clearly as assumptions, and estimate the time for a concrete multiplication under at least two sharding choices. Say which resource is the bottleneck in each case, and what you would change to improve it. ```hint Take the maximum, then justify it Estimate each resource's time separately, then decide whether they overlap, in which case the largest dominates, or run in sequence and add up. ``` #### What This Part Should Cover - Clearly stated hardware assumptions - Correct arithmetic, with units carried through - Identification of the bottleneck and a concrete improvement - A sanity check of the result against the roofline ### What a Strong Answer Covers - Precise counting with a clear distinction between operations, memory bytes and network bytes - Correct use of the roofline model, including the ridge point - Sharding analysis that identifies the required collective operations and their cost - Quantitative conclusions rather than qualitative guesses ### Follow-up Questions - How does tiling for on-chip memory change the real memory traffic, and how large must a tile be to reach the ideal? - For a transformer's feed-forward layer, which sharding would you choose for training and which for low-batch inference? - How do lower-precision formats change each of the three limits? - Two-dimensional sharding splits the work along two dimensions at once. When does it beat splitting along one?