В пайплайне перевода видео в Яндекс Браузере скорость работы в синтезе речи долго оставалась узким местом: одно видео — это сотни фраз, каждая генерируется отдельно, и любая сэкономленная миллисекунда становится заметной в масштабах сервиса. 

Внутри TTS работает каскад из трёх компонентов: языковая модель предсказывает аудиотокены по тексту, диффузионный декодер восстанавливает мел-спектрограмму из латентов, а вокодер превращает её в звуковую волну. Долгое время самой тяжёлой была языковая модель, но после её оптимизации на первый план вышел декодер латентов — его forward pass запускается на каждом шаге семплинга диффузии, а шагов — десятки. Именно его мы и взялись ускорять.

Меня зовут Цырен-Доржо Цыбиков, я ML-инженер команды TTS в Яндексе и ментор проекта Даниила Маслова, студента бакалавриата кафедры АД ФПМИ МФТИ и ШАД. В рамках проектного курса мы запустили отдельную исследовательскую ветку по оптимизации диффузионного декодера латентов в TTS-пайплайне перевода видео. 

Год назад мы уже рассказывали на Хабре, как заменили старый многоголосый синтез на zero-shot TTS на базе сильно переработанной архитектуры Tortoise и научились сохранять тембр и интонацию спикера. Тогда речь шла о модели целиком: языковой части, подборе аудиопромптов, переносе тембра между языками. В этот раз спускаемся на уровень глубже — к одному конкретному компоненту.

Для проектного курса такая задача хорошо подходит. С одной стороны, она реальная: декодер действительно занимает заметную долю времени в TTS-пайплайне. С другой — она не блокирует ближайший релиз. У студента есть пространство для нормального R&D: можно проверять гипотезы, гонять бенчмарки, смотреть в профайлер и не работать в режиме жёсткого релизного дедлайна.

Что именно мы оптимизируем

Диффузионный декодер в Tortoise — это стек слоёв DiffusionLayer. Каждый такой слой устроен как модифицированный трансформерный блок и состоит из двух основных частей:

• residual-ветка: пара Conv1D, SiLU и Adaptive GroupNorm, куда через scale/shift добавляется временной эмбеддинг шага диффузии;

• attention-блок: GroupNorm, QKV-проекция, Multi-Head Self-Attention и out-проекция.

Внутри attention-блока используется авторская реализация QKVAttentionLegacy. Важная особенность Tortoise в том, что позиции не зашиваются в Query и Key через RoPE и не сводятся к ALiBi. Вместо этого используется обучаемый аддитивный RelativePositionBias, который прибавляется к логитам внимания перед softmax.

Это важная деталь для оптимизации. Такой bias нужно отдельно вычислять и приводить к нужной форме для добавления к логитам attention. Кроме того, он мешает большинству быстрых реализаций attention, которые рассчитаны на более стандартные варианты позиционного кодирования.

Первым делом мы запустили torch.profiler на инференсе декодера в конфигурации, близкой к продовой: с реальным размером батча, длиной последовательностей и числом шагов сэмплинга.

Из профиля сразу стало видно две вещи:

Первая — основное время уходит в AttentionBlock, а внутри него — в QKVAttentionLegacy. Механика умножений там собрана через einsum с полной материализацией матрицы внимания. Это значит, что мы тратим время не только на сами умножения, но и на хранение промежуточной матрицы QKᵀ, а PyTorch не может заменить эту цепочку операций одним оптимизированным кернелом.

Вторая — на каждом forward pass заново строится RelativePositionBias, хотя он не зависит ни от шага диффузии, ни от конкретного сэмпла.

После профилирования мы составили пошаговый план работы и разделили его на четыре этапа:

1. оптимизации внутри attention, не требующие переобучения большой модели и совместимые с текущими весами;

2. переход на RoPE как способ разблокировать FlashAttention;

3. поиск более современных архитектур, которые смогут дать лучшее соотношение качества и времени инференса.

4. дистилляция classifier-free guidance как способ убрать второй forward pass на каждом шаге сэмплинга.

Этап 1. Бесплатные оптимизации: SDPA и кэширование позиционного bias

Главная цель первого этапа — получить ускорение без дообучения модели. Если веса остаются совместимыми с уже обученным чекпоинтом, новую реализацию можно проверять и выкатывать гораздо проще: без затрат на тренировочные кластеры и без риска просадить качество синтеза из-за изменения обучения.

Шаг 1.1. От рукописного attention к SDPA

QKVAttentionLegacy внутри Tortoise — это рукописная реализация Multi-Head Self-Attention, которая воспроизводит логику модуля, но имеет две основные проблемы с точки зрения производительности.

Во-первых, матрица внимания материализуется в памяти целиком. На длинных последовательностях это становится одним из основных потребителей GPU-памяти.

Во-вторых, PyTorch видит цепочку операций как набор независимых мелких кернелов и не может объединить их в одну более эффективную операцию.

Стандартное решение — переехать на torch.nn.functional.scaled_dot_product_attention, далее SDPA. Это API, под которым PyTorch сам выбирает один из доступных backend:

