Reading Grzegorz Kwasniewski et al., Red-blue pebbling revisited: near optimal parallel matrix-matrix multiplication (2019, arXiv v3). Original paper
Suppose a matrix multiplication requires a trillion multiply-adds. Adding processors seems as though it should keep shortening the runtime. Yet each processor must receive its inputs, and partial results from different processors may need to be combined. As arithmetic gets faster, movement can become the bottleneck.
COSMA begins with a question: what reuse does finite memory permit? It then seeks a work decomposition with as little communication as possible and turns that decomposition into parallel matrix multiplication.
Its focus is classical matrix multiplication, where each output is an ordinary sum of products. The same optimality claim does not include Strassen-like algorithms.
Three indices form a three-dimensional computation space
Let
Each output entry is
A scalar multiply-add corresponds to a triple $(i,j,\ell)$. The whole computation can therefore be pictured as an $m\times n\times k$ box containing $mnk$ multiply-add positions.
Suppose a processor receives an $a\times b\times c$ sub-box. The three data sets it needs are exactly the projections of that box onto three coordinate planes:
| Data | Local shape | Elements |
|---|---|---|
| $A$ | $a\times c$ | $ac$ |
| $B$ | $c\times b$ | $bc$ |
| Partial sums of $C$ | $a\times b$ | $ab$ |
The local work is $abc$, and the boundary volume of inputs and final local outputs is $ac+bc+ab$. A large computation volume with a small boundary means substantial reuse.
Why a square root appears in the lower bound
Let fast memory hold $S$ elements. If three square tiles of side length $t$ all reside there, then
Loading $\Theta(t^2)$ numbers supports $\Theta(t^3)$ multiply-adds. Each transferred number supports $\Theta(t)$ multiply-adds, so the reuse scales as $\sqrt S$, rather than $S$.
The whole multiplication has $mnk$ multiply-adds. The leading term of the classical I/O lower bound is therefore
This square-tile example explains the exponent $1/2$, but does not prove a bound for all schedules. A rigorous proof must handle irregular tiles, partial sums, recomputation, and reuse across phases.
COSMA starts from red-blue pebbling and graph partitions to make the reuse analysis more precise, paying particular attention to constants hidden by asymptotic notation.
A better sequential schedule: retain outputs and stream inputs
Keeping three large tiles simultaneously is not the only strategy. We can retain an $a\times b$ output tile in fast memory while streaming column segments of $A$ and row segments of $B$ through it.
For a fixed $\ell$, the update is an outer product:
Reading $a+b$ numbers completes $ab$ multiply-adds, giving reuse efficiency
For a fixed output area $ab=R$, the arithmetic–geometric mean inequality gives
so efficiency is greatest at $a=b=\sqrt R$:
Most of the memory should therefore hold an output tile as close to square as possible, supported by smaller input buffers.
Each output tile is written only after traversing the full $k$ dimension. Ignoring edges, divisibility, and buffer space, tiles with side length roughly $\sqrt S$ give leading input traffic
Add $mn$ elements for the final output. This explains the coefficient 2 that the paper studies closely, and why repeatedly writing partial sums wastes communication.
In a real schedule, not all $S$ positions can hold output: inputs and newly computed values also need space. The paper explicitly handles this executability constraint. Its near-optimal factor has the form
which approaches one as $S$ grows. A figure such as 0.03% describes closeness to a communication guarantee within the model; it does not mean runtime is within 0.03% of hardware peak performance.
From fixed partitions to proofs that account for reuse
The classical Hong–Kung argument cuts an execution by I/O count and uses partition boundaries to limit the work in each phase. COSMA further distinguishes values already resident at phase entry, values that must be newly loaded, and results that must be saved.
Its reuse lemma uses $X$ as the partition parameter, $R(S)$ as an upper bound on the reuse set, and $T(S)$ as a lower bound on necessary stores, obtaining
where $H(X)$ is the minimum number of subcomputations in a valid partition.
This is more precise than simply assigning $S$ transfers to every phase. Existing values can be reused, but all of fast memory cannot simultaneously be reserved for reuse while additional inputs occupy space outside that budget. Executability must be counted too.
The paper first analyzes greedy schedules that reuse data promptly, then argues that non-greedy schedules cannot beat the corresponding lower bound. Greedy has a formal definition here; it does not mean an arbitrary local heuristic.
Any concrete bound also requires checking the input/output convention and dimension regime. Traffic changes if inputs already occupy the target memory level, or the entire problem fits there. Usually we should also check the compulsory cost
under the usual convention that inputs start in slow memory and outputs must be written back. The communication leading term cannot be extrapolated indefinitely without these conditions.
When should parallel multiplication split the third dimension?
Suppose there are $p$ processors, each with local capacity $S$. With balanced work, each processor performs
multiply-adds.
Splitting only the $i,j$ dimensions gives each processor a final output tile accumulated over the complete $\ell$ dimension. Its partial sums need no cross-processor reduction. But as processor count grows, output tiles shrink and input reuse declines.
Splitting $\ell$ as well lets several processors compute different contributions to the same output tile:
This preserves larger local output tiles but requires communication to reduce partial sums. The choice between two and three dimensions should follow the problem shape, processor count, and available memory.
A local optimization we can solve by hand
Consider a local work domain $a\times a\times b$. Its work is $a^2b=W$, and its ideal boundary communication is
Substitute $b=W/a^2$:
The first term falls as $a$ grows because input reuse improves. The second rises because the partial-result tile grows.
Differentiate:
Setting the derivative to zero gives
When memory permits, the ideal local domain is a cube, with
But the output tile must satisfy $a^2\le S$. If the cube’s face does not fit, cap $a$ at $\sqrt S$ and extend the accumulation dimension. The ideal solution becomes
Its continuous geometric cost is
This is an ideal model for understanding the design. It ignores integer processor grids, edge tiles, input buffers, and dimension limits. A local domain cannot exceed a short matrix dimension; the actual decomposition must be clipped and redistributed.
The two branches cannot simply be compared and the smaller selected unconditionally. The cubic branch is feasible only if its face fits in memory.
A numerical example: memory changes the optimal shape
Suppose a processor is assigned $W=2^{30}$ multiply-adds.
With enough memory, its local domain is
and the ideal boundary volume is
elements.
If local capacity is only $S=2^{16}$, then
The domain becomes a long column, and its boundary volume becomes
For the same amount of work, smaller local memory requires more communication. These numbers describe ideal boundary costs, rather than exact measured traffic on a particular machine.
Why minimum communication is not always minimum runtime
Each message incurs startup latency as well as transmission time. A common model is
where $r$ is the number of communication rounds, $\alpha$ the latency per round, $w$ the element width in bytes, and $\beta$ the effective bandwidth.
Combining many small messages into a larger one reduces $r$, but requires more buffer space. Receiving $h$ input groups at once changes the local constraint to
Communication volume, message count, and memory capacity must therefore be balanced together. COSMA also overlaps communication with computation, preallocates buffers, optimizes broadcasts and reductions, and uses local BLAS.
Processor count affects grid shape too. The paper gives an example in which 65 processors are difficult to arrange into a suitable regular grid, while 64 can form a $4\times4\times4$ grid and may communicate less. Using every available processor and achieving the shortest runtime are different objectives.
What the experiments establish
The paper compares COSMA with ScaLAPACK, CARMA, and CTF on the CPU partition of Piz Daint, covering square and elongated matrices at several scales. Its abstract reports speedups up to 12.8 times, averaging 2.2 times over those baselines, and performance reaching 88% of the platform’s peak compute rate.
These are results for the 2019 paper’s particular platform, baselines, and experiments. They are not guarantees for every modern GPU or matrix shape. The valuable lesson is the complete chain from lower bound, through decomposition, to implementation.
What this suggests for inference performance research
To estimate a performance limit without relying on an existing kernel, COSMA suggests a method: describe the permitted scalar computation and reuse, derive minimum communication under memory constraints, then connect it to device bandwidth.
The argument still assumes classical matrix multiplication. Independence from an existing implementation does not eliminate algorithmic assumptions.
If several Transformer matrix multiplications are fused, one operator’s output may remain in fast memory for the next. We cannot mechanically add compulsory writeback costs for isolated operators. The lower bound must describe the permitted fusion scope and residency state.
COSMA’s central idea becomes a concrete question: how should this memory be organized so that every transferred number participates in as much useful computation as possible? Answering that before selecting a fixed processor grid gives the decomposition a reason to exist.
Where to read in the original
- Sections 2–3: machine model, classical matrix multiplication, and the COSMA overview.
- Section 4, Lemmas 2–4: partition lower bounds that account for reuse.
- Section 5, especially 5.2.7: sequential schedules, near-optimal constructions, and executability.
- Section 6.3, Equations 31–32: geometric optimization of local domains.
- Sections 7 and 9: implementation optimizations, platform conditions, and experimental results.
The outer-product derivation and numerical examples are original explanations. The local geometric model builds intuition; strict guarantees require the paper’s machine and computation assumptions.