One kernel for every head

A loop over attention heads, each with its own projection:

outs = [x[h] @ w[h] for h in range(heads)]
out = torch.stack(outs)

and the batched form:

out = torch.bmm(x, w)

Same arithmetic. One kernel instead of heads kernels plus a stack.

Why it is worth more than the launch count

A matmul kernel is tuned for large matrices: it tiles the work, keeps tiles in registers and shared memory, and needs enough of it to hide memory latency. A small matmul cannot fill the machine, so most of the hardware idles and the kernel is bounded by launch overhead rather than by arithmetic.

Batching gives the kernel enough work to schedule properly. The speedup is usually larger than the reduction in launches, which is not what you would predict from counting.

The general shape

many small identical operations   ->   one batched operation

It is why fused QKV projections exist (one matmul instead of three), why grouped convolutions are one kernel, and why nn.MultiheadAttention keeps all the heads in one tensor rather than in a list of modules.

It is also the reason the shapes track spent so long on merging and splitting axes: reshape and permute are what let you get the data into the layout a batched kernel wants, for free.

When not to

When the operations differ in shape. bmm needs a rectangular batch, so genuinely ragged work has to be padded, and if the padding is most of the tensor you have traded compute for launches badly. Measure.

Your task

def project_heads(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor

x is (heads, n, d) and w is (heads, d, m). Return (heads, n, m), applying each head’s own matrix, in at most two dispatched operations.