• math — референсный вариант, где всё разворачивается явно;

• memory-efficient — xFormers-style реализация без полной материализации QKᵀ;

• cuDNN — backend для новых GPU и свежих версий библиотек;

• FlashAttention.

Под капотом выбор backend зависит от размеров и типов тензоров, наличия маски и версий библиотек.

В нашем случае главным ограничением был RelativePositionBias. SDPA принимает attn_mask — произвольный тензор, который добавляется к логитам перед softmax. memory-efficient backend умеет работать с такой маской, а FlashAttention с произвольной аддитивной маской из коробки не совместим. Поэтому на первом шаге мы перешли на SDPA с memory-efficient backend, сохранив RelativePositionBias.

Шаг 1.2. Кэшируем то, что не меняется

Вторая проблема — постоянный пересчёт RelativePositionBias на каждом forward pass. Как мы уже выяснили, на фиксированной длине последовательности эта аддитивная маска одинакова на всех шагах диффузии и для всех элементов батча. Более того, при меньшей длине последовательности можно брать соответствующий срез уже посчитанной матрицы.

Решение простое: достаточно один раз посчитать маску для максимально возможной длины последовательности и дальше переиспользовать её. Само по себе кэширование не выглядит впечатляющей оптимизацией, но именно оно убирает из профиля хвост мелких einsum-операций.

Результаты первого этапа

Замеры мы прогоняли на одной H200 на трёх уровнях:

1. отдельный QKVAttention;

2. AttentionBlock;

3. DiffusionLayer целиком.

На каждом уровне сравнивали legacy-реализацию, переезд на современный API — Modern, варианты с разными backend SDPA (math, efficient, flash), переход на RoPE — об этом следующий раздел — и отдельно версии с кэшированным RelativePositionBias.

В числах на уровне QKVAttention переход с legacy-реализации на Modern + Cache + Efficient SDPA сокращает время примерно с 0,77 мс до 0,28–0,30 мс, то есть даёт ускорение около 2,5×.

На уровне DiffusionLayer целиком ускорение скромнее — это ожидаемо, потому что attention уже не единственный источник вычислительной стоимости. Дополнительный бонус — почти двукратное снижение потребления GPU-памяти под активации: memory-efficient backend не материализует QKᵀ.

Важно, что всё это получилось без изменения весов. Старый чекпоинт продолжает работать, мы меняем только способ вычисления attention внутри слоя.

Этап 2. RoPE как способ разблокировать FlashAttention

SDPA с memory-efficient backend — это хорошо, но FlashAttention в этих экспериментах оставался недоступен из-за RelativePositionBias. Возник логичный вопрос: что если переехать на позиционные эмбеддинги, которые не требуют аддитива поверх логитов? Тогда attention-маски можно либо вообще не использовать, либо использовать только обычные causal/padding-маски, с которыми FlashAttention работает гораздо лучше.

Мы решили попробовать RoPE. В этой схеме позиции кодируются непосредственно в Q и K через поворот в комплексной плоскости, и в attention идёт «чистая» матрица QKᵀ, без дополнительных слагаемых. Это означает, что FlashAttention и cuDNN-backend становятся доступны.

Здесь начинается принципиальная разница с предыдущим этапом: RoPE — это не drop-in-замена. Модель училась с обучаемым RelativePositionBias и видит позиции по-другому. Переезд требует дообучения. То есть цена оптимизации повышается: нужны GPU-часы и метрики, нужно убедиться, что качество синтеза не просядет.

В бенчмарках RoPE-варианты отмечены как rope/efficient, rope/cudnn и rope/flash. На уровне QKVAttention FlashAttention с RoPE укладывается примерно в 0,57 мс, а memory-efficient backend с RoPE — около 0,55 мс.

Лучший общий результат всё равно остаётся за modern_cache/efficient — около 0,28 мс. Это важный негативный результат.

Он говорит о том, что на конкретных размерах тензоров в нашем декодере memory-efficient backend PyTorch отрабатывает эффективнее, а сам по себе FlashAttention на этих масштабах не даёт значительного преимущества. Таким образом, дообучать текущую модель только ради FlashAttention не имеет практического смысла.

Для проектного курса это был один из самых полезных моментов. Гипотеза выглядела красиво: заменить RelativePositionBias на RoPE, открыть FlashAttention и получить ускорение. Но бенчмарки показали, что первый этап уже дал очень сильный baseline.

В реальном R&D важно не влюбляться в гипотезу только потому, что она модная или хорошо выглядит на бумаге. Если эксперимент поставлен честно, отрицательный результат тоже полезен: он снижает неопределённость и помогает команде не тратить GPU-часы на направление, которое не даёт практического выигрыша.

В итоге RoPE оказался полезен не как самостоятельная оптимизация текущей архитектуры, а как мост к следующему шагу — более серьёзной замене декодера.

Этап 3. В поисках новых архитектур

