Обслуживание Ideogram v4 менее чем за секунду без потери качества

В @fal мы обслуживаем Ideogram V4 на высокой скорости. При разрешении 1K генерация, которая раньше занимала 2,75 с, теперь выполняется за 0,44 с — это ускорение в 6 раз без видимой потери качества. Мы добились этого, снизив стоимость генерации изображения по всем возможным направлениям, и полезно рассматривать эти направления как единое уравнение: стоимость изображения — это количество проходов вперед (forward passes), выполняемых диффузионным трансформером, умноженное на стоимость каждого прохода.

Содержание

• FP4 и оптимизации на системном уровне
• Дистилляция с учетом квантования
• Дистилляция временных шагов (Timestep Distillation)

Два фронта. Сделать каждый проход вперед дешевым: запустить трансформер в FP4 и объединить небольшие операции вокруг каждого матричного умножения (matmul), чтобы избежать лишних обращений к памяти. Сократить количество проходов: дистиллировать классификатор-свободное руководство (classifier-free guidance) из двух ветвей в одну и сократить множество шагов шумоподавления до нескольких. Объедините все это, и вы получите то же изображение, что и в полной модели bf16, при значительно меньших вычислительных затратах.

В этой статье мы пройдем весь путь, примерно в том порядке, в котором мы его строили. Мы начнем с прохода вперед в FP4 и слияния ядер (kernel fusion), которое ускоряет его, затем рассмотрим проблему качества, вызванную 4-битным представлением, и дистилляцию с учетом квантования, которая ее решила, и, наконец, дистилляцию временных шагов, которая сократила количество шагов. FP4 работает быстро, но «наивный» FP4 выглядел заметно хуже, чем bf16, и даже тогда он все еще выполнял много шагов, умноженных на две ветви руководства. Вот как мы решили все эти проблемы.

Запуск трансформера в FP4

Самый быстрый путь вывода — запуск диффузионного трансформера в NVFP4. FP4 — это 4-битный формат с плавающей запятой: один бит знака, два бита экспоненты, один бит мантиссы, поэтому каждое значение может принимать только одно из нескольких уровней. Взамен вы получаете значительное ускорение и гораздо меньший объем памяти, что позволяет нам повысить параллелизм и снизить задержку на одно изображение. На Blackwell этот формат поддерживается на аппаратном уровне, поэтому матричные умножения выполняются нативно в 4 битах.

Но 4-битное матричное умножение окупается только в том случае, если именно на него уходит основное время, а в диффузионном трансформере это часто не так. Каждое большое GEMM сопровождается небольшой операцией: RMSNorm на пути внимания, gated-SiLU в FFN. Без слияния эта небольшая операция заставляет выводить весь результат GEMM в HBM и обратно. Матричное умножение стало дешевле, но обращение к памяти, которое вы не устранили, съедает весь выигрыш. Поэтому первое, что нужно для быстрого пути FP4, — это объединить эти операции непосредственно в GEMM.

Слияние эпилога для окупаемости FP4

Кратко о слиянии эпилога

GEMM выполняет A @ B по тайлам (tiles). Внутри каждого тайла тензорные ядра (Tensor Cores) направляют умножение-сложение в аккумулятор, который находится в собственной памяти SM: регистры на Hopper, TMEM на Blackwell. Пока выполняется основной цикл, HBM еще не затронута для вывода; результат накапливается на чипе.

Но A @ B почти никогда не является концом слоя. В диффузионном трансформере FFN обычно представляет собой GEMM, за которым следует gated-SiLU (SwiGLU), а путь внимания — GEMM, за которым следует RMSNorm. Этот паттерн повторяется в LLM и DiT: большое GEMM, затем небольшая пост-операция. Сама операция дешева; проблема в том, что без слияния GEMM записывает весь тайл в HBM, второе ядро загружает его обратно, применяет операцию и записывает снова. Несколько FLOP превращаются в полное обращение к памяти плюс запуск второго ядра.

Слияние эпилога — это применение операций после GEMM непосредственно перед сохранением результата в HBM. Вывод попадает в глобальную память один раз, уже преобразованным — никакого второго ядра, никакого лишнего обращения. На Hopper и Blackwell эпилог может перекрываться с основным циклом следующего тайла, поэтому для небольших операций дополнительные вычисления часто почти незаметны.

