Профилирование в PyTorch, часть 3: внимание, это всё, что нужно профилировать

Это третий пост серии Hugging Face «Profiling in PyTorch», которая учит читать профайлер-трейсы и таблицы PyTorch и использовать их для оптимизации кода. В первых двух частях разбирали простые операции сложения/умножения и линейные слои; в третьей части перешли к механизму attention (внимание), ключевому и квадратично-затратному по времени блоку архитектуры Transformer. Эксперименты ставили на GPU NVIDIA A100-SXM4-80GB.

Сначала авторы написали наивный модуль внимания вручную: matmul(q, k.T) → масштабирование → маска → softmax → matmul с values, пять операций, которые ожидаемо появились в трейсе. Но при разворачивании GPU-дорожки трейса обнаружился лишний, незапланированный kernel, операция копирования памяти (memory copy). Причина, в PyTorch операция scores.masked_fill() создаёт копию тензора (out-of-place), а не изменяет его на месте. Замена на масштабированный вариант с завершающим подчёркиванием, masked_fill_() (in-place-версия), убрала лишний kernel полностью: один forward-проход стал короче на целый kernel. Авторы поясняют: out-of-place, поведение PyTorch по умолчанию, потому что autograd должен помнить значения тензора с forward-прохода для вычисления градиентов на backward; но в их случае код выполняется под torch.no_grad (без обратного прохода), поэтому in-place-операция безопасна и вдобавок экономит память на больших тензорах.

Далее авторы протестировали встроенную функцию F.scaled_dot_product_attention (SDPA), единую функцию PyTorch, которая скрывает всю логику attention и сама выбирает один из нескольких бэкендов (math, flash, efficient, cudnn) в зависимости от входных данных (тип данных, размерность головы, маска, железо). Бэкенды можно явно зафиксировать через torch.nn.attention.sdpa_kernel.

Первый протестированный бэкенд, math (эталонная реализация). Ожидалось, что единая функция окажется быстрее и проще наивного кода, но результат оказался обратным: math-бэкенд оказался в 3,7 раза медленнее и запускает 20 GPU-kernel'ов на forward-проход вместо 5 у наивной версии. Причины три. Во-первых, math-бэкенд не использует Tensor Cores (специализированные блоки A100 для матричного умножения), по именам kernel'ов видно, что наивная версия использует быстрый bfloat16 Tensor Core matmul (сигнатура s16816), а math-бэкенд переводит тензоры в FP32 и считает через обычный sgemm на CUDA-ядрах, удваивая объём данных ради точности. Во-вторых, при флаге is_causal=True math-бэкенд заново строит причинную маску на каждом вызове (через aten::ones, aten::tril, aten::where и другие операции), тогда как в наивной версии маска строилась один раз и переиспользовалась. В-третьих, math-бэкенд вызывает aten::_safe_softmax вместо обычного softmax, защищённую версию, которая избегает NaN в строках, полностью замаскированных значением -inf (обычный softmax в такой ситуации выдаёт 0/0 = NaN); наивная реализация авторов эту защиту не делала и потенциально могла тихо возвращать NaN. Вывод авторов: math-бэкенд, это осторожная, безопасная по типам и по NaN эталонная реализация, чья задача не в скорости, а в корректности; именно эта тщательность и делает его медленным.

В противоположность этому efficient-бэкенд запускает всего один слитный (fused) kernel, fmha_cutlassF_bf16_aligned_64x64_rf_sm80, вместо 20 у math-бэкенда. Название kernel'а расшифровывается как fused multi-head attention (все примитивные операции attention объединены в одну), построенный на библиотеке CUTLASS от NVIDIA. На этом фрагменте исходный материал источника обрывается, часть про cudnn- и flash-бэкенды в предзагруженном тексте отсутствует.

Ключевые факты

Наивная ручная реализация attention (matmul → масштаб → маска → softmax → matmul) в трейсе PyTorch неожиданно показала лишний kernel копирования памяти из-за того, что masked_fill() создаёт копию тензора, а не изменяет его на месте.
Замена на in-place-операцию masked_fill_() убрала лишний kernel полностью, экономия и по времени, и по памяти (актуально под torch.no_grad, где autograd не нужен).
Встроенная функция F.scaled_dot_product_attention с бэкендом math оказалась в 3,7 раза медленнее наивной реализации и запускает 20 GPU-kernel'ов вместо 5.
Причины медленности math-бэкенда: он не использует Tensor Cores A100 (считает в FP32 на обычных CUDA-ядрах вместо bf16 на Tensor Core), заново строит причинную маску на каждом вызове и использует защищённый от NaN aten::_safe_softmax.
Бэкенд efficient сворачивает все операции attention в один слитный kernel fmha_cutlassF_bf16_aligned_64x64_rf_sm80 на базе библиотеки NVIDIA CUTLASS.

Почему это важно

Материал показывает на конкретном примере, что «удобный» и вроде бы оптимизированный API-вызов (F.scaled_dot_product_attention с автоматическим или дефолтным выбором бэкенда) может оказаться в разы медленнее самописного наивного кода, если PyTorch выбирает безопасный, но неэффективный бэкенд math. Причём проблема не видна без профайлера: по одному коду не скажешь, что вызов создаёт лишнюю копию тензора в памяти или заново строит маску на каждом forward-проходе.

Кому это важно

Инженерам и исследователям, которые пишут или оптимизируют модели на PyTorch, особенно трансформеры с механизмом attention, то есть практически все, кто работает с LLM и подобными архитектурами на GPU. Полезно и тем, кто хочет освоить чтение профайлер-трейсов PyTorch как навык, материал построен как обучающий разбор, а не просто отчёт о результатах.

Как это применить

Практические выводы из материала: заменять out-of-place операции (masked_fill) на in-place-варианты (masked_fill_) там, где не нужен backward-проход (например, под torch.no_grad), это убирает лишние копии в памяти. При использовании F.scaled_dot_product_attention стоит явно проверять и, если нужно, фиксировать бэкенд через sdpa_kernel, а не полагаться на автоматический выбор PyTorch, бэкенд math задуман как эталонный и безопасный, а не быстрый. Скрипты для всех экспериментов из поста выложены в открытом виде и их можно повторить самостоятельно, в том числе на инфраструктуре Hugging Face (Spaces Dev Mode, Hugging Face Jobs).

Можно ли доверять

Это методический разбор с открытым кодом (все скрипты названы и доступны), конкретными числами (3,7x, 20 против 5 kernel'ов) и снимками профайлер-трейсов на реальном железе (NVIDIA A100). Материал, часть образовательной серии Hugging Face, а не рекламный анонс продукта, что повышает доверие к выводам. Ограничение: предзагруженный текст источника обрывается на середине раздела про efficient-бэкенд, поэтому часть материала (о cudnn- и flash-бэкендах) в этом пересказе не отражена.

Риски и подводные камни

Главный риск, который иллюстрирует сам материал, доверять «удобным» высокоуровневым вызовам без профилирования: F.scaled_dot_product_attention без явного выбора бэкенда может неожиданно использовать медленный path (FP32 на CUDA-ядрах вместо bf16 на Tensor Cores) и потерять кратный запас производительности. Отдельный риск для тех, кто копирует наивные реализации attention без safe-softmax, тихое появление NaN в полностью замаскированных строках, что может испортить обучение или инференс без явной ошибки.

«Его задача, не быть быстрым, а всегда работать.»

— блог Hugging Face о профилировании PyTorch, про бэкенд math в scaled dot product attention