Reactor и Amazon Neuron Science довели генерацию видео на Trainium до реального времени

Что именно сделали Reactor и Neuron Science

Amazon Science 25 сентября 2026 года опубликовала разбор совместной работы команды Neuron Science и платформы Reactor: инженеры довели генерацию видео до реального времени на ускорителе Trainium, а не просто ускорили офлайн-рендер. Линия задачи восходит к идее world models — моделей, которые учатся сжатому представлению пространства и времени, — и к опубликованной в 2018 году работе «World Models» Дэвида Ха и Юргена Шмидхубера; сам блог Amazon Science называет Шмидхубера первым, кто предложил эту концепцию ещё в 1990 году. За прошедшие годы такие модели вышли за пределы исследовательских стендов: они используются в робототехнике, на транспорте, в климатическом моделировании, разработке игр и научной симуляции.

Reactor — платформа, которая даёт разработчикам, дизайнерам и исследователям разворачивать, использовать и масштабировать интерактивные ИИ-модели в реальном времени. Её соучредитель и технический директор Брайс Шмидтхен формулирует задачу коротко: «Речь о реальном времени, о низкой задержке и о том, чтобы делать это максимально эффективно на масштабе». Эффективность в его описании — это всё, от планирования инференса на конкретном чипе, в данном случае Trainium, до плотной «упаковки» каждого прямого прохода модели.

Команда Neuron Science занимается тем, что ищет новые способы включить генеративные модели на Trainium: новая архитектура, новый алгоритм оптимизации или новый способ генерировать и оптимизировать код для запуска этих моделей. Principal applied scientist Джун Ву описывает эту работу так: они находят возможность, собирают прототип и делают его пригодным для переноса в производственный конвейер. Вместе с ним в проекте участвовали senior applied scientists Мейсон Фу и Линфан Ю.

Поводом послужил сдвиг в самой генерации видео. По наблюдению Ву, рынок переходит от коротких видео фиксированной длины к видео бесконечной и динамической длины, а это требует autoregressive diffusion моделей — они соединяют последовательное предсказание следующего токена из больших языковых моделей с итеративным уточнением диффузионных моделей. Именно такие модели позволяют превращать нажатие клавиши в часть входных данных: модель, опираясь на предыдущий сгенерированный кадр, решает, каким будет следующее движение или следующая сцена. Команды рассмотрели несколько техник генерации видео и сошлись на Rolling Forcing — он стабильно выдаёт качественные тридцатисекундные ролики и при этом относительно невелик по размеру.

Почему в этой задаче всё решает не качество, а поток

Сложность тут не в качестве картинки, а в том, что модель обязана работать в реальном времени. Классическая генерация видео отрисовывает весь клип целиком и возвращает его позже; Rolling Forcing и подобные модели работают потоково. Каждый кадр генерируется и сразу же потребляется — показывается пользователю или подаётся назад как следующий вход. Именно так эти модели и используют, поэтому опоздавший кадр ломает весь опыт: генерация обязана идти с опережением таймлайна воспроизведения. Ву объясняет дополнительную трудность так: качественные диффузионные модели генерации требуют очень длинной последовательности, а значит большого расхода памяти; Rolling Forcing сам по себе невелик, но его последовательность очень длинна. Сочетание маленькой модели, очень длинной последовательности и жёсткого порога по частоте кадров и делает реальное время настолько требовательным.

Ещё одним доводом в пользу Rolling Forcing стала его способность выдавать 16 кадров в секунду — стандарт воспроизведения видео, а значит и соответствие требованиям по задержке. Здесь и появляется ценность Trainium: он приносит производительность и память, нужные, чтобы выдерживать длинную последовательность на скорости и давать генерацию выше 16 кадров в секунду, а не просто однажды выдать хорошие кадры. В итоге команды занялись включением Rolling Forcing как представителя всего класса autoregressive diffusion моделей генерации видео на Trainium.

Три узких места, которые не вытягивает общий компилятор

