Amazon and Reactor tune Trainium for real-time video generation

Why real-time world models need new optimization

World models turn text, image, audio and video inputs into a latent space and can predict state changes and respond to actions. Amazon Science says they are playing growing roles in robotics, transport, climate modeling, video game development, scientific simulation and design prototyping. That growth raises the demands on hardware and infrastructure.

Reactor, which lets developers, designers and researchers deploy and scale real-time interactive AI models, knows those demands well. Co-founder and chief technology officer Bryce Schmidtchen put it this way: "It's about real time, it's about low latency, and it's about doing that as efficiently as you can at scale," he said. Reactor already relies on the global reach of AWS, which Schmidtchen said makes a platform of its scale "possible and reliable."

Autoregressive diffusion and the Rolling Forcing model

The Neuron Science team, which explores new techniques for enabling generative AI models, noticed a shift from short, fixed-length video generation to infinite-length or dynamic-length video. According to principal applied scientist Jun Wu, that shift gave rise to autoregressive diffusion models. These combine the sequential next-token prediction of large language models with the iterative refinement of diffusion models. A user keystroke can serve as input, and the model decides the next scene based on the previously generated frame.

Wu, together with senior applied scientists Mason Fu and Lingfan Yu, worked with Reactor on the problem. The teams chose Rolling Forcing, citing its ability to consistently generate high-quality 30-second videos and its relative size. The difficulty, the article says, is not quality but timing. These models stream: each frame is shown to a user or fed back as the next input, so a late frame breaks the experience. Rolling Forcing is relatively small, but its sequence length is very large, and it needs to sustain 16 frames per second. Amazon says Trainium brings the performance and memory to deliver real-time generation above 16 fps.

Three problems for generic compilers

The teams took a kernel-centric, bottom-up approach aimed at making life easier for developers. They focused on three challenges that real-time video generation poses for generic compilers:

• Dynamic shapes. Frames are generated and evicted in a sliding window. "This means attention lengths vary across forward passes, and there are two distinct passes per window: denoising, then cache cleanup," Wu said, which makes static compilation hard.
• Unusual memory access patterns. Operations such as rotary position embeddings (RoPE) and attention transposes are workload-specific. In video, RoPE has to handle three axes (height, width and time) instead of the single position used in LLMs. Wu said the 3D rotary embedding produces many tiny data transfers when compiled generically.
• Heavy cache management. Every layer reads and writes a rolling KV cache, so on-device copies need high throughput.

Because these costs recur on every forward pass, a small delay can add up to a latency failure in real-time generation.

Results with the Neuron Kernel Interface

To address the bottlenecks, the team used NKI. Wu described it this way: "NKI lets developers write compute kernels that run directly on NeuronCore hardware, with precise control over how data moves between memory and compute engines." He said it gives developers a self-service way to replace specific bottleneck operations with hardware-tuned implementations when profiling shows that compiling the model alone is not enough.

The reported gains are large. The 3D-RoPE kernel went from five seconds to 1.8 milliseconds. Cache copies dropped from 23 milliseconds to 1.9 milliseconds per layer. Attention transposes were eliminated entirely.

What it means

According to Amazon Science, the work lays a foundation for future models to reach real-time capability on Trainium. For developers building interactive video and world-model applications, the main takeaway is that targeted kernels written with NKI can remove specific hotspots without waiting for general-purpose compilers to handle workload-specific patterns.