Key takeaways
- 07x playback rendering rate.
- Live world modeling introduces unforgiving serving constraints: batch size is fixed at one to support user steering, and sequential temporal dependencies prevent out-of-order rendering.
- Chaining video diffusion, 1664×936 upscaling at 18 fps, and audio generation on shared silicon avoided the infrastructure costs of multi-cluster deployments.
- Ulysses context parallelism partitioned token sequences across accelerator chips, while custom Pallas splash attention from MaxDiffusion bypassed memory bandwidth bottlenecks.
- A selective optimization methodology using torchax enabled rapid PyTorch bring-up in 30 days, while rewriting only the audio model in pure JAX yielded a 5.8x speedup.
Serving real-time world models in production requires beating an unforgiving clock: each multimodal video and audio chunk must finish rendering before previous playback concludes. AI startup Visko resolved this challenge by deploying Orbis, its steerable live world model, onto a single Google Cloud Tensor Processing Unit (TPU) slice. By chaining video diffusion, super-resolution upscaling (1664×936 at 18 frames per second), and audio generation, the engineering pipeline renders content at 1.07x playback speed with roughly 7% latency headroom. Achieving this sub-playback deadline at a strict batch size of one required combining Ulysses context parallelism, custom Pallas splash attention kernels from MaxDiffusion, and selective native JAX rewrites for bottlenecked audio stages.
The technical deployment, documented in an engineering breakdown on the Google Cloud Blog by Visko CEO Will Yin and Google Cloud developer relations engineer Ivan Nardini, transitioned from kickoff in August 2026 to production handover in 30 days. Visko publicly demonstrated Orbis on September 1, 2026, targeting use cases spanning interactive entertainment, robotics simulation, virtual companions, and live commerce. While traditional video generators run batch-oriented inference over minutes, live world models operate under hard real-time latency service-level objectives (SLOs) that demand fundamentally different architectural trade-offs.
The Latency Constraints of Streaming World Models
Most generative video systems, including foundational models from Runway, OpenAI, or Kuaishou, operate asynchronously. A user submits a prompt, waits in a job queue, and receives a finalized, static MP4 file after several seconds or minutes of rendering. In this asynchronous paradigm, latency is forgiving; system architects can optimize for overall throughput by batching concurrent generation requests across expansive accelerator clusters.
Live world models eliminate the luxury of a waiting room. In Visko’s Orbis system, the model streams an ongoing visual world that viewers steer interactively. For example, during a live session depicting an animal resting by a hearth, an audience member can submit a text prompt introducing a food bowl; the environment responds on the fly, rendering the object and corresponding behavioral transitions without resetting playback. To maintain continuous playback without frame buffering or world freezes, two rigid technical constraints govern system design:
- Batch size of one: Because each live stream represents an independent user navigating a unique, steerable path through latent space, requests cannot be batched across users. Infrastructure engineers cannot mask individual inference latency by aggregating queries into large batches, which is the conventional method for saturating accelerator tensor cores.
- Strict sequential temporal dependency: The generative diffusion pipeline maintains a rolling memory buffer of preceding video latents to preserve spatial consistency and object identity across time. Consequently, chunk N+1 cannot begin processing until chunk N completes. Rendering cannot be parallelized across future timestamps, and frames cannot be rendered out of order.
These structural boundaries transform inference serving from a tokens-per-second throughput optimization problem into an unforgiving per-chunk latency budget.
Chaining Multimodal Generation on a Single TPU Slice
To keep infrastructure operational costs manageable for real-time consumer and enterprise workloads, Visko designed its multimodal pipeline to run entirely within a single Cloud TPU slice running in native 16-bit Brain Floating Point (BF16) precision. Rather than provisioning disparate accelerator clusters for vision and sound, the pipeline chains three distinct generative models sequentially on shared silicon:
- Video diffusion transformer: Generates continuous spatiotemporal latent representations based on conditioning prompts and rolling historical latents. This stage represents approximately two-thirds of the total per-chunk compute budget.
- Super-resolution upscaler: Decodes the generated latents and performs spatial upscaling, outputting high-definition video at a resolution of 1664×936 running at 18 frames per second.
- Audio generation model: Synthesizes synchronized, context-aware audio tracks matching the visual dynamics generated by the diffusion model.
Managing this pipeline required a hybrid software framework that balanced developer velocity with hardware-level execution speed. The team structured the stack around three open-source components: Google’s JAX numerical computing framework, the XLA (Accelerated Linear Algebra) optimizing compiler, and torchax—Google’s interoperability tool that dispatches PyTorch models to JAX and XLA through Python-level interceptors without requiring an architectural rewrite.
| Pipeline Stage | Primary Function | Implementation Runtime | Compute Share & Primary Bottleneck | Optimization Applied |
|---|---|---|---|---|
| Video Diffusion | Spatiotemporal latent denoising with rolling memory | PyTorch via torchax | ~66% of per-chunk budget; attention memory bandwidth | Ulysses context parallelism and custom Pallas splash attention |
| Spatial Upscaler | Latent decoding to 1664×936 resolution at 18 fps | PyTorch via torchax | Secondary compute share; pixel reconstruction latency | XLA compilation via torchax without full rewrite |
| Audio Synthesis | Synchronized ambient and event audio generation | Native pure-JAX | Critical end-of-chunk bottleneck in initial PyTorch port | Pure-JAX reimplementation delivering 5.8x speedup |
Three Optimizations That Unlocked 1.07x Playback Speed
Initial deployment of the full PyTorch pipeline onto TPUs did not achieve real-time playback. Reaching 1.07x playback speed—leaving a 7% timing buffer to accommodate real-world network and compute jitter—required three distinct optimizations targeting memory bandwidth, sequence distribution, and execution overhead.
1. Ulysses Context Parallelism Across Accelerator Chips
In diffusion transformers handling high-dimensional video, spatial and temporal patches expand into substantial token sequences. Because batch size is constrained to one, traditional data parallelism is impossible, while standard tensor parallelism introduces high communication frequency across weight matrices. Ring attention presents latency overheads when token sequences are moderate in size.
Instead, the engineering team applied Ulysses context parallelism across the TPU slice. Under this layout, the input token sequence is split evenly across all chips within the slice during feed-forward operations. Immediately prior to the attention layer, high-speed all-to-all Inter-Chip Interconnect (ICI) communications regroup the tensors: each individual chip receives the full sequence length, but only for a subset of attention heads (one-quarter of the heads in a four-way partition). After attention is computed across the full sequence for those assigned heads, a second all-to-all communication restores the sequence-partitioned layout for subsequent feed-forward layers.
2. Custom Pallas Splash Attention from MaxDiffusion
Standard attention layers exhibit quadratic memory complexity with respect to sequence length, quickly saturating High Bandwidth Memory (HBM) on AI accelerators. To prevent memory bandwidth bottlenecks, the team incorporated a custom Pallas splash attention kernel from MaxDiffusion, Google Cloud’s open-source reference library for diffusion architectures.
Pallas serves as JAX’s specialized kernel language, granting fine-grained control over TPU hardware memory hierarchies, similar to how Triton operates on graphics processing units. Splash attention adopts FlashAttention tiling concepts for TPU architecture: instead of allocating and materializing an intermediate quadratic attention matrix in HBM, the kernel partitions queries, keys, and values into compact blocks that fit inside the TPU’s ultra-fast on-chip Vector Memory (VMEM). Intermediate softmax operations are computed online within VMEM, eliminating round-trip memory transfers and allowing the diffusion model to stay within its aggressive latency budget.
3. Selective Pure-JAX Rewrite Delivering 5.8x Audio Speedup
A frequent dilemma for AI startups migrating from research prototypes to production hardware is whether to rewrite extensive PyTorch codebases in JAX from scratch. Visko adopted a selective optimization strategy. The video diffusion and upscaling components ran adequately under torchax once paired with the Pallas splash attention kernel, meeting latency targets without code refactoring.
However, profiling the complete pipeline revealed that the audio synthesis stage incurred disproportionate dispatch overhead under torchax, threatening the strict playback deadline. Rather than undertaking an end-to-end rewrite of the entire system, engineers isolated the audio model and rewrote it in native, pure-JAX. This targeted migration yielded a 5.8x execution speedup for the audio stage, shifting the entire multi-stage pipeline from lagging playback to rendering at 1.07x real-time speed.
Operational Lessons for Real-Time AI Systems
The engineering techniques established during the Visko deployment offer several practical takeaways for platform teams and machine learning engineers deploying streaming generative workloads:
- Native BF16 avoids drift over rolling horizons: While modern 8-bit floating-point (FP8) quantization reduces memory pressure and accelerator workload, quantization errors accumulate across chained frames in continuous streaming sessions. Over extended rollouts, minor numerical imprecision compounds into visible visual drift and artifacting. Maintaining native BF16 precision preserved scene fidelity across long-horizon interactions while meeting latency limits.
- Graph compilation caching stabilizes cold starts: Compiling XLA computation graphs on initial boot requires several minutes of processing before an accelerator slice can begin inference. To prevent severe downtime during autoscaling events or worker node failures, teams must serialize and cache compiled XLA artifacts. Caching reduces subsequent container and worker restart durations to under 60 seconds.
- Rapid prototyping with torchax informs targeted tuning: Relying on torchax allowed Visko to migrate complex PyTorch models to Cloud TPUs in days. Using the Cloud TPU Profiler against the live pipeline allowed engineers to identify the single bottleneck—audio dispatch—rather than prematurely rewriting working diffusion layers.
- Flexible compute allocation: For early-stage validation, provisioning TPU Flex-start VMs allowed benchmarking and kernel tuning without committing to multi-year capacity reservations, reflecting wider shifts in Google Cloud TPU hardware infrastructure and developer access models.
Bottom line
Visko’s co-engineering effort with Google Cloud demonstrates that live, steerable video world models can be served from a single accelerator slice in real time, challenging the assumption that interactive world simulation requires massive, multi-node GPU clusters. Achieving 1.07x playback speed with 7% latency headroom at batch size one validates the viability of real-time interactive media, robotics simulation environments, and dynamic digital companions.
Nevertheless, engineering teams evaluating this architecture must weigh practical constraints. A 7% latency margin leaves narrow room to absorb sudden compute variance or network packet drops before frame stalls affect user experience. Furthermore, serving interactive streams at batch size one means compute costs scale linearly with concurrent users, without the margin expansion typical of batched offline APIs. For technical leaders, the key architectural takeaway is methodological: leverage framework bridges like torchax to validate hardware viability immediately, rely on hardware-native tiling kernels for attention bottlenecks, and reserve complete framework rewrites strictly for measured pipeline constraints.
