A kernel-centric path to real-time video generation on Trainium

Summary

Reactor and the AWS Neuron Science team collaborated to optimize real-time video generation on AWS Trainium. They focused on autoregressive diffusion models, which combine sequential next-token prediction with iterative diffusion to generate video frames that respond to user input. The teams used the Rolling Forcing model as a proxy, addressing challenges such as dynamic shapes, unusual memory access patterns, and heavy cache management that generic compilers struggle with. Optimizations included splitting SDPA heads across four cores and sequences across two, and batching common components of the diffusion and cache update phases to improve hardware utilization. The collaboration achieved real-time video generation above 16 frames per second on a single Trainium2 chip, a result that generalizes across autoregressive diffusion models.

Takeaways

1. 1

   Rolling Forcing, a streaming autoregressive diffusion model, requires generating and consuming each frame in real time, with generation remaining ahead of playback to avoid breaking the user experience.

2. 2

   Real-time video generation on Trainium benefits from a hybrid sharding approach: heads are split across 4 cores and sequences across 2, maintaining both mathematical correctness and data layout while improving compute utilization.

3. 3

   Batching the common components of the diffusion and cache update phases—rather than executing them separately—improves hardware utilization on Trainium’s multi-core architecture by avoiding underutilized partitions.

FAQ