CosyVoice 3 ускоряем инференс

Если взять инференс из коробки, то он получается не особо быстрый и не влезает в реалтайм:
https://github.com/FunAudioLLM/CosyVoice

Меня интересует именно стриминг режим, поэтому все ниже происходит с параметром stream=True.
Оффлайн инференс происходит быстрее, так как отсутствует оферхед на переключение между моделями каждый чанк.
Все тесты и замеры проводил на RTX 3090.

Ориентироваться будем в основном на 3 метрики:
TTFB (Time to first byte) - задержка до первого чанка готовой речи, как быстро можем начать воспроизводить речь.
RTF (Realtime factor) - отношение времени генерации речи к длине ее воспроизведения. Если 8 секунд речи получили за 4 секунды, значит RTF = 0.5.
TPS (Tokens per second) - касается только LLM части модели, сколько токенов в секунду успевает генерировать LLM.

Что имеем

Модель внутри состоит из 3-х модулей:
LLM (Qwen2 0.5B)
FLOW (DiT)
HIFT (vocoder)

Добавляем логи на каждый шаг, замеряем время/скорость и начальные значения:
LLM работает где-то 6 секунд
FLOW отрабатывает 3 секунды
HIFT 1 секунда

TTFB: 2.2 sec
RTF: 1.5
TPS: 16

Это то что имеем из коробки, без каких-либо оптимизаций.

Вайб

Конечно же сразу отправляем Cursor (Opus-4.5) искать в чем затык, сами лезем глазами смотреть как устроен инференс, вдруг что-то заметим.
Из коробки можно включить fp16 и tensorRT, включаем.

Cursor возвращается с проблемами и их решением:
forward_one_step - вытаскиваем весь hidden state (долго), нужно только последнее состояние.
nucleus_sampling - бегаем в цикле обычным питоном (долго).

Запускаем, видно +20-30% ускорения, но все еще RTF > 1 в двух из трех кейсах.

Больше всего съедает времени LLM, на нее и будем смотреть.
Мы включили fp16, это должно было дать в идеале х2, включился ли он? Да, но только для FLOW.
Исправляем, включаем fp16 для LLM части, получаем примерно +10-20%, TPS 20-24 - х2 vs х1.2 странно?

Бенчмарк LLM

Курсорим бенчмарк, где загружаем только LLM часть, она самая долгая, поэтому экспериментировать будем с ней. Тут я отключил TensorRT для FLOW, для чистоты эксперимента.

Решаю попробовать комбинации:
fp16 - всегда включен
sdpa - дефолтное внимание
flash attension 2 - попробуем его
sdpa + torch.compile - база

Запускаем:
sdpa: ~49 TPS
flash attention 2: 30-35 TPS - стало хуже
sdpa + torch.compile: 95-105 TPS - вот это нам и надо

Докручиваем

Применяем fp16 + sdpa + torch.compile (default mode, другие у меня под wsl не завелись), запускаем полную генерацию:
TTFB: 1.5 sec
RTF: 0.8
TPS: ~35

Уже лучше, но в бенчмарке мы видели TPS 100, что-то поджимает LLM и не дает ей работать.
Еще раз анализируем библиотеку, находим цикл while и внутри time.sleep(0.1).

Код написан так, что мы ждем первые 25 токенов от LLM, потом запускаем FLOW и HIFT на эти токены, получаем чанк аудио. Но если LLM сгенерировала 24 токена, то мы ждем еще 100ms, а пока работает FLOW и HIFT, LLM стоит на месте и ждет.
Разнести модели по разным контекстам можно, но это позже, мы тут вайбкодим вообще-то. Поэтому просто уменьшаем паузу с 0.1 до 0.005.

Запускаем генерацию:
TTFB: 1.4 sec
RTF: 0.7
TPS: ~45

Вспоминаем что выключали TensorRT (он ускоряет FLOW, а по коду LLM ждет FLOW), включаем обратно, запускаем:
TTFB: 0.8 sec
RTF: 0.42
TPS: ~65

Такие метрики уже вкуснее, для начала думаю хватит.

Итого

Генерация трех фраз:
Средний TTFB: 0.796 сек
Средний RTF: 0.427
Общая длительность: 23.160 сек
Общее время: 9.897 сек

TTFB: 2.2 -> 0.8 sec
RTF: 1.5 -> 0.4 sec

Исправленный инференс я планирую закинуть на github, будет отдельный пост.
Думаю можно ускориться еще, так как LLM зажимается последовательностью пайплайна, а в бенчмарках мы получили 100 TPS. Предполагаю еще +30% за счет нормального распараллеливания можно найти.
Примеры речи этой модели можно послушать в телеграме – xVibeNot