Reading attention traces before chasing another faster kernel

The development

The third installment of Hugging Face’s PyTorch profiling series examines attention through profiler tables and execution traces. It begins with an explicit implementation and compares increasingly optimized paths, using the trace to connect mathematical operations to the work actually launched on a GPU.

A straightforward attention implementation creates scores from queries and keys, scales and masks them, applies softmax and combines the result with values. Seeing those stages separately makes it easier to understand why an optimized implementation can have a very different execution pattern while computing the same general operation.

The published examples used a particular NVIDIA A100 configuration. Their main value for another machine is the investigative method, not an assumption that the same timings or preferred backend will transfer unchanged to a different GPU, tensor shape or software stack.

Begin with a hypothesis you can disprove

A useful profiling session starts with a specific question. Perhaps the model spends too long in attention, perhaps temporary tensors consume too much memory, or perhaps the GPU repeatedly waits between short operations. These are different problems. Collecting a large trace without deciding what to investigate often produces more information than insight.

Choose an input shape that reflects the application. Sequence length, batch size, number of heads and numeric precision all influence the work. Record them with the trace so that a later comparison does not accidentally compare different workloads. A small reproducible case is easier to interpret than an entire application captured without boundaries.

Separate initialization from steady execution

The first call may include setup that does not recur on every request. Warm the workload before collecting steady-state timing, while recording cold-start behavior separately if it matters to deployment. Mixing the two can make a fast implementation look slow or hide startup costs that users will encounter.

GPU work is asynchronous relative to the host. A timing method must account for completion rather than merely measuring how quickly the CPU submitted operations. Profiler traces help reveal that relationship, but their own overhead is another reason to confirm conclusions with a lighter measurement outside the trace.

Read both the gaps and the busy regions

A long operation is an obvious target, yet the empty spaces between operations can be equally important. Repeated launches, host-side preparation or synchronization may limit throughput even when individual kernels look efficient. Ask whether the device is performing useful work continuously and which dependency forces the next pause.

Memory behavior deserves a separate view. Intermediate tensors can increase peak allocation without dominating the timing table. If the real objective is fitting a longer context into memory, the fastest isolated kernel is not necessarily the most useful improvement. Define success in terms of the constraint the application actually faces.

Keep correctness beside performance

For every implementation change, compare outputs on representative inputs using a tolerance appropriate to the precision. Include masking and edge cases rather than testing only a convenient square tensor. A speedup caused by omitting required work is not an optimization.

Finish with a compact experiment record: the hypothesis, the controlled change, the observed effect and any tradeoff. That discipline makes the next profiling session cumulative. Instead of repeatedly trying fashionable kernels, the team builds a map of where its own workload spends time and which changes reliably improve it.