Two hundred kernels, or one

Everything else in this track is worth a few percent. This one is worth two orders of magnitude, and it is nearly always available.

out = torch.empty_like(x)
for i in range(x.numel()):
    out[i] = x[i] * 2

versus

out = x * 2

On 64 elements: 257 dispatched operations against 1. Measured, not estimated.

Where the time goes

Not the arithmetic. Multiplying 64 floats is free on any hardware built this century. The cost is the per-operation overhead, paid once per element by the loop:

a Python function call
argument parsing and type dispatch
a kernel launch
on a GPU, a launch latency of a few microseconds

A few microseconds times 64 is nothing; times a million is a minute. And the vectorised version pays it once regardless of size, which is why the gap widens with the data rather than staying constant.

Recognising it

Any for loop whose body operates on one element is this. The tells:

for i in range(len(x)):        # indexing a tensor by a Python integer
for row in matrix:             # iterating a tensor at all
[f(v) for v in x.tolist()]     # a list comprehension over elements

Loops over a small fixed number of layers or heads are fine. The problem is loops proportional to the data.

When it genuinely resists

Some computations really are sequential: a recurrence where step n needs step n - 1‘s output. Even then, look for the scan formulation first, as the reductions track showed with cumsum, and look for a way to batch the independent dimension.

Your task

def normalise_rows(x: torch.Tensor) -> torch.Tensor

Divide each row by its own maximum absolute value, leaving a row of all zeros alone. Use at most six dispatched operations regardless of the input’s size.

The starter loops over rows and then over elements.