---
title: "A kernel-centric path to real-time video generation on Trainium"
description: "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."
---

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

[AWS](https://yomu.fyi/company/aws) · Stephen Zorio · Sep 25, 2026

**Type:** Problem & solution

## 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.

## Context

Real-time interactive video generation using autoregressive diffusion models presents unique challenges due to dynamic sequence lengths, sliding window operations, and the need for low-latency frame generation. The collaboration between Reactor and the Neuron Science team aimed to optimize these workloads for Trainium hardware, which offers high memory and parallelization but requires careful kernel-level optimization to meet real-time constraints.

## Approach / What changed

The teams adopted a kernel-centric, bottom-up methodology focusing on dynamic shapes, memory access patterns, and cache management. They used Rolling Forcing as a proxy model and applied optimizations including hybrid sharding (SDPA heads split across 4 cores, sequences across 2), batching common components of diffusion and cache update phases, and optimizing memory access for RoPE and attention transposes. They achieved real-time generation (>16 fps) on a single Trainium2 chip.

## Takeaways

- 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.
- 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.
- 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

### What is the Rolling Forcing model?

Rolling Forcing is an autoregressive diffusion model for video generation that streams frames one by one. It generates 16 frames per second, meeting the standard for video playback, and uses a sliding window to continuously generate and evict frames.

### What specific optimizations were applied to achieve real-time video generation on Trainium?

The teams addressed dynamic shapes, memory access patterns, and cache management. Key optimizations included hybrid sharding of attention heads and sequences, and batching common components of the diffusion and cache update phases to improve hardware utilization on Trainium's 16-chip, 8-core per chip architecture.

### Does the optimization approach generalize beyond Rolling Forcing?

Yes. The teams emphasized that the techniques are common for any autoregressive diffusion model requiring real-time streaming. Rolling Forcing served as a small proxy that includes all building blocks (encoder, DIT, VAE, decoder) of more robust systems.

**Tags:** [AWS](https://yomu.fyi/topic/aws), [Inference](https://yomu.fyi/topic/inference), [Real-time](https://yomu.fyi/topic/real-time), [Video Generation](https://yomu.fyi/topic/video-generation)

- Source: [AWS](https://www.amazon.science/blog/a-kernel-centric-path-to-real-time-video-generation-on-trainium)
- Source URL: https://www.amazon.science/blog/a-kernel-centric-path-to-real-time-video-generation-on-trainium
- Ingested by Yomu: 2026-09-25T15:00:08.073Z

[Read original post](https://www.amazon.science/blog/a-kernel-centric-path-to-real-time-video-generation-on-trainium)