Visko Streams Real-Time World Models on Google Cloud TPUs in 30 Days

AI startup Visko partnered with Google Cloud to deploy Orbis, its interactive live world model, onto Cloud TPUs in just 30 days. The core challenge was not average throughput but beating a strict real-time clock, where every video and audio chunk must render faster than playback speed to prevent the world from freezing. By optimizing the pipeline for a batch size of one and strict sequential dependencies, the team achieved a rendering speed of 1.07x playback.

Architecting the Multimodal Pipeline

The Orbis pipeline chains three distinct generative models on a single Cloud TPU slice running in BF16 precision:

• Core video diffusion model: Generates the base video latents.
• Super-resolution upscaler: Enhances resolution to 1664x936 at 18 fps.
• Audio generator: Produces synchronized audio.

To balance development speed with performance, the stack relies on three open-source pillars:

• JAX: For high-performance array computation and XLA compilation.
• torchax: An interoperability layer allowing PyTorch models to run on TPUs via JAX without a full rewrite.
• Pallas: A custom kernel programming language for TPUs, used to write specialized attention kernels.

Key Optimizations for Real-Time Streaming

To meet the latency budget, the engineering team implemented three specific optimizations:

1. Ulysses Context Parallelism: At batch size 1, the team split the token sequence across TPU chips rather than splitting model weights. High-speed all-to-all interconnect exchanges regrouped tensors so each chip processed the full video sequence for a subset of attention heads, proving faster than tensor parallelism or ring attention for this workload.
2. Custom Pallas Splash Attention: The team integrated a custom splash attention kernel from MaxDiffusion, Google Cloud's open-source reference library. This kernel applies FlashAttention-style tiling to keep computations in on-chip VMEM, avoiding the memory bandwidth bottleneck of materializing full quadratic attention matrices in HBM.
3. Selective JAX Rewrite: While torchax successfully brought the video and upscaling models to TPU quickly, the audio stage remained a bottleneck. Rewriting only the audio model from torchax to native JAX resulted in a 5.8x speedup, bringing the total pipeline under the playback deadline.

Practical Lessons for Startups

The project highlighted several strategies for ultra-low-latency generative AI:

• Native BF16 Precision: Unlike offline generation, continuous streaming suffers from quality degradation if precision is reduced too aggressively. BF16 provided the necessary balance between speed and visual fidelity.
• Strategic Porting: Use torchax as a fast bridge to get PyTorch models running, then profile to identify true bottlenecks before investing in native JAX rewrites.
• XLA Caching: Caching compiled XLA graphs reduced worker restart times from minutes to seconds, crucial for iterative development.