Оптимизировать нейросети вслепую — значит гадать. Чтобы точно понять, что тормозит и как работает код на GPU, нужны инструменты профилирования. Hugging Face в третьей части своей серии «Профилирование в PyTorch» показывает, как именно механизм внимания (attention) проявляет себя в трассировках, и как разные реализации влияют на производительность.
Первая часть серии учила основам torch.profiler на базовых операциях, а вторая — профилированию многослойных перцептронов (MLP) и слиянию операций (fused kernels). Теперь фокус на механизме внимания (attention), известном своей квадратичной сложностью. Цель — показать, как разные подходы к его оптимизации выглядят в профайлере.
На примере наивной реализации attention видно, что он состоит из нескольких примитивных операций: умножения матриц (q, k.T), масштабирования, применения маски, softmax для нормализации и финального умножения (attn, v). Профайлер чётко показывает эти дискретные операции, включая два умножения матриц (matmul), масштабирование, маскирование и softmax.
Понимание того, как каждая операция отображается на трассировке CPU и GPU, позволяет инженерам выявлять узкие места. Это помогает сравнивать производительность различных оптимизаций, таких как inplace-операции, Scaled Dot Product Attention (SDPA) и оптимизированные CUDA-кернелы. Для экспериментов использовалась NVIDIA A100-SXM4-80GB GPU, а скрипты доступны через Dev Mode with Spaces или Hugging Face Jobs pipeline.
Серия «Профилирование в PyTorch» не даёт готовых рецептов, но формирует критически важный навык: умение самостоятельно интерпретировать трассировки и принимать обоснованные инженерные решения. Это позволяет инженерам точно видеть, как их изменения влияют на работу GPU, а не просто догадываться об эффекте.