CUTLASS предоставляет это через деревья посетителей эпилога (Epilogue Visitor Trees, EVT). Вы пишете visit(), который выполняется для каждого фрагмента аккумулятора, а коллективный эпилог обрабатывает загрузки и сохранения. Механика хорошо объяснена в статье Colfax об EVT. Краткая версия: поточечные слияния (смещение, ReLU, приведение типов) выполняются легко: каждый выходной элемент независим, нет ничего, что нужно редуцировать, и CUTLASS поставляет готовые посетители для них.

В этой статье мы сосредоточимся на слиянии RMSNorm. Она нормализует вектор по его собственному RMS: y = x · rsqrt(mean(x²) + eps) · weight. В отличие от поточечной операции, ей требуется rstd по всей голове, то есть редукция. Именно эта редукция делает слияние сложным, и чтобы понять почему, нам нужно посмотреть, как эпилог видит свои данные.

Тайлы и фрагменты на Blackwell

GEMM работает с тайлами CTA: на Blackwell (SM100) наш NVFP4 GEMM использует тайл CTA 256×256. Эпилог не обрабатывает весь этот тайл сразу; он проходит по нему меньшими тайлами эпилога 128×128, давая 4 подтайла эпилога на CTA:

Каждый подтайл эпилога распределяется между 128 потоками, поэтому каждый поток содержит фрагмент из 128 элементов (FragmentSize = 128). Обратный вызов visit() получает по одному фрагменту за раз — 128 значений, за которые отвечает этот поток. Это единица работы, которую видит эпилог, и это важно для того, что можно объединить.

Это делает слияние RMSNorm с head_dim=128 тривиальным: голова из 128 элементов — это в точности фрагмент одного потока. Когда вызывается visit(), поток уже держит все 128 значений одной головы в своих регистрах. Чтобы вычислить RMSNorm, он суммирует квадраты по этим 128 значениям (без перетасовки между потоками, без второго фрагмента, без дополнительного прохода), масштабирует и сохраняет. Редукция помещается в один вызов visit().

Ideogram v4 не так удобен. Его внимание использует head_dim=256, поэтому одна голова больше не помещается в один фрагмент; она охватывает два, и простая редукция за один вызов visit() перестает работать.

Почему head_dim=256 — это сложно

Чтобы понять, почему это проблема, посмотрите, как работает эпилог. После основного цикла коллективный эпилог CUTLASS проходит по 4 подтайлам тайла CTA в фиксированном порядке: Tile 0 → Tile 1 → Tile 2 → Tile 3, и для каждого подтайла он вызывает visit() один раз на фрагмент, сохраняет результат и переходит дальше. Это один проход вперед; нет возможности вернуться к предыдущему подтайлу.

RMSNorm требует rstd по всей голове, то есть редукции: rstd = rsqrt(mean(x²) + eps), затем y = x * rstd * weight. Чтобы вычислить это, нам нужна сумма квадратов по обеим половинам, но они приходят в разных вызовах visit(): первая половина в Tile 0, вторая в Tile 1. К тому времени, как мы видим Tile 1 и можем вычислить rstd, Tile 0 уже был посещен и сохранен. Нет способа вернуться назад и исправить это. Остальная часть этого раздела посвящена тому, как мы все же объединили это: график повторного посещения в два прохода, который сохраняет первую половину в регистрах.

Слияние с повторным посещением

Коллектив эпилога CUTLASS работает с фиксированным порядком цикла в один проход: посетить каждый подтайл один раз, сохранить, перейти дальше. Это не решает нашу проблему: нам нужно увидеть Tile 1, прежде чем мы сможем нормализовать Tile 0, что означает необходимость возврата. Поэтому мы написали эпилог, который позволяет посетителю проходить по подтайлам в порядке, необходимом для редукции (включая возвращение к тайлу, который он уже видел), вместо стандартного прохода вперед.

График для каждой группы строк состоит из трех касаний:

