In 2018, Jürgen Schmidhuber — who in 1990 was the first person to propose world models as a machine learning concept — published a paper with David Ha. They wrote about “a predictive world model” which could “extract useful representations of space and time” and use that to train agents to drive, among other things. The growth in the capability and real-world applications of world models in the eight years since that seminal paper is staggering.
The emergence of massive training datasets, paired with variational autoencoders, and both diffusion and autoregressive transformers enable today’s world models to convert text, images, audio, and video inputs into latent space that is used to build and maintain breathtaking, immersive world models which can predict state changes and respond to actions. These models have vital and expanding roles in robotics, transport, climate modeling, video game development, scientific simulation, and design prototyping.
The sharp increase in the power and potential of world models has been accompanied by an increased need for hardware and infrastructure capable of supporting them. That is a challenge that Bryce Schmidtchen, co-founder and chief technology officer of Reactor, is well acquainted with.
“It's about real time, it's about low latency, and it's about doing that as efficiently as you can at scale,” Schmidtchen observed. Reactor is a platform that allows developers, designers, and researchers to deploy, use, and scale real-time interactive AI models. “Efficiency means everything from how you schedule the inference on the given chip, in our case Trainium, to how you think about maximally bin packing every forward pass of the model.”
The challenge of efficiency is made even more acute by Reactor’s global presence. “We think about GPU clusters in terms of regions — we have hundreds all over the world. The latent that turns into a pixel that comes off the GPU needs to be able to get to the network without hopping around through Kubernetes. And you need to do this in a way where it doesn't matter if it's a well-supported cluster or bare metal in a closet.
“And then,” Schmidtchen continued, “you need to tie this to world-class networking and media streaming that supports different codecs, different resolutions, and allows that to connect to APIs and SDKs across different languages that can support all different types of applications.”
The global reach of AWS already helps Reactor to achieve many of its goals. “Everything they have at the inference layer, the level of scale that we're able to achieve in different regions is what makes a platform of this scale possible and reliable,” Schmidtchen observed.
Now, a recent collaboration between Reactor and the Amazon Neuron Science team has forged a path to even greater efficiency for world models. This is a look at the optimization strategy those teams pursued, and how that laid the foundation for future models to achieve real-time capability on Trainium.
The rise of autoregressive diffusion models
“The Neuron Science team explores new techniques for generative AI model enablement optimization,” explained Jun Wu, a principal applied scientist on the Neuron team. “This might be a new model architecture or new algorithm for optimization or a new way to generate and optimize the code for running those models on Trainium. We identify opportunities to build a prototype and make it feasible to be ported to the production pipeline.”
The team spotted one such opportunity around the usage of diffusion models in video generation. “We noticed an evolution from generating shorter, fixed-length videos to infinite-length or dynamic-length videos,” Wu explained. “Generating high-quality video frames at those lengths gave rise to the usage of autoregressive diffusion models.”
Those types of models, which combine the sequential next-token prediction found in large language models with the iteration and refinement of diffusion models, are essential for users who want to generate video where they can navigate the generated environment.
“Autoregressive diffusion means a user keystroke can be absorbed as input to the model and, conditioned on the previously generated video frame, the model can correctly decide the next move or the next scene,” Wu observed. “Most of the interactive video generation models we are seeing today are using this.”
Wu and his team, Mason Fu and Lingfan Yu, both senior applied scientists, saw a chance to optimize how those models are deployed for real-time video generation on Trainium, and saw the opportunity to do this in collaboration with Reactor. “We had been working on real-time video and interactive video generation for over 18 months at that point,” Schmidtchen noted. The teams considered various video generation techniques and aligned on utilizing Rolling Forcing, citing both its ability to consistently generate high-quality 30-second videos and its relative size.
The hard part with these models isn't quality, it's that they have to run in real time. Unlike traditional video generation, which renders a full clip offline and returns it later, models like Rolling Forcing are streaming. Each frame is generated and immediately consumed, shown to a user or fed back in as the next input. That is how developers and customers actually use them, and a frame that arrives late breaks the experience, because generation has to stay ahead of the playback timeline. “High-quality generation diffusion models pose the challenge of a very long sequence, which requires a lot of memory consumption,” Wu explained. “Rolling Forcing is relatively small, but its sequence length is very large.” That combination of a small model, a very long sequence, and a hard frame-rate floor is what makes real time so demanding. Rolling Forcing's ability to generate 16 frames per second, the standard for video playback, meant it also met latency requirements. This is where Trainium adds value, bringing the performance and memory to sustain that long-sequence workload at speed and deliver real-time generation above 16 fps rather than merely producing good frames eventually. So the teams set about enabling Rolling Forcing, as a proxy for autoregressive diffusion video generation on Trainium.
A developer-friendly approach
The Neuron Science and Reactor teams adopted a kernel-centric, bottom-up methodology aimed at making life easier on developers. They focused on three challenges — dynamic shapes, unusual memory access patterns, and heavy cache management — that real-time video generation poses for generic compilers. Each of those challenges recurs on every forward pass, so what might be an insignificant delay on other workloads can compound into latency failures in real-time video generation.
For example, the challenges posed by dynamic shapes are partly rooted in the shifting nature of scenes in real-time video generation: Imagine a shot showing a pitcher, alone on the mound, that pans out to show the other players and then thousands of fans in the stands, all within seconds. Real-time video generation also involves frames which are generated and evicted in a continuous sliding-window process. “This means attention lengths vary across forward passes, and there are two distinct passes per window: denoising, then cache cleanup,” Wu said. “All of that makes static compilation hard.”
In addition to sliding KV cache copies, video generation workloads contain operations — rotary position embeddings (RoPE) and attention transposes — whose memory access patterns are workload-specific, making them challenging for any general-purpose compiler to fully optimize. For example, in video generation RoPE must contend with three axes (height, width, and time) rather than the single position it accounts for in LLMs. “The 3D rotary embedding interleaves odd and even elements along the innermost dimension, producing many tiny data transfers when compiled generically,” Wu explained.
Finally, because KV caching happens at every layer — each one reads and writes a rolling KV cache — high throughput is required for on-device copies. The need to read old cache contents and write new ones at every layer for every step acts as another significant drag on memory.
Using the Neuron Kernel Interface
To help solve for this additional complexity, the Neuron Science team turned to the Neuron Kernel Interface (NKI).
“NKI lets developers write compute kernels that run directly on NeuronCore hardware, with precise control over how data moves between memory and compute engines,” Wu explained. He noted that NKI gives developers a self-service path to fix hotspots directly, replacing specific bottleneck operations with hardware-tuned implementations where profiling shows that simply compiling models is insufficient.
“In the work we did, the 3D-RoPE kernel went from five seconds to 1.8 milliseconds, cache copies from 23 milliseconds to 1.9 milliseconds per layer, and attention transposes were eliminated entirely by fusing them into the attention kernel,” Wu noted. “For real-time workloads where every millisecond matters, that direct hardware access is what makes production-grade performance achievable.”
Additionally, NKI-Dev-Suite — an agent for generating NKI kernels — produced a working 3D-RoPE kernel on its first attempt. “The combined effect: the pipeline used 11 GB of high-bandwidth memory, while the standard eager-mode path ran out of memory,” Wu noted.
Hybrid sharding strategy
As established, video diffusion models produce token sequences far longer than text models under the real-time requirement. That makes the challenge of self-attention more acute. “Self-attention here operates on 23,400 query tokens attending to 32,760 context tokens, and accounts for about 70% of compute time,” Yu observed. “No single core handles this efficiently without distributing the work.”
To address this, the team used a hybrid sharding strategy entailing sequence parallelism (SP), or partitioning data sequentially, and tensor parallelism (TP), which shards tensors along a specific dimension to distribute computation across multiple devices. For certain non-self-attention parts of the module, the team utilized sequence parallelism. However, for the self-attention portions, only the hybrid approach sufficed.
“LLM attention is causal and unidirectional — each token attends only to previous tokens, the KV cache grows monotonically, and there's one attention type per layer with a single cache policy,” Wu said. Rolling Forcing, however, has two attention types per block: self-attention for spatiotemporal consistency across frames and cross-attention for text conditioning and bidirectional attention within the active window, since all frames are jointly refined from noise.
That, combined with a dual-policy KV cache (a sliding window for recent context plus a permanent attention sink for global context), two forward passes per window (denoising writes noisy KV entries, then a cache update pass overwrites them with clean values, ensuring future windows always attend to clean context), and 3D video token structure constrains how the sequence can be split across cores.
Yu explained that TP alone presents a math problem. “The WAN diffusion transformer model has 12 attention heads, but we have eight Neuron cores per chip, so it's not divisible. Using TP alone means you would have to pad, but padding wastes computation.”
SP alone, on the other hand, can break the 3D structures because, as Yu noted, “You cannot guarantee the sequence partition will be right at the boundary of a frame. Your data must be at least within the granularity of a frame, but if you partition on the frame boundary, those operations won't work.” The hybrid approach splits heads across 4 cores and sequences across 2, keeping both the math and the data layout correct. “The VAE decoder used spatial W-axis sharding, achieving a super-linear 8.25 times speedup,” Yu said.
Model structure changes
The Reactor and Neuron Science teams also optimized parts of the Rolling Forcing model code to run more efficiently on Trainium. “Basically, the model has two phases: one is diffusion, the other is the cache update,” Yu said. “Those are executed in two separate runs, but the issue is the cache update phase has significantly less computation, so if you execute it in a separate round, it has far less hardware utilization. This is because we also have to shard it, and so the high-level principle is that the less data you feed to the chip, the worse hardware utilization you have.”
The team optimized the model code so that the components each of those phases have in common were batched together.
“When we encounter components that are slightly different, we split again and then handle the different components separately. But for most of the pipeline, they are batched together,” Yu explained. “This is specifically useful for Trainium, because each instance has 16 chips and each chip has eight cores. And if you partition your computation across too many cores, each call will just have a small amount of partition computation, and that's underutilizing capacity.”
The result
After employing these, and other optimizations, Reactor and the Neuron Science teams were able to successfully generate a correct video on the first end-to-end run, utilizing Trainium to deliver real-time models. And, the teams emphasized, those results are generalizable.
“Rolling Forcing was the pipeline, it's a very small model,” said Yahav Biran, a principal solutions architect. “You can iterate quickly on it, but it's still a robust system end to end. It has all the complexity that you have in a robust system: the encoder, the DIT, the VAE, the decoder. So basically, if you take a more robust system, it is operating on the same building blocks.”
“We're building common techniques for models that employ autoregressive diffusion which also have a requirement for real-time interaction,“ Yu added. “We're not optimizing a single model only. We're building common techniques for supporting all models with the requirements of real-time streaming.”
The future
Schmidtchen said he is excited about the future this kind of work may enable. “In the not-so-far future, every pixel will be generated in real time, interactively,” he said. “Whole stories can be created by world models in real time: stories that react, that you can engage with, that can even change their entire landscape on the fly.”
He also noted that the work Amazon is doing, and has already done, will do a great deal to make those visions a reality.
“Trainium, is clearly showing a tremendous commitment from AWS and Amazon overall,” Schmidtchen noted. “There's a clearer roadmap of higher performance, better cost performance, more scale globally. In this future where you have real-time interactive AI that needs to be distributed at scale to consumers, physical AI, and more, Trainium is very well positioned to work very well at the inference layer—in terms of its parallelization and its memory and its software stack—and integrate nicely with AWS's global scale infrastructure. We are excited to continue working closely with Amazon as we explore the untapped potential of world models. This is just the beginning.”