Команды Neuron Science и Reactor выбрали «ядерный», восходящий подход — от отдельных вычислительных ядер к целому конвейеру — и исходили из того, что разработчику должно быть проще. Они сосредоточились на трёх трудностях, которые генерация видео в реальном времени создаёт для универсальных компиляторов: динамические формы, нестандартные схемы доступа к памяти и тяжёлое управление кэшем. Каждая из них повторяется на каждом прямом проходе, поэтому задержка, незаметная на других нагрузках, в потоковом видео накапливается до срыва требований по времени.

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

Вторая трудность — схемы доступа к памяти. Помимо копирования скользящего KV-кэша, в нагрузке есть операции, чьи обращения к памяти специфичны именно для этой задачи: rotary position embeddings (RoPE) и транспонирование для внимания — универсальному компилятору их полностью оптимизировать тяжело. В генерации видео RoPE работает не с одной позицией, как в языковых моделях, а сразу с тремя осями: высотой, шириной и временем. Трёхмерное вращательное вложение перемежает чётные и нечётные элементы по самой внутренней размерности, что при общей компиляции порождает множество крошечных пересылок данных.

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

Neuron Kernel Interface: спуск на уровень ядер

Чтобы разобраться с этой дополнительной сложностью, команда Neuron Science обратилась к Neuron Kernel Interface (NKI). Официальная документация AWS Neuron описывает NKI как «bare-metal» язык и компилятор для прямого программирования устройств Trainium и Inferentia: он даёт управление на уровне инструкций, чтобы разрабатывать, оптимизировать и запускать собственные операторы на NeuronCore. Документация советует перед написанием ядра проверить библиотеку NKI — возможно, оптимизированное ядро для нужной операции уже есть.

Ву объясняет смысл инструмента так: NKI позволяет разработчикам писать вычислительные ядра, которые исполняются прямо на оборудовании NeuronCore, с точным контролем над тем, как данные перемещаются между памятью и вычислительными блоками. По его словам, это даёт разработчику самостоятельный путь к починке горячих точек — конкретные операции-узкие места заменяются реализациями, настроенными под железо, там, где профилирование показывает, что одной компиляции модели недостаточно.

Эффект команда измеряет прямо в миллисекундах. Ядро трёхмерного RoPE сократило своё время с пяти секунд до 1,8 миллисекунды, копирование кэша — с 23 миллисекунд до 1,9 миллисекунды на слой, а транспонирование для внимания исчезло вовсе: его вплавили в само ядро внимания. Как замечает Ву, для нагрузок реального времени, где важна каждая миллисекунда, именно прямой доступ к оборудованию и делает достижимой производительность производственного уровня.

Отдельно стоит результат инструментария: NKI-Dev-Suite — агент для генерации ядер NKI — выдал рабочее ядро трёхмерного RoPE с первой попытки. Суммарный эффект оказался заметен и по памяти: конвейер уложился в 11 ГБ высокоскоростной памяти, тогда как стандартный путь в eager-режиме упирался в нехватку памяти.

Гибридный шардинг: 12 головок внимания на восемь ядер

Видеодиффузионные модели порождают последовательности токенов, куда более длинные, чем текстовые, и это обостряет проблему self-attention. Линфан Ю приводит конкретные величины: внимание здесь работает с 23 400 токенами запроса, которые смотрят на 32 760 токенов контекста, и на него уходит около 70 % времени вычислений. Ни одно отдельное ядро не справится с этим эффективно без распределения работы.

Для этого команда применила гибридную схему шардинга: sequence parallelism (SP), то есть последовательное разбиение данных, и tensor parallelism (TP), который делит тензоры по определённой размерности, распределяя вычисления между устройствами. Для некоторых частей модуля, не относящихся к self-attention, хватило последовательного параллелизма. А вот для участков с self-attention подошёл только гибрид.

Причина — в устройстве внимания у Rolling Forcing. Ву напоминает, как устроено внимание в языковых моделях: оно причинное и однонаправленное, каждый токен смотрит только на предыдущие, KV-кэш растёт монотонно, а на слой приходится один тип внимания с единственной политикой кэша. У Rolling Forcing на каждый блок приходится два типа внимания: self-attention для пространственно-временной согласованности между кадрами и cross-attention для текстового обусловливания и двунаправленного внимания внутри активного окна, поскольку все кадры уточняются из шума совместно.

