Rebuilt every step

for batch in loader:
    mask = torch.triu(torch.ones(seq, seq), diagonal=1).bool()
    out = attend(batch, mask)

The mask depends on seq and nothing else. It is rebuilt from scratch every iteration, for the entire run.

Why this is not obvious in PyTorch specifically

In ordinary code a hoistable expression is usually cheap and the compiler hoists it anyway. Here it is a tensor allocation plus a couple of kernels, Python has no optimiser that will move it, and the line looks like setup rather than like work.

It also often lives inside forward, where “once per call” reads as “once”, and a training run calls forward a hundred thousand times.

The usual suspects

causal masks           depend on the sequence length
position encodings     depend on the length and the dimension
arange indices         depend on a shape
a scaling constant     d ** 0.5, recomputed per attention call
.to(device) on a constant   a transfer per step

All of them belong in __init__ as a buffer, which is what the modules track was for: registered, so it moves with the model, and built once.

self.register_buffer("mask", torch.triu(...), persistent=False)

persistent=False because it is derivable from the config, so saving it would waste space and bake a sequence length into the checkpoint.

The caveat

Caching by shape needs the shape to be stable, or the cache needs a key. A mask cached for seq=512 and reused for a shorter batch has to be sliced, not reused whole:

mask = self.mask[:seq, :seq]      # a view, free

which is the views track paying for itself again.

Your task

def run(steps: int, seq: int) -> dict

Apply a causal mask to steps successive score tensors, building the mask once. Return the final result and the total number of tensors allocated.