415.tech
AI & tech, from the frontlines of Silicon Valley
Hugging Face profiles PyTorch attention on the A100, from naive kernels to flash

Hugging Face profiles PyTorch attention on the A100, from naive kernels to flash

The third post in Hugging Face's 'Profiling in PyTorch' series reads torch.profiler traces to compare attention implementations on an A100 — a naive hand-built version, an in-place variant, and the four SDPA backends. The payoff is buildable: a developer can spot hidden kernels in the trace — an out-of-place masked_fill quietly adds a memory copy per forward pass — and pin the fastest path, with flash at 146.8µs against efficient's 277.9µs and cuDNN's 186.3µs.

Source: huggingface.co ↗

Post on XEmail

With a one line change we shaved a whole kernel off each forward pass.

Hugging Face

Why this matters

  • → Hidden memory copies in masked_fill silently degrade attention performance per layer.
  • → SDPA backend selection dramatically varies kernel counts and latency across A100.
  • → Profiler traces reveal optimization wins invisible to raw code inspection.
Profiling attention kernels