Session 2, Section 4 · a matrix multiply, one beat at a time
Systolic array heartbeat
The weights are preloaded into the processing elements (PEs), and the activations are pushed in from the left one beat at a time, on a diagonal. Partial sums flow from top to bottom and become results at the bottom edge. Change the token count M and count how many PEs work on each beat.
Streaming Y = X · W through the array
PEs working this beat
Multiply-accumulates so far
Overall utilization (run to the end)
Reuse
View
Check
Checks whether the array output equals the product worked by hand. In each cell, the small text at top left is the weight w held in that PE, at top right the activation x passing through now, and in the center the partial sum after that PE.
Check against the text ()
In all three views, checks that the array output equals X·W worked by hand, and that the active PEs and cumulative MACs at the initial state (fourth beat) and the utilization at TPU v1 size match the text and figure of Session 2, Section 4.
Item
Computed
Expected
Scaled to TPU v1 size
Token count M
Beats M+K+N−2
Multiply-accumulates
Utilization
With M = 1 (batch-1 decode), most beats go to filling and draining. In prefill, with thousands of tokens, that cost is buried and utilization climbs to nearly 90%. This counts in beats what Section 4 of the text says: "with small batches or long, thin matrices (decode!), utilization drops sharply."
What this toy simplifies
Only one tile is streamed. Real chips feed the next tile right behind it, overlapping fill and drain time. Even so, with a small M each weight is used only M times and then discarded, so the gain is small.
The time to load weights into the PEs is not counted.
Only one dataflow, weight-stationary, is shown. Others exist, such as output-stationary.
The values are small integers chosen for easy checking. The real TPU v1 does 8-bit integer multiply-accumulate.