К этому добавляются KV-кэш с двумя политиками — скользящее окно для недавнего контекста плюс постоянный «сток» внимания для глобального контекста, — два прямых прохода на каждое окно (на денойзинге записываются «шумные» записи KV, а проход обновления кэша перезаписывает их чистыми значениями, чтобы следующие окна всегда опирались на чистый контекст) и трёхмерная структура видеотокенов. Всё это ограничивает то, как последовательность можно разрезать между ядрами.

Ю объясняет, что один только TP упирается в арифметику: у диффузионного трансформера WAN двенадцать головок внимания, а на чипе восемь ядер Neuron — число не делится, и при чистом TP пришлось бы добавлять дополнение, которое расходует вычисления впустую. Один только SP ломает трёхмерную структуру: как замечает Ю, нельзя гарантировать, что граница разбиения последовательности совпадёт с границей кадра, а данные должны быть как минимум в гранулярности кадра — иначе эти операции не сработают. Гибридный подход решает обе задачи: головки делятся между четырьмя ядрами, а последовательность — между двумя, так что и арифметика, и раскладка данных остаются корректными. Декодер VAE при этом использовал пространственный шардинг по оси W и дал сверхлинейное ускорение в 8,25 раза.

Правки в самой модели и что из этого следует

Команды оптимизировали не только ядра, но и код самой модели Rolling Forcing. Ю описывает её устройство так: у модели две фазы — диффузия и обновление кэша; они исполнялись двумя отдельными запусками, а проблема в том, что у фазы обновления кэша значительно меньше вычислений, поэтому в отдельном раунде её загрузка оборудования оказывается куда ниже — сказывается и необходимость шардинга. Общий принцип, который он формулирует: чем меньше данных вы подаёте на чип, тем хуже используется оборудование.

Решение — сгруппировать друг с другом те компоненты, которые обе фазы делят между собой. Там, где компоненты слегка различаются, команда снова разделяет их и обрабатывает отдельно, но большую часть конвейера удалось объединить в пакет. По словам Ю, это особенно полезно именно для Trainium: в каждом инстансе шестнадцать чипов, в каждом чипе восемь ядер, и если разнести вычисление по слишком большому числу ядер, каждый вызов получит лишь небольшой фрагмент работы — это и есть недоиспользование ёмкости.

Результат: применив эти и другие оптимизации, команды успешно сгенерировали корректное видео уже на первом сквозном прогоне. Запустив Rolling Forcing на одном чипе Trainium2, участники получили генерацию видео в реальном времени — и, как они подчёркивают, результат обобщается на другие autoregressive diffusion модели. Principal solutions architect Яхав Биран объясняет выбор такой «маленькой» модели как раз её полнотой: Rolling Forcing — это законченная система со всей сложностью настоящей системы — энкодер, DIT, VAE, декодер, — поэтому более крупная система работает на тех же строительных блоках. Ю добавляет, что команда не оптимизирует одну конкретную модель: она строит общие техники для моделей с autoregressive diffusion, которым нужно реальное взаимодействие, и поддерживает все модели с требованиями потоковой передачи в реальном времени.

Trainium в этом рассказе — не фон, а условие задачи: AWS описывает его как ускоритель, спроектированный ради одной цели — лучшей экономики высокопроизводительного обучения и инференса на масштабе. Шмидтхен формулирует, что дальше: «В не столь далёком будущем каждый пиксель будет генерироваться в реальном времени, интерактивно». Целые истории, по его словам, смогут создаваться world-моделями в реальном времени — истории, которые реагируют, с которыми можно взаимодействовать и которые могут менять весь свой ландшафт на ходу. Он отмечает и более ясную дорожную карту Trainium: рост производительности, лучшая стоимость производительности, больше масштаба по миру; в будущем, где интерактивный ИИ нужно распределять по потребителям и в физический мир, Trainium, по его оценке, хорошо позиционирован на уровне инференса — и по параллелизации, и по памяти, и по программному стеку.