• Tile 0, проход 0: это первая половина головы. Мы вычисляем ее сумму квадратов и переходим дальше; мы пока не можем нормализовать без второй половины.
• Тайл 1, проход 1: прибывает вторая половина заголовка. Теперь у нас есть обе половины: мы вычисляем сумму квадратов для этой половины, прибавляем её к сумме первой половины и, наконец, получаем rstd. Мы нормализуем и сохраняем эту половину.
• Тайл 0, проход 2 (повторное посещение): мы возвращаемся к тайлу 0 и масштабируем его с помощью уже известного rstd. Это тот самый дополнительный штрих, который не может обеспечить однопроходный эпилог.

Всё остаётся на чипе: никаких обращений к глобальной памяти, никакого второго ядра. Пользовательский эпилог реализует схему с тремя обращениями; посетитель (visitor) переносит состояние редукции между этими обращениями и применяет масштабирование, когда rstd готов. Вместе они объединяют нормализацию до того, как данные попадут в HBM.

Слияние Gated-SiLU (SwiGLU)

Другой распространённый паттерн в диффузионных трансформерах (и LLM, таких как LLaMA) — это gated-SiLU. Один GEMM выдаёт выход C, затем слой завершается разделением его пополам и умножением:

Это сокращает ширину выхода вдвое; две половины становятся одной. История та же, что и с RMSNorm: математика простая, но неслитое второе ядро требует полного обращения к HBM. В CUTLASS нет встроенного эпилога для этого паттерна «пара + редукция N/2», поэтому для его слияния требуется пользовательский посетитель.

Проблема заключается в компоновке. Для выходного столбца n пара состоит из gate = C[:, n] и up = C[:, n + N/2], разделённых половиной N. Эти два значения могут легко оказаться в разных фрагментах, разных подтайлах или даже разных CTA. Синхронизация между тайлами разрушила бы параллельную модель эпилога.

Вместо этого мы один раз переставляем столбцы B во время упаковки весов. Затем GEMM выдаёт чередующиеся пары [gate0, up0, gate1, up1, ...] вместо двух разделённых половин. Поскольку смежность гарантирована, пользовательский посетитель EVT применяет SiLU и умножение к соседним элементам в одном фрагменте и сохраняет редуцированный результат: одно ядро, никакой межтайловой коммуникации.

Полное руководство (упаковка весов, хуки посетителя и настройка EVT) приведено в статье «Создание эффективных ядер с помощью слияния эпилогов» (Crafting Efficient Kernels with Epilogue Fusion).

Подвох: FP4 снижает качество изображения

Слияние эпилога даёт нам быстрый прямой проход в FP4. Само по себе это также делает результат менее качественным. Стандартный вывод FP4 выглядел заметно хуже, чем bf16, особенно в плане цвета. Красные оттенки блекли, насыщенность падала, а изображения с обилием цветов (логотипы, продукты, растения) выглядели менее насыщенными по сравнению с эталоном bf16.

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

Методы постобработки (коррекция яркости, перебалансировка цвета) существенно не помогли, так как ошибка была заложена в латентные представления до декодирования, а не просто была проблемой отображения. Нам нужно было исправить это внутри модели, а не после неё.

Что не сработало

Прежде чем прийти к окончательному рецепту, мы попробовали несколько вещей, которые научили нас тому, что важно, а что нет:

• Постобработка вывода. Нашим первым порывом было просто исправить цвета после генерации. Коррекция насыщенности, регулировка яркости, перебалансировка цвета на декодированном изображении. Ничего из этого не сработало. Обесцвечивание — это не проблема отображения; это проблема латентной траектории. К тому моменту, когда у вас появляются пиксели для исправления, денойзер уже сделал немного неверные шаги на многих итерациях, и цветовая информация просто потеряна. Её невозможно «раз-запечь» из финального изображения.
• Наивная дистилляция. Мы попробовали самое простое: квантовать студента до FP4, запустить учителя bf16 на том же входе и минимизировать MSE между их предсказаниями. Градиент не проходил через квантование — студент был заморожен в своём квантованном состоянии и просто сравнивался с учителем. Функция потерь снижалась, изображения не становились лучше. Именно тогда мы на собственном горьком опыте узнали, что функция потерь при обучении диффузии не предсказывает качество изображения. Вы можете наблюдать, как она падает тысячи шагов, а результат почти не меняется.