К этому моменту у нас было два вывода. Первый: из текущей архитектуры можно получить заметное ускорение без переобучения. Второй: переходить на RoPE только ради FlashAttention не стоит — выигрыш не окупает стоимость дообучения старой модели.

Логичный следующий шаг — задаться вопросом: а не поменять ли архитектуру декодера принципиально?

Главный претендент на замену — DiT, Diffusion Transformer. Он впервые появился в задачах генерации изображений как замена U-Net-бэкбона: вместо свёрточных блоков используются трансформерные, а conditioning задаётся через Adaptive LayerNorm — параметры нормализации зависят от эмбеддинга шага диффузии. Сначала на DiT перешли ведущие модели генерации изображений, а затем похожие идеи появились и в современных TTS-системах, например F5-TTS и CosyVoice3.

Есть и ещё одно практическое следствие такого перехода: наша версия DiT не предполагает произвольной аддитивной маски поверх attention-логитов. Позиции кодируются через RoPE, что означает полную совместимость с FlashAttention — то, чего мы так и не смогли добиться с RelativePositionBias.

Блок DiT почти в 2,5 раза тяжелее по числу параметров, чем исходный DiffusionLayer. Но параметры не всегда линейно переводятся в реальные FLOP на инференсе. Основная «лишняя» нагрузка DiT сосредоточена в небольшом MLP, который предсказывает параметры scale и shift для AdaLN из эмбеддинга шага диффузии. Это сравнительно дешёвая операция.

Основная вычислительная нагрузка по-прежнему приходится на attention и FFN, для которых существуют хорошо оптимизированные GPU-кернелы. В итоге вычислительная стоимость одного прохода растёт скромнее, чем можно было бы предположить, глядя только на разницу в числе параметров.

По результатам экспериментов DiT-декодер при том же числе шагов сэмплинга работает немного медленнее оптимизированной Tortoise-реализации — это ожидаемая плата за более крупный блок. При этом DiT показывал более низкий loss и выглядел перспективнее по качеству на предварительных прослушиваниях.

На практике это даёт дополнительное пространство для оптимизации: можно сократить число шагов генерации или использовать модель с меньшим числом слоёв — и тем самым вернуться в конкурентоспособный диапазон по времени инференса уже с более сильной архитектурой.

Этап 4. Дистиллируем classifier-free guidance

Помимо оптимизаций attention, мы отдельно проверили дистилляцию classifier-free guidance. В исходном декодере CFG требовал двух forward pass’ов на каждом шаге сэмплинга: conditional и unconditional. Это улучшало следование условию, но почти удваивало стоимость инференса. Мы обучили student-модель напрямую воспроизводить guided-предсказание teacher-модели с CFG, но за один forward pass.

Для проектного курса это была хорошая исследовательская задача: понятный baseline, понятная метрика успеха и практическая ценность в случае удачи. В итоге именно она дала самый сильный выигрыш. В side-by-side-замерах дистиллированная модель оказалась статистически неотличима от teacher-модели с CFG: пользователи не заметили, что вместо двух вызовов исходной модели мы делаем один вызов student-модели. При этом диффузионный декодер ускорился примерно в 2 раза, и эта оптимизация в итоге дошла до продакшена видеоперевода

Выводы

В этом проекте мы прошли путь от низкоуровневых оптимизаций attention до дистилляции модели и проверки новых архитектур. Вот что мы поняли.

1. Начинать нужно с профилирования: torch.profiler сразу показал две основные проблемы: legacy attention через einsum и пересчёт RelativePositionBias на каждом forward pass.

2. FlashAttention не всегда автоматически быстрее. На наших размерах тензоров RoPE + FlashAttention не обогнал SDPA memory-efficient с кэшированным bias. Это отрицательный, но полезный результат: мы поняли, что тащить RoPE в старую архитектуру только ради FlashAttention невыгодно.

3. Самый большой практический выигрыш дала CFG-дистилляция. Вместо двух forward pass’ов teacher-модели на каждом шаге мы обучили student-модель воспроизводить guided-предсказание за один проход. SBS относительно teacher оказался серым, а декодер ускорился примерно в 2 раза. В отличие от части исследовательских гипотез, эта оптимизация дошла до прода видеоперевода.

4. Проектный курс хорошо подходит для таких задач. Здесь есть реальная боль, baseline, метрики и пространство для R&D без риска заблокировать основной релиз. Команда получает проверенные гипотезы и иногда — как в случае CFG-дистилляции — продовый результат.

Для меня как ментора этот проект оказался хорошим примером задачи, которую удобно выносить в проектный курс ШАД. Это не изолированная учебная лабораторная и не критичный кусок продового релиза, где нельзя ошибаться. У студента есть возможность пройти через настоящий инженерный R&D: от профилирования и простых оптимизаций до проверки архитектурных гипотез. А команда получает ответы на вопросы, до которых в основной продовой ветке часто не доходят руки.

Если у вас в команде есть похожие задачи — с неопределённостью, но с понятной инженерной рамкой, baseline и метриками, — из них часто получаются сильные студенческие проекты, а команда ШАДа как раз в поисках новых менторов и проектов.