Skip to content

← Performance step 3 of 7

Easy Primitives

Build a reusable mask outside the loop

Jump to the task

If a value stays the same throughout a loop, you may be able to compute it once and reuse it. Here the sequence length is fixed, so the mask does not need rebuilding for each score matrix:

import torch

seq = 3
mask = torch.triu(torch.ones(seq, seq), diagonal=1).bool()
for step in range(2):
    scores = torch.full((seq, seq), float(step + 1))
    out = scores.masked_fill(mask, 0.0)
print(out.tolist())  # [[2.0, 0.0, 0.0], [2.0, 2.0, 0.0], [2.0, 2.0, 2.0]]

True entries are strictly above the diagonal. For masked_fill, True means replace that position, so those entries become zero. This is a matrix-masking exercise, not a complete attention implementation: setting future logits to zero would not exclude them from a later softmax.

Check what can change

The mask is reusable here because its shape, device, dtype and contents stay suitable for every iteration. In a different workload, the sequence length or device may change. Include those dependencies when deciding whether a cached tensor can be reused.

For a module with a fixed maximum sequence length, a registered buffer can hold such a mask. A nonpersistent buffer moves with the module but is omitted from its state_dict, which is useful when the mask can be recreated from the configuration. Slice it for shorter inputs, and handle longer inputs explicitly. A basic slice shares storage; it is not literally free of all overhead.

Not every repeated expression belongs in a buffer. A Python scaling constant can remain a number. A value derived from trainable parameters may become stale after an update, and retaining its graph across iterations can be wrong. Also, calling to(device) on a tensor already on that device need not transfer or copy it. Check the actual operation before treating it as repeated work.

Your task

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

For positive steps and seq, keep the starter's loop. At step s, beginning with zero, create a seq-by-seq score matrix filled with float(s + 1), then zero the entries strictly above its diagonal. Build the mask once before the loop and reuse it. Return only the last matrix as a nested list under final; do not skip the earlier score calculations.

The harness adds hoisted, which is True when its storage-output counter is less than 3 * steps + 3. With the supplied operations, the starter counts 5 * steps, while moving mask construction outside the loop gives 2 * steps + 3. The supplied cases use at least three steps.

This counter checks a storage-output budget. It is not a measurement of all allocator calls or runtime, and it cannot by itself prove where a mask was created. The intended exercise is to move the invariant construction, not to replace the simulation with a formula for its final result.