Это подсказало нам две вещи: ошибку квантования нужно устранять во время обучения (а не после), а окончательное суждение должно быть визуальным (а не числовым).

Дистилляция с учётом квантования

Примерно в то время, когда мы боролись с этим, для LLM и VLM была опубликована техника под названием «дистилляция с учётом квантования» (QAD) paper. Идея проста, если смотреть назад: вместо того чтобы просто квантовать модель и надеяться на лучшее, вы дистиллируете полноточного учителя в квантованного студента. Студент учится предсказывать то, что предсказал бы учитель, но в условиях ограничения собственного шума квантования. Учитель — это замороженный bf16, студент — квантованный, и вы обучаете студента соответствовать выводу учителя.

В той статье использовалась дивергенция Кульбака-Лейблера (KL) на логитах, потому что речь шла об LLM. Мы предсказываем не токены; мы предсказываем поле скоростей над латентным пространством. Поэтому мы заменили функцию потерь: MSE между предсказанием скорости студента и предсказанием скорости учителя. Тот же принцип, другая цель.

Критическое отличие от наивной дистилляции, которую мы пробовали ранее: в QAD студент не просто работает в FP4 и сравнивается с учителем. Студент обучается через собственное квантование. Градиенты текут обратно через округление FP4 с помощью прямого оценочного метода (straight-through estimator), поэтому веса активно адаптируются к шуму квантования, который они производят. В этом весь смысл — веса перемещаются в ту часть пространства параметров, где округление до 4 бит не вредит результату. Наивная дистилляция просто сравнивает студента с замороженным квантованием с учителем и надеется на успех; QAD позволяет квантованию стать частью оптимизации.

При прямом проходе веса и активации студента в bf16 квантуются до FP4 для матричного умножения, а затем деквантуются обратно. Выход несёт реальный шум квантования от этого цикла. Замороженный учитель bf16 создаёт целевую скорость. Функция потерь — это MSE между зашумлённым предсказанием студента и чистой целью учителя. При обратном проходе округление FP4 отсекается от графа автограда: функция округления имеет нулевой градиент почти везде, поэтому её нельзя дифференцировать напрямую, и прямой оценочный метод просто обходит её, позволяя градиентам течь в полной точности. Прямой проход несёт шум квантования; обратный проход игнорирует его. В процессе обучения веса сходятся к области пространства параметров, где округление до 4 бит не ухудшает результат.

При инференсе ничего из этого не нужно. Мы экспортируем обученные веса, загружаем их в производственный конвейер, применяем реальное статическое квантование FP4 с реальными ядрами и генерируем изображения. Студент работает просто как условная ветвь с guidance = 1 — без учителя, без CFG, без STE.

Сворачивание CFG в один проход

Classifier-free guidance (CFG) — это способ, с помощью которого диффузионные модели ориентируются на промпт. Во время инференса вы запускаете трансформер дважды на каждом шаге: один раз с условием по промпту (conditional) и один раз без текста (unconditional). Это удваивает вычислительные затраты на каждом шаге шумоподавления. Мы поняли, что можем включить это в дистилляцию и получить бесплатно.

Во время обучения замороженный bf16-учитель выполняет обе ветки и выдает направленную скорость CFG. Студент представляет собой только условную ветку и обучается напрямую предсказывать эту направленную скорость. При инференсе мы запускаем только обученную условную ветку с guidance = 1 — один прямой проход вместо двух. Таким образом, помимо ускорения FP4, мы сокращаем затраты на трансформер вдвое, полностью исключая CFG из пути обслуживания.

Это та часть, которая выходит за рамки статьи QAD. Эта статья восстанавливает точность квантования и на этом останавливается. Мы восстанавливаем точность квантования и дистиллируем ветку CFG, поэтому итоговый граф инференса становится меньше (fp4) и короче (одна ветка).

Что заставило QAD работать

Чтобы QAD действительно восстановил визуальное качество, потребовалось принять несколько решений помимо основного рецепта. Вот те, которые имели наибольшее значение:

• Комбинация функций потерь. Мы используем функцию потерь flow-matching в сочетании с функцией потерь дистилляции, чтобы использовать качество реальных данных. Одной лишь чистой дистилляции было недостаточно для восстановления полного визуального качества, к которому мы стремились.
• График guidance. Мы не использовали фиксированный guidance для целевого значения учителя. Мы использовали более высокий guidance на шумных и средних шагах и более низкий ближе к чистому результату, потому что модель выигрывает от более сильного управления в начале и более мягкого по мере сходимости.

Поскольку QAD восстанавливает качество, а CFG свернут в одну ветку, два из трех множителей затрат устранены. Последний — самый значительный.

Дистилляция по количеству шагов (Timestep distillation)

Теперь мы сделали каждый прямой проход настолько дешевым, насколько это возможно: 4-битные матричные умножения в FP4, объединенные эпилоги, чтобы нормы и активации не требовали обратного прохода, и одна ветка управления вместо двух. Единственное, чего мы не коснулись, — это количество прямых проходов, то есть количество шагов шумоподавления, и это самый важный рычаг из всех.

Диффузионный трансформер достигает своего качества за счет шумоподавления латентного представления в течение многих итераций, и каждая итерация — это полный прямой проход. Если вы можете научить модель выполнять ту же работу за несколько шагов, вы устраняете самый большой фактор, влияющий на задержку и пропускную способность. Это и есть дистилляция по количеству шагов. Это также часть конвейера с наибольшим количеством движущихся частей, наибольшим объемом литературы и наибольшим количеством способов незаметно испортить качество изображения.

Основа: Distribution Matching Distillation

Естественная первая идея — регрессия: запустить учителя, запустить студента, минимизировать MSE между ними. Это то, что делает работа QAD, упомянутая выше, и это работает, когда студент и учитель идут по одному пути. Но студент, работающий за несколько шагов, не идет по тому же пути, что и учитель, работающий за много шагов. Вы не можете требовать от него попадания в те же промежуточные латентные представления, потому что он пропускает большинство из них. Регрессия привязывает студента к траектории, с которой он должен иметь возможность сойти.

Distribution Matching Distillation (DMD) переосмысливает цель: вместо того чтобы сопоставлять учителя сэмпл за сэмплом, сопоставляйте распределение с распределением. Нас не волнует, что делает студент на любой отдельной траектории; нас волнует, чтобы облако изображений, которое он создает, выглядело как облако, созданное учителем.

Механизм поддерживает две сети оценок (score networks): реальную оценку (замороженный учитель) и поддельную оценку (вторая сеть, отслеживающая текущее распределение студента). Чтобы обновить студента, DMD берет сэмпл студента, добавляет шум и спрашивает обе сети оценок, куда двигаться. Разница между двумя оценками — это градиентное направление, которое перемещает распределение студента к распределению учителя. Эта разница оценок, с точностью до константы, является градиентом дивергенции Кульбака-Лейблера между двумя распределениями на этом уровне шума. Никаких парных целей, никакой MSE для конкретного латентного представления, просто «сделайте так, чтобы эти два распределения совпали».

DMD2: версия, которую люди действительно используют

DMD2 приводит классический DMD к рецепту, от которого происходят большинство современных моделей генерации изображений за несколько шагов: он отбрасывает член регрессии (делая метод не зависящим от данных), использует обновления с двумя временными шкалами, чтобы поддельная оценка оставалась точной при изменении студента, добавляет GAN-потерю в качестве вторичного члена для повышения резкости и поддерживает многошаговую выборку, чтобы студент не был привязан к одному количеству шагов.

Целое семейство, каждое из которых преследует свой рычаг

DMD2 положил начало семейству последующих работ. Правильный выбор зависит от того, что вы оптимизируете: DMDR, Flash-DMD, AdvDMD и GNDM включают RL и сигналы вознаграждения в дистилляцию. DP-DMD сохраняет разнообразие. AMD улучшает стабильность. Decoupled DMD показывает, что член guidance является настоящим двигателем, в то время как сопоставление распределений действует как регуляризатор.

