Batched matmul and einsum

@ multiplies the last two axes as matrices and treats the rest as a batch. The weight has no batch axes, so it broadcasts: every sequence goes through the same $W$. This is nn.Linear on a whole batch.

Queries times transposed keys, for every sequence and every head at once: $(T \times d_h)(d_h \times T) = T \times T$ per $(b, h)$ pair. The leading axes $(B, h)$ must match (or broadcast).

Give each axis a letter. A letter that appears in the inputs but not after -> is summed over; the rest are kept, in the order you write them. That one rule covers dot products, matrix products, transposes and batches.

In the attention strings, q and k both index tokens: queries and keys. Writing bhkd for the keys means no explicit transpose is needed.

$A$ holds the attention weights (one row per query, one column per key). Write the einsum for the output $AV$ and its shape.

k (the key position) appears in both inputs and not in the output, so it’s summed over: each query’s output is a weighted sum of value rows. Same as A @ V.

The 2mkn rule generalizes: every output entry costs 2 FLOPs per summed index. For one sequence of 4,096 tokens with 32 heads of 128: $2 \cdot 32 \cdot 4096^2 \cdot 128 \approx 1.4 \times 10^{11}$ FLOPs for the scores of one layer. Note the $T^2$: attention cost grows with the square of the context.

• @ multiplies the last two axes; leading axes broadcast.
• einsum: letters missing from the output are summed.
• Scores "bhqd,bhkd->bhqk", output "bhqk,bhkd->bhqd".
• FLOPs = 2 × kept sizes × summed size; attention scores grow like $T^2$.