Visko’s Orbis: Live World Model Streaming on Cloud TPUs

Generative video is moving from static clip generation to interactive, real-time worlds, but serving a live world model is a hard-deadline problem rather than a throughput problem. Visko, an AI startup building steerable ‘Live Models,’ partnered with Google Cloud TPU engineers in August 2026 to port its flagship model Orbis onto Cloud TPUs. The goal was to stream continuous video and synchronized audio in real time, with each chunk rendered before the viewer finishes watching the previous one. In 30 days, the teams brought the full multimodal pipeline onto a single Cloud TPU slice running in BF16, reaching 1.07x playback speed while maintaining about 7% headroom to absorb latency variance.

Orbis starts from a single text or image prompt and keeps generating; viewers can enter steering prompts mid-stream without restarting. A test session shows a cat napping by a fireplace, and when the viewer steers to add a bowl of food, the bowl appears seamlessly and the cat reacts. Because each session is one viewer’s interactive stream, the batch size is effectively one, and the model’s rolling memory of recent video latents creates a strict sequential dependency — chunk N+1 depends on chunk N — so chunks cannot be pre-rendered out of order. Every chunk must be denoised, decoded from latents, upscaled to 1664×936 at 18 fps, and paired with generated audio within the playback deadline.

To keep infrastructure lean, the pipeline chains three models on one TPU slice: the core video diffusion model, a super-resolution upscaler, and an audio generator, sharing chips for each chunk. The software stack uses JAX for compiled accelerator code, torchax to run PyTorch models directly on TPUs, and Pallas for custom kernels. Video denoising accounts for nearly two-thirds of the compute budget, so most optimization effort focused there.

Three optimizations proved decisive. First, Ulysses context parallelism distributes the long spatiotemporal token sequence across TPU chips at batch size 1; all-to-all ICI exchanges switch between a sequence-split layout for feed-forward layers and a head-split layout for attention, and this beat tensor parallelism and ring attention. Second, a custom Pallas splash attention kernel from MaxDiffusion brings FlashAttention-style tiling to TPUs, keeping attention tiles in on-chip VMEM instead of materializing the full quadratic matrix in HBM. Third, the audio model was rewritten from torchax to native JAX, yielding a 5.8x speedup and bringing the pipeline under the playback deadline. The video and upscaling models stayed in torchax; the lesson is to use torchax as a fast bring-up bridge, profile the full pipeline, and rewrite only the components that dominate latency.

Field notes from the port include keeping native BF16 rather than FP8, since small quantization errors accumulate across chunks in a continuous stream; caching XLA compiled artifacts to reduce worker restarts from minutes to under a minute; and using TPU Flex-start VMs for short development runs. Google Cloud also points builders to the open-source splash attention kernel in MaxDiffusion, torchax for fast PyTorch migration, and the Google for Startups Cloud Program for credits and mentorship.

How Visko streams real-time world models on Google Cloud TPUs

View Original