Ветка, наиболее близкая к нашей работе, — это сопоставление траекторий и внетраекторное сопоставление. Стабильность обучения была серьезной проблемой: GAN-потеря — это хрупкая часть DMD2, поэтому мы хотели зайти как можно дальше без дискриминатора. Это подтолкнуло нас к сопоставлению распределений без GAN и без данных. Два ближайших публичных метода:

• TDM сопоставляет распределения в точках вдоль траектории учителя с целевой функцией, учитывающей количество шагов выборки, поэтому одна модель работает с несколькими значениями количества шагов. Она остается на траектории; ее гибкость обусловлена учетом количества шагов.
• CDM выходит за пределы траектории, обеспечивая сопоставление распределений на экстраполированных латентных представлениях из собственного поля скоростей студента, сохраняя надежность при изменении графика.

Наша собственная цель — это кастомный дизайн сопоставления распределений без GAN и без данных в том же духе. Целью была стабильная дистилляция, которая достигает высокого качества и остается гибкой при разном количестве шагов, а не какой-то один опубликованный рецепт.

Во-первых, дистилляция без GAN

Мы сделали первый проход без GAN: только наша кастомная цель сопоставления распределений, без дискриминатора. Это дало нам студента, работающего за несколько шагов, который был действительно хорош: с правильной композицией, точным следованием промпту и стабильным поведением при гибком графике шагов. На большинстве промптов вам было бы трудно отличить его от учителя. Просто он был не совсем производственного уровня. Последнего прироста резкости и точности мелких текстур еще не было. Сопоставление распределений делает распределение правильным; получение окончательной четкости каждого сэмпла — это другая задача.

Затем добавление GAN

Дискриминатор, обученный на реальных данных, знает, что такое «четкий и реальный», лучше, чем любая аппроксимация оценок. Проблема состязательного обучения заключается в том, что оно обычно рушится, когда студент начинает со слабых позиций: дискриминатор побеждает мгновенно, градиенты насыщаются, и генератор никогда не учится. Большая часть литературы тратит усилия на ограничение дискриминатора с помощью аппроксимированной R1-регуляризации (как в APT) или случайных проекций, что мы и использовали в нашей более ранней дистилляции FLUX.2 dev Turbo.

Для этой модели мы вообще не использовали стабилизатор, и он нам не понадобился. Поскольку мы начинаем с уже сильной дистиллированной модели-ученика, её выборки близки к реальным с первого шага, поэтому дискриминатор никогда не получает легкой победы, которая приводит к коллапсу. Обучение остается стабильным само по себе, а GAN восстанавливает остатки резкости и детализации.

Это противоположность тому, как DMD2 использует состязательную функцию потерь, где DMD2 накладывает GAN в качестве вторичного члена. Для нас это основная задача после обучения, и она работает, потому что ученик, с которого мы начинаем, уже силен. Недавний двухэтапный метод исследует похожую идею использования GAN в качестве функции потерь после обучения.

Возвращаясь к FP4

Ученик с малым количеством шагов работает в формате bf16. Чтобы его обслуживать, мы пропускаем его через тот же путь QAD: дистилляция с учетом квантования против теперь уже малошагового учителя, STE, экспорт, реальное статическое квантование FP4. Итоговая модель является малошаговой, 4-битной и однопоточной — на этот раз не нужно сворачивать CFG, поскольку малошаговый ученик является однопоточным по своей структуре.

К чему это нас приводит

Итоговый путь обслуживания — это малошаговая, однопоточная модель FP4 с объединенными ядрами. Каждый метод решает свою часть уравнения стоимости, и они суммируются: FP4 и слияние эпилогов делают каждый прямой проход дешевым, дистилляция CFG сокращает количество ветвей вдвое, дистилляция временных шагов сокращает количество шагов, а QAD восстанавливает качество, потерянное при 4-битном квантовании. Результатом является модель изображений, которая выдает тот же результат, что и многошаговый сэмплер bf16 с CFG, при доле вычислительных затрат, без видимой потери качества на промптах, которые раньше приводили к сбоям стандартного FP4.

Оба производственных уровня используют этот путь FP4, что является важной причиной, по которой они превосходят базовый уровень bf16. При разрешении 1K:

Fast сворачивает две ветви управления в одну; Instant дополнительно сокращает количество шагов, работая в 6,3 раза быстрее, чем Base.