Reactor and Amazon Optimize Real-Time Video Generation on Trainium Using NKI

Reactor and the Amazon Neuron Science team developed a kernel-centric optimization strategy to enable real-time interactive video generation on AWS Trainium. By leveraging the Neuron Kernel Interface (NKI) and a hybrid sharding approach, the teams successfully ran the Rolling Forcing autoregressive diffusion model, achieving frame rates above the 16 fps standard required for smooth playback.

Kernel-Level Optimizations

The teams addressed three primary bottlenecks—dynamic shapes, unusual memory access patterns, and heavy cache management—that generic compilers struggle to handle in real-time video workloads. Using NKI, they replaced inefficient operations with hardware-tuned kernels:

• 3D-RoPE: Latency dropped from 5 seconds to 1.8 milliseconds.
• Cache Copies: Reduced from 23 milliseconds to 1.9 milliseconds per layer.
• Attention Transposes: Eliminated entirely by fusing them into the attention kernel.

These optimizations allowed the pipeline to fit within 11 GB of high-bandwidth memory, whereas the standard eager-mode path ran out of memory.

Hybrid Sharding Strategy

Video diffusion models require processing long token sequences, with self-attention accounting for about 70% of compute time. A single core cannot handle the workload efficiently, and standard parallelism methods presented issues:

• Tensor Parallelism (TP) alone: The model has 12 attention heads, which is not divisible by the 8 Neuron cores per chip, leading to wasted computation via padding.
• Sequence Parallelism (SP) alone: Risked breaking the 3D video token structure if partitions did not align with frame boundaries.

The team implemented a hybrid strategy, splitting heads across 4 cores and sequences across 2. This approach maintained data layout correctness and achieved a super-linear 8.25x speedup for the VAE decoder using spatial W-axis sharding.

Model Structure Changes

To improve hardware utilization, the teams optimized the model code by batching common components between the diffusion and cache update phases. Since the cache update phase has significantly less computation, executing it separately resulted in poor utilization. By batching shared components and only splitting where necessary, the teams ensured that the 16 chips and 8 cores per instance were used more effectively.

The resulting pipeline successfully generated correct video on the first end-to-end run. The teams emphasize that these techniques are generalizable to other autoregressive diffusion models requiring real-time interaction, positioning Trainium as a viable platform for scalable, interactive world models.