Batch the head projections

When several matrix multiplications have matching shapes, express them as a batch instead of looping over them in Python. Each batch item can still use its own matrix.

import torch

x = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]])
w = torch.tensor([[[10.0], [1.0]], [[1.0], [10.0]]])
print(torch.bmm(x, w).tolist())  # [[[12.0]], [[43.0]]]

The first item computes 1*10 + 2*1; the second computes 3*1 + 4*10. This pairs corresponding items. It does not multiply every input by every weight matrix or share the first matrix across the batch.

Replace the Python loop

The starter indexes each head, multiplies its matrices and stacks the results. torch.bmm(x, w) expresses that work in one tensor operation. This can reduce Python dispatch overhead and give the backend more work to schedule together.

A tensor operation is not a promise of one GPU kernel. The chosen kernels and performance depend on the device, tensor shapes, layout, dtype and library implementation. Batching often helps, but neither a specific speedup nor improved hardware utilization follows from the dispatch count alone.

Layout transformations have costs too. permute creates a view for ordinary strided tensors, while reshape may copy when the existing strides cannot represent the requested shape. Measure the complete path, including any copies required to prepare the batch.

Shape requirements

bmm takes two rank-3 tensors with equal leading batch sizes. Their inner matrix dimensions must match:

(heads, n, d) @ (heads, d, m) -> (heads, n, m)

If different heads need different matrix sizes, they cannot be represented by these two rectangular batches directly. Padding, grouping compatible shapes or using another implementation are options with different costs. Padding is not automatically the best choice.

Multiply rows by columns

For A with shape (M, K) and B with shape (K, N), the matrix product has shape (M, N). Each output entry is a dot product of one row of A and one column of B:

$$C_{ij}=\sum_{k=1}^{K}A_{ik}B_{kj}.$$

For example, multiplying row [1, 2] by column [3, 4] gives 1 * 3 + 2 * 4 = 11. Elementwise multiplication would instead give [3, 8]; it does not include the sum.

If examples are rows of X, with shape (B, D), and weights have shape (D, H), then XW produces H features for each of the B examples.

Keep batch axes separate from matrix axes

For batched matrix multiplication, the last two axes hold each matrix. Leading axes identify batches of matrices. A general matrix-multiply operation may broadcast compatible batch axes. A stricter batched operation may require them to match exactly. Use the operation specified by the exercise.

Einstein notation, often used with einsum, names the axes explicitly:

• bij,bjk->bik multiplies corresponding matrices in batch b and sums over j.
• bi,bj->bij forms an outer product for each example: output (b, i, j) is the product of input (b, i) and input (b, j).

An axis is summed out when its label is absent from the explicit output. A label that appears in both inputs is not necessarily summed: b stays in both outputs above.

Choose a multiplication order with small intermediate results

In exact arithmetic, (AB)C = A(BC), but the two orders can have very different costs. If the shapes are (1000, 10), (10, 10), and (10, 1), forming BC first needs a (10, 1) intermediate. Forming AB first needs a (1000, 10) intermediate and many more operations.

Floating-point rounding can also differ between orders, so do not expect every bit of the two results to match.