Профилирование Attention в PyTorch: от наивной реализации до SDPA и кастомных ядер

Профилирование Attention в PyTorch: от наивной реализации до SDPA и кастомных ядер
Замерьте профиль наивного attention: torch.profiler покажет гигантскую аллокацию для матрицы scores (BxHxSxS) и низкую утилизацию GPU.
Переключитесь на torch.nn.functional.scaled_dot_product_attention (SDPA) — встроенная оптимизация автоматически выбирает memory-efficient или FlashAttention backend.
Для сверхдлинных последовательностей подключите Flash Attention 3 (через xformers или PyTorch nightly) — память растёт почти линейно, а скорость кратна наивной реализации.
Если нужны кастомные маски или sparse attention, напишите ядро на Triton, которое обходит квадратичную матрицу и точно контролирует аллокации.
Читать статью →

Подписаться на канал