Profiling in PyTorch (Part 3): Attention is all you profile

A GPU can finish a calculation faster while keeping fewer execution groups active at once. This attention-profiling tutorial demonstrates that distinction on an NVIDIA A100: of the PyTorch attention backends it tested, its FlashAttention implementation was the fastest on that test shape despite a profiler-reported occupancy of about 13%. Attention compares positions in an input sequence. A straightforward implementation repeatedly writes and reads a large matrix of those comparisons to the GPU's main memory; FlashAttention processes smaller tiles and keeps intermediate work on the chip. It spends more registers and shared memory per block to avoid that memory traffic, which leaves fewer blocks resident at once. The example explains why occupancy alone cannot tell a developer whether an implementation is efficient. Profiling walkthrough