В этой статье демонстрируется профилирование механизмов внимания в PyTorch с использованием трассировок профилировщика для выявления узких мест производительности и возможностей оптимизации. Сравниваются наивные реализации внимания со встроенной реализацией Scaled Dot Product Attention (SDPA) от PyTorch для иллюстрации поведения ядер.
- Наивное причинное внимание включает несколько ядер, включая matmul, масштабирование, маскирование, softmax и неожиданное копирование памяти, вызванное операциями вне места.
- Замена masked_fill на in-place masked_fill_ устраняет ненужное ядро копирования памяти из трассировки GPU.
- Функция SDPA PyTorch автоматически диспетчеризирует к самому быстрому доступному бэкенду (например, Flash Attention или cuDNN) в зависимости от типов входных данных и оборудования.
- Профилирование отдельных бэкендов SDPA выявляет различные сигнатуры ядер, позволяя разработчикам проверять, какая реализация активна.
Понимание этих трассировок профилировщика помогает пользователям оптимизировать трансформеры, устраняя избыточные операции с памятью и используя эффективные бэкенды внимания PyTorch.